TRL

How a rollout runsCopy Anchor Link

Three pieces connect a trainer to a bundle:

InstallCopy Anchor Link

You need Python 3.11 or newer, Docker running, and a GPU for training:

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.

Train on GridSimCopy Anchor Link

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 bundleCopy Anchor Link

The bundle is a folder in a Hugging Face dataset. The hf command comes with the install:

hf download earakely-scale/GridSimEnv-AgentEnv --repo-type dataset --include "bundles/gridsim-v1/*" --local-dir .

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 handCopy Anchor Link

AgentEnv.local serves the bundle in your process, on your Docker. The first start builds the env's image, in about a minute:

from agent_env.training.trl import AgentEnv

env = AgentEnv.local("bundles/gridsim-v1", split="train")
len(env.dataset), len(env.dataset_for("eval"))

(2000, 200)

Then play one rollout the way TRL will. The task leaves the model one tool, submit_schedule:

environment = env.factory()
environment.reset(**env.dataset[0])
environment.submit_schedule(actions=[0] * 23)
environment.submit_schedule(actions=[0] * 24)
environment.get_reward()

'{"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 SFTCopy Anchor Link

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.

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 GRPOCopy Anchor Link

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:

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.

{"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 learnsCopy Anchor Link

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 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 hostCopy Anchor Link

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:

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)

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 instead of the host's Docker.

What the trainer getsCopy Anchor Link

env.dataset holds a row per task of the split, and env.dataset_for("eval") another split's rows:

env.dataset[0]

{
  "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:

import inspect

[name for name in vars(env.factory) if not name.startswith("_")]
inspect.signature(env.factory.submit_schedule)

['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 trainerCopy Anchor Link

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 AsyncGRPOCopy Anchor Link

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:

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

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()

CUDA_VISIBLE_DEVICES=1 python train_async.py

{"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.

What the model seesCopy Anchor Link

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:

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.

RewardsCopy Anchor Link

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:

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:

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 toolsetsCopy Anchor Link

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:

env = AgentEnv.local("my-bundle")
sorted(env.factory)               # ['email', 'email+bash', 'gridsim']
env.dataset[0]["environment"]     # 'email'

Deadlines and sizingCopy Anchor Link

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:

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 upCopy Anchor Link

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 supportedCopy Anchor Link

LimitsCopy Anchor Link

• 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.