# TRL (https://www.agentenvframework.com/docs/training/trl) > Train a model on a bundle's tasks with TRL's GRPOTrainer or AsyncGRPO, through agent_env.training.trl, starting with the GridSim bundle ## How a rollout runs 1. The trainer takes a row of `env.dataset`, one task of the split, and resets an environment from `env.factory` with it. 2. The server starts the row’s task as an episode: it deploys the env, loads the task, and opens the agent’s turn for your model. 3. Your model plays the turn: each tool call is a method on the environment, which runs it on the env through the episode. 4. Your model ends its turn with `end_turn`, and the task’s verifier scores the episode, as it would score any agent’s run. 5. `get_reward()` hands the score to the trainer as the reward; the episode is torn down, and the session takes the next row. Parts of the scene: - **Your trainer**: TRL’s `GRPOTrainer` or AsyncGRPO, with `train_dataset=env.dataset` and `environment_factory=env.factory`. It owns the model, the generation and the update; the environment owns the reward. - **A dataset row**: One task of the split: its prompt as chat messages, and the `split`, `index` and `task_id` that `reset` plays. The prompt is yours to change; the task stays the same. - **An environment**: An instance of the class `env.factory` builds, holding one session on the server. Its methods are the turn’s tools plus `end_turn`; `reset` starts a task and `get_reward` reads its score. TRL reuses it rollout after rollout. - **The session**: One OpenEnv WebSocket session per environment, so one episode at a time. Every call carries a deadline the server ends it at: a missed one ends the episode, unscored, and the session goes on. - **The OpenEnv server**: `agent_env.openenv.build_app(bundle)` serves every task of a bundle as an OpenEnv env, the split `all` plus one per eval. `AgentEnv.local(bundle)` starts it in the trainer’s own process. - **The env**: The task’s env, deployed fresh for this episode in a sandbox, here a container on your Docker. Its tools are what the model calls. - **The agent’s turn**: The task prompts one agent, and your model plays it from outside: the episode deploys a proxy in the agent’s place and hands the turn’s prompt and tools to the environment. - **The verifier**: The task’s own verifier step scores the episode once the turn ends, the same score an evaluation run records. With several scores, `reward=` picks the one to train on. - **The reward**: `get_reward()` returns the score, or `None` when the episode ended in an error, which TRL leaves out of the update. GRPO compares it with the other rollouts of the same row. Three pieces connect a trainer to a bundle: | Piece | What it does | | ------------------------------------- | ------------------------------------------------------------------------------------------------------------------- | | `agent_env.openenv.build_app(bundle)` | Serves every task of a bundle as an [OpenEnv](https://www.agentenvframework.com/docs/openenv.md) env: the split `all`, plus one split per eval. | | `agent_env.training.trl.AgentEnv` | Reads the server's tasks into a dataset, and builds the `environment_factory` whose environments play the episodes. | | The Episode API | Runs a task with its one prompted agent played from outside, and scores it with the task's verifiers. | ## Install You need Python 3.11 or newer, Docker running, and a GPU for training: ```bash title="Terminal" pip install 'agentenv-framework[trl]' 'trl[vllm,peft]' 'vllm==0.29.*' ``` The `trl` extra brings OpenEnv and `datasets`. TRL and vLLM are yours: TRL 1.15 or later, which lets an environment own its reward, and vLLM 0.29 or older. > [!WARNING] > **vLLM 0.29 or older, for now** > > OpenEnv 0.10 and agent-env both need mcp 1, and vLLM 0.30 and later need mcp 2. Installing the > extra over a newer stack also moves `openai` below 3 (litellm caps it) and `pillow` and `tomlkit` > down (OpenEnv's gradio caps them). TRL, vLLM, transformers and torch stay where you pinned them. ## Train on GridSim GridSim is a bundle to start from. Each task runs one grid battery for one day in the ERCOT day-ahead market: the model reads the day's 24 hourly prices and submits a schedule, one integer an hour, to charge, sell or idle. The env grades the schedule against the day's exact optimum: 0 for nothing useful, 0.2 for idling all day, 1 for the optimum. The bundle has 2,000 train days from 2014 to 2024 and 200 eval days from 2025, on real ERCOT prices. ### Download the bundle The bundle is a folder in a Hugging Face dataset. The `hf` command comes with the install: ```bash title="Terminal" hf download earakely-scale/GridSimEnv-AgentEnv --repo-type dataset --include "bundles/gridsim-v1/*" --local-dir . ``` ```text title="bundles/gridsim-v1" envs/gridsim/ Dockerfile and env.toml: the battery env and its price packs artifacts/ gridsim-verifier, which grades a schedule from the env's report tasks/ 2,200 tasks, one JSON file each evals/ train (2,000), eval (200) ``` ### Play a rollout by hand `AgentEnv.local` serves the bundle in your process, on your Docker. The first start builds the env's image, in about a minute: ```python title="Python" from agent_env.training.trl import AgentEnv env = AgentEnv.local("bundles/gridsim-v1", split="train") len(env.dataset), len(env.dataset_for("eval")) ``` ```text title="Output" (2000, 200) ``` Then play one rollout the way TRL will. The task leaves the model one tool, `submit_schedule`: ```python title="Python" environment = env.factory() environment.reset(**env.dataset[0]) environment.submit_schedule(actions=[0] * 23) environment.submit_schedule(actions=[0] * 24) environment.get_reward() ``` ```text title="Output" '{"submitted":false,"error":"expected exactly 24 integers, got 23"}' '{"submitted":true,"feasible":true,"truncated_at":null,"violation":null,"soc_path":[0,0,…,0],"profit_usd":0.0}' 0.2 ``` The env refuses a schedule that isn't 24 hours, and an idle day scores 0.2. To beat it, the model has to charge when power is cheap and sell when it's dear, within the battery's limits. ### Warm start with SFT Played as it comes, Qwen3-4B idles. On the 200 eval days it scores 0.18: it submits an idle day on most rollouts, and beats idling on 1%. GRPO from it has nothing better to reinforce: in three steps of 32 rollouts, one beat idle. A short SFT on the train days' best schedules gives it a start. The rows are published beside the bundle: each train day's prompt, a one-line rationale and the day's optimal `submit_schedule` call, with the tools `env.factory` offers and thinking off, so SFT trains on the exact context GRPO generates from. ```python title="sft.py" from datasets import load_dataset from peft import LoraConfig from trl import SFTConfig, SFTTrainer rows = load_dataset("earakely-scale/GridSimEnv-AgentEnv", data_files="sft/gridsim-v1/train.jsonl", split="train") trainer = SFTTrainer( model="Qwen/Qwen3-4B", args=SFTConfig( output_dir="sft", num_train_epochs=1, max_length=2048, per_device_train_batch_size=8, learning_rate=1e-4, lr_scheduler_type="cosine", warmup_steps=0.1, bf16=True, gradient_checkpointing=True, seed=0, ), train_dataset=rows, peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear", task_type="CAUSAL_LM"), ) trainer.train() trainer.model.merge_and_unload().save_pretrained("qwen3-4b-gridsim-sft") trainer.processing_class.save_pretrained("qwen3-4b-gridsim-sft") ``` ### Train with GRPO Then GRPO plays the bundle's train days from the SFT model, and scores the 200 eval days at the start and every 20 steps: ```python title="train.py" from peft import LoraConfig from trl import GRPOConfig, GRPOTrainer from agent_env.training.trl import AgentEnv env = AgentEnv.local("bundles/gridsim-v1", split="train", max_concurrent_envs=32) trainer = GRPOTrainer( model="qwen3-4b-gridsim-sft", args=GRPOConfig( output_dir="out", max_steps=40, per_device_train_batch_size=8, gradient_accumulation_steps=4, # 32 rollouts a step: 4 days, 8 rollouts each num_generations=8, max_completion_length=1024, max_tool_calling_iterations=2, chat_template_kwargs={"enable_thinking": False}, temperature=1.1, learning_rate=2e-5, lr_scheduler_type="cosine", warmup_steps=3, loss_type="dr_grpo", beta=0.0, epsilon_high=0.28, use_vllm=True, vllm_mode="colocate", vllm_gpu_memory_utilization=0.3, vllm_max_model_length=4096, bf16=True, gradient_checkpointing=True, eval_strategy="steps", eval_steps=20, eval_on_start=True, num_generations_eval=1, per_device_eval_batch_size=32, seed=0, ), train_dataset=env.dataset, eval_dataset=env.dataset_for("eval"), environment_factory=env.factory, peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear", task_type="CAUSAL_LM"), ) trainer.train() ``` There are no `reward_funcs`: each environment owns its reward through `get_reward()`, which TRL reads once the rollout ends. ```json title="Output (step 12 of the GridSim run, trimmed)" {"loss": -0.0452, "grad_norm": 0.0272, "reward": 0.6883, "reward_std": 0.3281, "rewards/gridsim-v1/mean": 0.6883, "tools/call_frequency": 2.0, "tools/failure_frequency": 0.0, "completions/mean_length": 520.88, "step_time": 186.98} ``` The reward column is named after the environment class, here the bundle's name. The server needs a session for each rollout of a step, here 32. ### What it learns | Qwen3-4B | eval reward (200 days) | |---|---| | base | 0.180 | | after SFT (GRPO step 0) | 0.569 | | GRPO step 20 | 0.562 | | GRPO step 40 | 0.614 | Mean train reward of each GRPO step's 32 rollouts, by step: 1: 0.543, 2: 0.865, 3: 0.642, 4: 0.717, 5: 0.637, 6: 0.514, 7: 0.443, 8: 0.526, 9: 0.718, 10: 0.809, 11: 0.499, 12: 0.688, 13: 0.636, 14: 0.699, 15: 0.736, 16: 0.618, 17: 0.802, 18: 0.630, 19: 0.861, 20: 0.400, 21: 0.511, 22: 0.684, 23: 0.603, 24: 0.587, 25: 0.606, 26: 0.576, 27: 0.715, 28: 0.649, 29: 0.712, 30: 0.722, 31: 0.677, 32: 0.555, 33: 0.642, 34: 0.677, 35: 0.697, 36: 0.664, 37: 0.658, 38: 0.759, 39: 0.834, 40: 0.352. | Qwen3-4B, on the 200 eval days | Eval reward | Beats idle | At the optimum | | ------------------------------ | ----------- | ---------- | -------------- | | Base | 0.18 | 1% | 0% | | After SFT | 0.57 | 75% | 17% | | After 40 GRPO steps | 0.61 | 79% | 14% | SFT does most of the work, from 0.18 to 0.57. The 40 GRPO steps add 0.05 on the eval, within what one rollout a day can tell apart (about ±0.06), and the train reward stays level. GRPO on this task is slow: [GridSim's published runs](https://huggingface.co/datasets/earakely-scale/GridSimEnv-AgentEnv) took the same model to 0.80 on these 200 days, with SFT and then 330 steps of 64 rollouts. GridSim's greedy rule scores 0.76 on them. The run took four hours on one L40S with 4 vCPUs, an AWS g6e.xlarge, at about $8.50 on demand: 18 minutes for the base eval, 26 for SFT, and three hours for GRPO with its three evals. A GRPO step took about 190 seconds, most of it the envs: each step deploys 32 GridSim envs at once, so give the machine that serves them more CPU cores. ## Serve the bundle on its own host To keep the envs off the GPU machine, serve the bundle where they can deploy, and point the trainer at it. The rest of `train.py` stays the same: ```python title="serve.py, on the env host" import uvicorn from agent_env.openenv import build_app app = build_app("bundles/gridsim-v1", max_concurrent_envs=48) uvicorn.run(app, host="0.0.0.0", port=8000) ``` ```python title="train.py, on the GPU machine" env = AgentEnv("http://env-host:8000", split="train") ``` Each environment holds one session on the server, so `max_concurrent_envs` must cover every rollout in flight. `build_app(..., sandbox="modal")` deploys each episode's envs on another [sandbox provider](https://www.agentenvframework.com/docs/sandboxes.md) instead of the host's Docker. ## What the trainer gets `env.dataset` holds a row per task of the split, and `env.dataset_for("eval")` another split's rows: ```python title="Python" env.dataset[0] ``` ```json title="Output (prompts cut)" { "prompt": [ {"role": "system", "content": "You operate a grid battery in the ERCOT day-ahead market (DAM) for one operating day…"}, {"role": "user", "content": "{\"task_id\": \"dam-v1-2021-05-06-HB_HUBAVG-1h-r\", \"tier\": \"foresight\", \"market\": \"ERCOT DAM\"…"} ], "split": "train", "index": 0, "task_id": "@local/…/gridsim-v1/dam-v1-2021-05-06-HB_HUBAVG-1h-r" } ``` `env.factory` is the class TRL builds environments from. Its methods are the tools the task's turn offers, typed and documented from each tool's JSON Schema, plus `end_turn`: ```python title="Python" import inspect [name for name in vars(env.factory) if not name.startswith("_")] inspect.signature(env.factory.submit_schedule) ``` ```text title="Output" ['submit_schedule', 'end_turn'] (self, *, actions: list[int]) -> str ``` TRL reads each method the way transformers' `get_json_schema` does, so the model sees every tool as the env describes it. A tool whose name isn't a Python identifier, or is `reset` or `get_reward`, is refused when the class is built, with the fix in the message. ## Which trainer 1. Each rollout is an episode on the server. `GRPOTrainer` plays a batch of rows at once; AsyncGRPO keeps up to `max_inflight_tasks` going. 2. `GRPOTrainer` trains once its whole batch is scored, with generation paused; AsyncGRPO trains on scored rollouts while the rest keep playing. 3. AsyncGRPO sends the new weights to its vLLM server after each step, and drops rollouts more than `max_staleness` steps old. Parts of the scene: - **GRPOTrainer**: `trl.GRPOTrainer`: generation and training take turns on the same GPUs, here with vLLM colocated. The simplest to run, and the one to start with. - **AsyncGRPO**: `trl.experimental.async_grpo.AsyncGRPOTrainer`: a vLLM server on its own GPU generates while the trainer trains on another, so neither waits for the other. - **The GPUs**: `GRPOTrainer` can train and generate on one GPU. AsyncGRPO needs two at least: vLLM on one, the trainer on another, and its packed attention runs on A100, L40S, A10G, L4 and H100 but not on Blackwell. - **A rollout**: One environment playing one row’s episode, from `reset` to `get_reward`. Its length is the deploy, the model’s turns and the tool calls, so rollouts end at different times. - **The generation batch**: `GRPOTrainer` plays a batch of rollouts together and waits for the slowest. The server needs a session for each: `max_concurrent_envs` at least the batch. - **Rollouts in flight**: AsyncGRPO’s rollout process keeps `max_inflight_tasks` rollouts going, each on an environment from its pool. The server needs a session for each one. - **A train step**: The policy update on scored rollouts. GRPO compares each rollout’s reward with the rest of its row’s group, so a group needs rewards that differ to teach anything. - **Weight sync**: After a step, AsyncGRPO sends the trainer’s weights, or with `--enable-lora` just the LoRA adapter, to the vLLM server, so the next rollouts use the new policy. - **Staleness**: A rollout started under an older policy is still trained on, unless it is more than `max_staleness` steps old, when it is dropped and its environment goes back to the pool. `GRPOTrainer` is the one to start with: it trains and generates on the same GPUs, one after the other. AsyncGRPO overlaps them, with a vLLM server on its own GPU, and takes the same `env.factory`. ## Train with AsyncGRPO AsyncGRPO needs two GPUs at least: vLLM on one, the trainer on another. The example is the run that measured it, on the PortSim training pack: berth planning at the Port of Barcelona, a multi-turn task with three tools, served on its own host. Start vLLM first: ```bash title="Terminal 1: vLLM on GPU 0" CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 vllm serve Qwen/Qwen3-4B \ --port 8001 --max-model-len 32768 --logprobs-mode processed_logprobs \ --weight-transfer-config '{"backend":"nccl"}' --enable-lora --max-lora-rank 16 --max-loras 6 ``` ```python title="train_async.py" from peft import LoraConfig from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer from agent_env.training.trl import AgentEnv def main(): env = AgentEnv("http://env-host:8000", split="train-standard", tool_timeout=300, reset_timeout=1800) trainer = AsyncGRPOTrainer( model="Qwen/Qwen3-4B", args=AsyncGRPOConfig( output_dir="out", per_device_train_batch_size=8, gradient_accumulation_steps=2, num_generations=8, max_completion_length=2048, max_tool_calling_iterations=8, chat_template_kwargs={"enable_thinking": False}, dtype="bfloat16", bf16=True, gradient_checkpointing=True, vllm_server_base_url="http://127.0.0.1:8001", max_inflight_tasks=32, ), train_dataset=env.dataset, environment_factory=env.factory, peft_config=LoraConfig(r=16, lora_alpha=32, target_modules="all-linear", task_type="CAUSAL_LM"), ) trainer.train() if __name__ == "__main__": main() ``` ```bash title="Terminal 2: the trainer on GPU 1" CUDA_VISIBLE_DEVICES=1 python train_async.py ``` ```json title="Output (PortSim, one step's metrics, trimmed)" {"step": 3, "loss": -0.2684, "grad_norm": 0.0896, "rewards/portsim-dock-v1": 0.0, "reward_std": 0.0055, "tools/call_frequency": 8.0, "tools/failure_frequency": 0.0, "perf/weight_sync_s": 0.41, "sample/staleness_mean": 2.0} ``` On four L40S, this trained 5 steps of 16 PortSim rollouts in 391 seconds, with 32 in flight. > [!WARNING] > **What AsyncGRPO needs** > > * **The main guard.** AsyncGRPO plays rollouts in a process it spawns, which imports your script > again. `env.factory` pickles into it, and the process builds its own environments. > * **Room in the context.** AsyncGRPO asks vLLM for `max_completion_length` tokens on every turn, > however long the conversation has grown, and a request past `--max-model-len` fails the run. Size > `--max-model-len` for the longest conversation plus `max_completion_length`, and cap the turns > with `max_tool_calling_iterations`. > * **A supported GPU.** Its packed attention loads FlashAttention 3, built for A100, L40S, A10G, L4 > and H100, not for Blackwell. > * **Sessions.** `max_inflight_tasks` rollouts at once, so the server's `max_concurrent_envs` at > least that. ## What the model sees The model starts from the row's `prompt`, which is yours to change, and calls the turn's tools as the environment's methods. Its reply reaches the task only through `end_turn(response=...)`: TRL hands an environment no completion. `end_turn` answers with the next turn's prompt, after its system prompt when that changes, or once the last turn is over: ```text title="end_turn(response=...)" The episode has ended. ``` A turn the model leaves open is ended with an empty reply before the episode is scored, so tell the model, in the system prompt, to end its turn with `end_turn`. When the played agent works in a sandbox, its turns also offer `bash`. ## Rewards The reward is the task's score, the one an evaluation run records. A task that records one score has its reward. With several, pick one by its step id, or combine them: ```python title="train.py" env = AgentEnv.local("my-bundle", reward="grade") env = AgentEnv.local("my-bundle", reward=lambda scores: 0.8 * scores["grade"] + 0.2 * scores["style"]) ``` `get_reward()` returns `None`, and TRL leaves the rollout out of the update, when its episode ended in an error: a deploy that failed, a call that missed its deadline, or a server that went away. Each one is logged: ```text title="Log" WARNING agent_env.training.trl: rollout left unscored: submit_schedule didn't answer within 60s, and the episode ended, since what the call does may still be under way ``` One error stops training instead: a rollout whose task offers other tools than its environment class, since the model was shown the class's tools. It means the bundle's toolsets don't separate those tasks; see [bundles with several toolsets](#bundles-with-several-toolsets). ## Bundles with several toolsets A task's toolset is its played agent's envs and their versions, its role, the tool-access changes on it, and whether it works in a sandbox. When the bundle's tasks give the played agent different tools, `env.factory` is a dict of classes, one per toolset, and each row names its `environment`. TRL routes each row to its class: ```python title="train.py" env = AgentEnv.local("my-bundle") sorted(env.factory) # ['email', 'email+bash', 'gridsim'] env.dataset[0]["environment"] # 'email' ``` ## Deadlines and sizing | Setting | Default | What it bounds | | --------------------- | ------- | ----------------------------------------------------------------------------------------------------------------------- | | `tool_timeout` | 60 s | One tool call. The server ends a call that misses it, and the episode with it, unscored. | | `reset_timeout` | 600 s | A reset, through the deploy and the opening of the first turn, and each `end_turn`. Past it, the episode ends unscored. | | `max_concurrent_envs` | 8 | The sessions the server allows at once: one per rollout in flight. | A deadline ends only its episode: the session goes on to the next rollout. A server at its session limit refuses an environment's first reset, with the fix: ```text title="Error" the server at http://env-host:8000 is at its limit of 8 sessions at once, and the trainer holds one per rollout it generates at once: serve the bundle with MAX_CONCURRENT_ENVS (or max_concurrent_envs=) at least that many, or generate fewer rollouts at once ``` ## Cleaning up `env.close()` ends the episodes the environments hold and stops a server `AgentEnv.local` started. When the trainer exits, or crashes, the same happens on its own: TRL never closes its environments. A local server gets up to seven minutes to tear its episodes down. A process killed outright tears nothing down, so its sandboxes run until their provider stops them. ## What's supported | | | | ----------------------------- | ---------------------------------------------------------------------------------------------------------------------------- | | `GRPOTrainer` | TRL 1.15 and later. A CI job trains a step through the helper on TRL 1.16 and the latest, every week. | | AsyncGRPO | Takes the same `environment_factory`, trained on GPUs with TRL 1.16. | | Multi-turn tasks | `end_turn` hands the model each next prompt. | | Bundles with several toolsets | TRL's multi-environment mode, set up for you. | | Sandboxes | `sandbox=` deploys every episode where `agent-env run --sandbox` would, from your Docker to [any provider](https://www.agentenvframework.com/docs/sandboxes.md). | ## Limits * **vLLM 0.29 or older,** until OpenEnv and agent-env move to mcp 2. * **One deploy per reset.** Every rollout deploys its envs fresh, and the deploy counts in its time. * **AsyncGRPO's context.** Size `--max-model-len` for whole conversations, as above.