AgentExecutionEngine

from transformers import AutoTokenizer
from rllm.engine import AgentExecutionEngine
from rllm.agents import MyAgent
from rllm.environments import MyEnv

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")

engine = AgentExecutionEngine(
    # Agent and environment classes
    agent_class=MyAgent,
    env_class=MyEnv,
    agent_args={"system_prompt": "You are a helpful assistant"},
    env_args={"max_turns": 3},
    
    # Backend configuration
    engine_name="openai",  # or "verl" for training
    tokenizer=tokenizer,
    rollout_engine_args={
        "base_url": "http://localhost:30000/v1",
        "api_key": "your-api-key"
    },
    
    # Sampling parameters
    sampling_params={
        "temperature": 0.6,
        "top_p": 0.95,
        "model": "Qwen/Qwen3-4B"
    },
    
    # Execution parameters
    n_parallel_agents=128,      # Number of concurrent agent-env pairs
    max_steps=10,               # Max steps per trajectory
    max_response_length=4096,   # Max tokens in entire trajectory
    max_prompt_length=2048,     # Max tokens in prompt
    
    # Performance tuning
    trajectory_timeout=300,     # Timeout per trajectory (seconds)
    retry_limit=3,              # Retries on failure
    max_workers=64,             # Thread pool size for env operations
)

import asyncio
from rllm.data import DatasetRegistry

# Load tasks
tasks = DatasetRegistry.load_dataset("math", "test").get_data()

# Execute all tasks
trajectories = await engine.execute_tasks(tasks)

# Each trajectory contains complete interaction history
for traj in trajectories:
    print(f"Task: {traj.task}")
    print(f"Reward: {traj.reward}")
    print(f"Steps: {len(traj.steps)}")
    
    for i, step in enumerate(traj.steps):
        print(f"  Step {i}: {step.model_response}")
        print(f"  Reward: {step.reward}")

# Initialize agents and environments
envs = [MyEnv.from_dict(task) for task in tasks]
agents = [MyAgent(**agent_args) for _ in tasks]
engine.update_envs_and_agents(envs, agents)

# Generate trajectories as async generator
async for trajectory in engine.trajectory_generator(mode="Text"):
    print(f"Generated trajectory with reward: {trajectory.reward}")
    # Process trajectory for training

# Trajectories can terminate due to:
# - ENV_DONE: Environment signals completion
# - MAX_STEPS: Reached max_steps limit
# - TRUNCATION: Exceeded max_response_length
# - TIMEOUT: Trajectory timeout exceeded
# - ENV_TIMEOUT: Environment step timeout

# Access termination info in trajectory
for traj in trajectories:
    if traj.info.get("termination_reason") == "TRUNCATION":
        print(f"Trajectory {traj.uid} was truncated")

engine = AgentExecutionEngine(
    # ... other args ...
    retry_limit=3,  # Retry failed trajectories up to 3 times
)

# Retries happen automatically for:
# - API errors
# - Timeout errors
# - Environment errors (if recoverable)

# In "Token" mode, trajectories are assembled into training format:
# prompt_tokens: Initial prompt from first step
# response_tokens: Concatenated completions + environment messages
# response_masks: Binary mask indicating which tokens contribute to loss

async for token_result in engine.trajectory_generator(mode="Token"):
    prompt_tokens = token_result["prompt_tokens"]      # torch.Tensor
    response_tokens = token_result["response_tokens"]  # torch.Tensor
    response_masks = token_result["response_masks"]    # torch.Tensor (0/1)
    trajectory_reward = token_result["trajectory_reward"]

import asyncio
from transformers import AutoTokenizer
from rllm.engine import AgentExecutionEngine
from rllm.agents.frozenlake_agent import FrozenLakeAgent
from rllm.environments.frozenlake import FrozenLakeEnv
from rllm.data import DatasetRegistry
from rllm.utils import compute_pass_at_k

# Setup
model_name = "Qwen/Qwen3-4B"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Initialize engine
engine = AgentExecutionEngine(
    agent_class=FrozenLakeAgent,
    env_class=FrozenLakeEnv,
    agent_args={"max_steps": 10, "use_accumulate_history": True},
    env_args={"max_steps": 8, "is_slippery": False},
    engine_name="openai",
    tokenizer=tokenizer,
    rollout_engine_args={
        "base_url": "http://localhost:30000/v1",
        "api_key": "None",
    },
    sampling_params={
        "temperature": 0.6,
        "top_p": 0.95,
        "model": model_name
    },
    max_response_length=16384,
    max_prompt_length=4096,
    n_parallel_agents=256,
)

# Load tasks
tasks = DatasetRegistry.load_dataset("frozenlake", "test").get_data()

# Execute
results = asyncio.run(engine.execute_tasks(tasks))

# Evaluate
compute_pass_at_k(results)