[earakely-scale]

Summary

TRL's trainers can now train on a bundle served by agent_env.openenv.build_app, through TRL's own environment_factory, with no changes to TRL.

• GRPOTrainer is verified with TRL 1.16, including on a GPU (below).
• AsyncGRPO takes the same environment_factory, but hasn't been run yet.
• Either needs TRL 1.15 or later.

from trl import GRPOConfig, GRPOTrainer
from agent_env.training.trl import AgentEnv      # pip install 'agentenv-framework[trl]'

env = AgentEnv("http://localhost:8000")          # or AgentEnv.local("./bundle", max_concurrent_envs=64)
trainer = GRPOTrainer(model="Qwen/Qwen3-4B", args=GRPOConfig(max_tool_calling_iterations=12),
                      train_dataset=env.dataset, environment_factory=env.factory)

What's new:

• agent_env/training/trl.py, behind a new trl extra (agentenv-framework[openenv] and datasets). TRL itself stays the trainer's to pin. Later trainers' modules go in agent_env/training/ too, each behind its own extra.
• agent_env/training/_tools.py turns an env's tools into typed, documented Python methods from their JSON Schema, in the form transformers' get_json_schema reads, which is how TRL shows tools to a model.

env.dataset has one row per task of the split, read from the server's Task API:

• the prompt as chat messages, which the user may change;
• the split and index that reset plays;
• the task_id that each rollout is checked against.

env.dataset_for(split) reads any other split.

env.factory is a generated class. Its methods are every tool the tasks' turns offer, plus end_turn.

• Probe: to list the tools, building it walks one task of each toolset through its turns, ending each but the last with an empty reply, then ends the episode unscored, so its verifiers never run.
• Sessions: each instance holds one session on the server, over OpenEnv's WebSocket (GenericEnvClient), and plays one episode at a time.
• Reset: reset returns at once, because GRPOTrainer resets one rollout at a time. The first tool call or get_reward waits for the deploy.
• Turns: end_turn answers with the next turn's prompt, preceded by its system prompt when that changed. A turn the model leaves open ends with an empty reply. TRL hands an environment no completion, so the model's reply reaches the task only through end_turn(response=...).
• Tool deadlines: the client sends a tool's deadline in the call's action metadata (timeout_s), and the server ends the call at it. Cancelling the call doesn't stop what it does, so the server ends the episode too: the rollout is left unscored (None), and the session is kept for the next one.

The environment owns the reward through get_reward, which both trainers read, so no reward_funcs are needed.

• It returns None when the episode ended in an error or played a different task than its row.
• It raises when a rollout's tools don't match its class's. That's a grouping bug no rollout fixes, so training stops with the reason.

Fails early, with the fix in the message:

• no server at the URL;
• an unknown split;
• a tool or argument name TRL can't call;
• a server at its session limit for an environment's first session, with the error naming MAX_CONCURRENT_ENVS / max_concurrent_envs=.

A later session, opened in the background after one was dropped, is retried while the server is full. If it still fails, only that rollout is lost.

Changes to the server's Task API:

• toolset: each task's description names the played agent's envs, the env versions the task deploys, the agent's role, its modify_env_tool_access changes, and whether its turns offer bash. Tasks with different toolsets get one class each (factory is then a dict), and each row names its environment.
• turns: each task's description says how many turns it has.
• Tool deadlines: AgentEnvEnvironment honors a tool call's timeout_s metadata. A call that misses it ends the episode with an error.

Also in this PR:

• AgentEnv.local(bundle, ...) serves build_app in process on a loopback port. A trainer that exits or crashes ends the episodes it left running and stops that server. It waits up to 30 s for the sessions to close, then up to 7 minutes for the server to tear down their sandboxes: a teardown's three parts (changelogs, recordings, sandboxes) take at most 2 minutes each, plus a minute. It warns if the server is still tearing down. close() gives the server the same time. A hard kill (OOM, SIGKILL, a preempted node) runs no teardown, so those sandboxes last as long as their provider lets them.
• build_app has its app construction split into _app(served, ...), so the unit tests serve hand-built tasks through the same server.
• Dependency tests: they now allow an extra that includes another (trl includes openenv). A new check imports agent_env.openenv and agent_env.training.trl with nothing installed beyond their own extra.
• THIRD_PARTY_NOTICES.md adds the five runtime packages the extra brings: datasets, dill, multiprocess, pyarrow and xxhash.
• transformers joins dev, for the get_json_schema test.

OpenEnv 0.10 requires fastmcp<4, and so mcp 1. A trainer with OpenEnv installed therefore runs vLLM 0.29 or older, which TRL 1.16 allows. Once OpenEnv releases FastMCP 4 and agent-env supports mcp 2, the helper runs on vLLM 0.30+ unchanged.

What the docs site should cover

A "Train with TRL" page under the OpenEnv docs:

• Install: pip install 'agentenv-framework[trl]', with TRL 1.15+ installed by the user; vLLM 0.29 or older while OpenEnv needs mcp 1.
• Serve the bundle: openenv import and uv run server with MAX_CONCURRENT_ENVS, or build_app, or AgentEnv.local(bundle) in the training process.
• The trainer: AgentEnv(url, split=...), with env.dataset / env.dataset_for(split) as the datasets and env.factory as the environment_factory. No reward_funcs: the environment owns the reward.
• Session limit: the server must allow a session for every rollout the trainer generates at once.
• Reaching the task: the model's reply reaches the task only through end_turn(response=...). A turn left open is scored as an empty reply, so tell the model to end its turn with end_turn.
• Multi-turn tasks: end_turn returns the next prompt; new system instructions come first.
• Mixed bundles: tasks with different tools take TRL's multi-environment mode automatically.
• Timeouts: tool_timeout (a call that misses it ends the episode, unscored) and reset_timeout (deploying, and ending each turn).
• Unscored rollouts: what leaves a rollout unscored (None), and what stops training.

Verified on a GPU

GRPOTrainer with Qwen3.5-9B (LoRA) trained on the PortSim training pack, served by build_app on an EC2 g7e:

• 5 steps of 16 rollouts, plus an eval;
• 98 rollouts, with 0 unscored and 0 tool failures;
• 4 of the 5 steps with a non-zero gradient.

Five steps don't measure learning.

Tests

• tst/unit/training/trl_test.py runs against the real OpenEnv server over the Episode API's fake steps, driven the way TRL drives it. It covers:
  • the dataset and the class shape;
  • a full rollout on a reused environment;
  • multi-turn tasks, including a tool only a later turn offers and a changed system prompt;
  • episodes left unscored, and a row whose prompt the user changed;
  • a tool that misses its deadline, which ends the episode unscored and keeps the session, on a server allowing one session;
  • a session refused after one was dropped;
  • a rollout whose tools aren't its class's;
  • an end_turn that waits for a slow verifier;
  • capacity, and no server;
  • multi-environment mode;
  • a tool name TRL can't call;
  • local(), with teardown at close and at exit, waiting out a slow teardown.
• tst/unit/training/tools_test.py checks the generated methods, and that transformers' get_json_schema reads each one back as the env describes it.
• tst/integration/training/trl_test.py (slow tier, local Docker) has 2 tests:
  • rollouts on one reused environment score as the items verifier scores them, and a two-turn task hands over its second prompt;
  • a task whose agent works in its sandbox gets a +bash toolset of its own, plays through bash, and scores 1.0. Nothing is left running.
• The full unit tier: 11,156 passed, 66 skipped.

🤖 Generated with Claude Code

Confidence Score: 5/5

The PR appears safe to merge; no blocking finding remains.

Summary

Adds AgentEnv so TRL can train on bundle tasks served through OpenEnv.

• TRL trainers can use task prompts, tools, and rewards from an OpenEnv server.
• Task descriptions include each task's toolset and turn count.
• Tool calls that exceed their deadline end the episode.
• A training process can serve a bundle locally and stop its server when training ends.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Bundle tasks] --> B[OpenEnv server]
  B --> C[AgentEnv dataset and factory]
  C --> D[TRL rollout]
  D --> E[Tool calls and end_turn]
  E --> B
  B --> F[Episode reward]
  F --> D

Loading

_{Reviews (9) · Last reviewed commit: "fix(training): give a local server a tea..." · Reviewed by Greptile}

[socket-security]

Review the following changes in direct dependencies. Learn more about Socket for GitHub.

Diff	Package	Supply Chain Security	Vulnerability	Quality	Maintenance	License
	pypi/​transformers@​5.21.0
	pypi/​datasets@​5.1.0

View full report

[AlexS-scale]

Had Claude run a review, can you take a look please:

Must fix

1. A tool timeout can crash the training run (trl.py:292, 343, 257, 428)

• What happens: when a tool misses tool_timeout, _fail closes the client's session. But OpenEnv's /ws handler processes one message at a time, so the server doesn't see the disconnect until the hung call returns. That can take up to the server's 600s tool timeout, and the session slot stays taken until then.
• The crash: the environment's next reset opens a new session and hits the limit. reset raises, and GRPOTrainer doesn't catch errors from reset, so the run dies. The error message then tells the user to raise MAX_CONCURRENT_ENVS, which wasn't the problem.
• Repro: I reproduced it against the PR with a 15-line test: the existing _Hanging tool, tool_timeout=0.5, max_concurrent_envs=1, then a second reset.
• Why defaults hit it: local() defaults to 8 sessions, TRL typically runs 8 rollouts per process, and the error message advises sizing capacity exactly.
• Fix, either one:
  • Send the deadline in the action's metadata (OpenEnv actions have that field) and have AgentEnvEnvironment._call honor it. The server then cancels the call at the same moment, and the session can be kept instead of dropped.
  • Or raise the capacity error only on an environment's first session. A reconnect after a failure should retry in the background and lose that one rollout, not the run.

2. Checking a row by its prompt text silently unscores every rollout (trl.py:270, _prompt_of at 485)

• What happens: _begin requires the server's system prompt and prompt to match the dataset row exactly. A user who edits dataset["prompt"], for example by adding a formatting system prompt (common in GRPO), gets every rollout unscored. The only signal is a warning per rollout.
• Fix: compare the task id instead. The Task API already returns it, and the reset observation's metadata carries task.id.

Worth fixing (can be follow-ups)

3. Tool mismatches fail silently, rollout by rollout.
   • Mechanism: tasks are grouped by a toolset key built from the agent's envs, its role, its tool-access changes and bash (environment.py:51). A rollout whose first turn offers different tools than the probe saw is left unscored (trl.py:275).
   • What the key misses: it ignores env_version on deploy_env. It may also miss apply_server_config extensions, though I haven't confirmed those can change the tool list.
   • Why it matters: this is a deterministic configuration bug, yet training carries on with those tasks never scored. Fail loudly instead.
4. Startup cost.
   • Today: building factory plays a full episode per toolset, one after another, including the verifiers on empty replies, and every training process repeats it.
   • Cost: with heavy envs and several toolsets, that's minutes per launch.
   • Options: probe in parallel up to the session limit, or stop after the last turn's tool listing so the verifiers never run. That needs the Task API to say how many turns a task has.
5. The dependency test got looser (test_runtime_dependencies_are_imported.py). I added import datasets to agent_env/openenv/app.py and the test still passed, though pip install '[openenv]' would then break.
6. The integration test isn't marked int_test_slow. It builds a real image and starts a registry, using the same fixtures as tst/integration/openenv/build_app_test.py, which is marked slow. AGENTS.md puts real image builds in the slow tier.
7. Document that the model's final text never reaches the env. TRL's get_reward() takes no arguments, so a turn the model doesn't end with end_turn(response=...) is scored as an empty reply. Verifiers that judge the reply score 0 until the model learns to call end_turn. This will be users' first surprise.
8. Nits:
   • episode_loop, on_loop and tool_method are public names, but all lists only AgentEnv.
   • threading._register_atexit is a private CPython API. The docstring justifies it, but process exit can block for up to about 10 minutes.
   • The PR doesn't say what the docs site should cover, which AGENTS.md asks for.
   • Nothing tests against a real TRL install, and TRL marks environment_factory as experimental. A TRL release could break this with no CI signal, so a small optional contract test would be worth adding.

[earakely-scale]

Thanks for the thorough review, all addressed in 6b6cb5b:

1. Tool deadlines now travel in the call's action metadata and the server ends the call at it; the model gets "x didn't answer within Ns", the episode and session go on. Only an environment's first session raises on capacity; a reconnect after a drop retries while the server is full and otherwise costs one rollout. Your repro (hung tool, tool_timeout=0.5, max_concurrent_envs=1, second reset) is now a test.
2. Rows carry task_id; rollouts are checked against the played task, not the prompt.
3. Toolsets include deployed env versions, and a tool mismatch makes get_reward raise.
4. The Task API reports turns; the probe stops after the last turn's listing, so no verifiers run.
5. New test: agent_env.openenv / agent_env.training.trl import with only their own extra installed (catches your import datasets example).
6. Integration tests are int_test_slow.
7. Documented in AgentEnv and in the PR's docs-site notes.
8. Loop helpers are private, exit waits ≤30 s per stage, and the README section moved to the PR's docs notes. A real-TRL contract CI job is filed as a follow-up.