[earakely-scale]

What

Train a bundle's own agents (claude-code, codex, opencode, any harness a task deploys) with TRL's harness rollouts: trl.experimental.async_grpo.HarnessRolloutWorker with harness_adapter=None, where the agent owns its loop. Each GRPO rollout runs one bundle task the way agent-env run does: its envs, its deployed agent, its prompts, its verifiers. The agent's model calls go through OpenEnv's capture proxy to the trainer's vLLM. The proxy records the exact token ids and log-probs TRL trains on, and the task's verifiers give the reward.

Pass functools.partial(AgentEnvSessionFactory, bundle, llm_url=..., model=...) as HarnessRolloutWorker's harness_session_factory (from trl.experimental.async_grpo.openenv_harness). Give the trainer Dataset.from_list(prompt_rows(bundle)) as its dataset and reward_funcs=[]. The docs page will carry the full recipe once the GPU run below pins it.

agent_env.training.harness (extra harness: OpenEnv 0.10 or later; TRL and vLLM are the trainer's to pin):

• prompt_rows(bundle, ...): the dataset, a row per task. TRL hands the session factory only a row's prompt, so each row names its task in a system line. The model never sees that prompt, since the task's agent builds its own.
• AgentEnvSessionFactory(bundle, llm_url=, model=, ...): OpenEnv's ResourceSessionFactory, built in TRL's rollout process. It does four things:
  • prepares the bundle as agent-env run does;
  • refuses an engine that can't return token ids (OpenEnv's require_llm), rather than producing untrainable rollouts;
  • serves the capture proxy on a free port, where local Docker sandboxes reach it, with its session routes behind a random admin key, so a sandbox can't list or read other rollouts;
  • opens one capture session per rollout, keyed by an API key the agent sends.
• Policy versions: a session keeps the model name TRL published when it was created (trl-policy-vN with LoRA), even after TRL moves on.
• Reward: reward= works as in build_app. A run that fails is None, which TRL leaves out of the update rather than scoring 0.
• Stopping: TRL closes the sessions still in flight when it stops. A session runs Task.run as an asyncio task of its own, which close() or the rollout's timeout cancels, and tears down in its own coroutine, which nothing cancels. So close() returns only once the run is torn down, no sandbox outlives the trainer, and a stop can't cut a teardown short. A second close() sends no cancel.
• Local Docker only: every deploy runs there; a remote sandbox can't reach the proxy on the trainer's loopback without a tunnel, so there's no sandbox argument.

Only the solver calls the policy

A task's rubric judge, its user simulator and envs that call a model must keep the configured endpoint. Otherwise a judge would ask the policy through the proxy, and its tokens would be trained on. So the factory doesn't use the run-wide litellm_base_url / litellm_api_key overrides:

• deploy_agent reads litellm_base_url and litellm_api_key from its own step_params over the run-wide overrides.
• The factory sets them for the task's solver deploys (solver_steps, the agent a step prompts).
• agent= swaps the solver's agent exactly as agent-env run --agent does: prepare_bundle takes agent, and solver_params (was _solver_params) gives the step params.
• prepare_bundle(solver_endpoint=True) tells preflight that the caller gives those deploys their endpoint. Training needs no endpoint configured, while a task with an agent judge still does.
• Left out: the tasks to train are chosen from the plan before the bundle is prepared (planned_tasks, new). A task whose solver no deploy_agent step deploys is left out with a warning, and so is one whose rubric judge is its solver. Nothing a left-out task alone needs can stop the rest.

Also

• Reward rules: reward_of, the reward rules build_app applies, moves out of app.py into agent_env.openenv.reward, which the harness imports.
• CI: harness.yml runs the harness tier on the locked and on the latest OpenEnv, whose capture internals (session registry, export, training trace) the harness reads. Like the TRL and verl jobs, it isn't a required check.

Rollouts waiting for a restarting env server now say so, once per server (_openenv.reopened); before, they waited silently.

Tests

• harness_test.py drives a bundle with real deploy_agent / prompt_agent configs against a stand-in vLLM (token ids, processed logprobs), with no model endpoint configured:
  • the solver alone gets the proxy, and the run-wide overrides carry none;
  • a task deploying an unprompted agent, and one whose judge is its solver, are left out without blocking the rest;
  • the call is captured as one trainable turn and scored;
  • the session keeps its policy version;
  • agent= swaps the solver;
  • a failed run is unscored;
  • a prompt naming another bundle's task is refused, and so is a bundle with no solver to train;
  • an engine without token ids is refused.
• tst/integration/training/harness_test.py (slow tier): the shipped opencode and claude-code run their smoke tasks in Docker through the proxy, against the stand-in (tst/util/fake_vllm.py), which has them run the shell command and then answer. With no endpoint configured, the rollout scores 1.0 and every captured turn trains.
• preflight_test.py: with solver_endpoint, the judge still needs the configured endpoint.
• test_step_overrides_agent_and_model.py: deploy_agent's endpoint is its step params', over the run-wide overrides, then its own.
• harness_test.py also covers stopping: close() on a rollout mid-task returns once it's torn down, with the run unscored; a rollout past its timeout is torn down before TimeoutError; and a session closed before it runs never runs.
• sessions_test.py: the restart warning appears once.

No-GPU check

On a Linux EC2 box with Docker, using TRL 1.17.1 and OpenEnv 0.10.0 (CPU torch). The real claude-code and opencode harnesses run their shipped smoke tasks (agents/tasks/*-smoke.json: run echo $((6 * 7)), answer 42) in local Docker, two rollouts at once per task, as a GRPO group runs in TRL's session pool. Their calls go through the capture proxy to a scripted stand-in for vLLM that returns token ids and processed logprobs. It answers the first call with a shell tool call and the call after the tool result with "42". Each capture then goes through TRL's own no-adapter path: _turns_from_training_trace, to_trace_entries, _messages_from_trace and _chain_to_sequences with fork_threshold_tokens=0.

	claude-code, 2 at once	opencode, 2 at once
exit, reward	0, 1.0 each	0, 1.0 each
turns captured	2: Bash call, then "42"	2: bash call, then "42"
TRL sequences	2, every completion token trained, a logprob per token	same
one rollout	~27 s	~24 s

• Sampling: TRL's training sampling (temperature 1, top_p 1, top_k -1, min_p 0, repetition_penalty 1) reached the engine on every agent call.
• claude-code: its Anthropic /v1/messages calls were translated to chat completions.
• Tally: the fork in TRL's drift tally comes from the stand-in's token rendering, not from the capture.
• Empty trace: a run that calls no model converts to 0 sequences without raising, so it can't stop TRL's worker.

The first version of this PR used the run-wide model overrides, and the check found that preflight refused the smoke tasks with no endpoint configured. Fixing that turned up the judge problem above.

Not yet

• GPU run: vLLM plus AsyncGRPO on a real model and a task with a learning signal. It follows review of this design.
• Remote sandboxes (Modal, E2B) can't reach a proxy on the trainer's loopback without a tunnel; local Docker only for now.

🤖 Generated with Claude Code

Confidence Score: 5/5

The changes since the previous review appear safe to merge.

Summary

Adds bundle-agent training through OpenEnv’s capture proxy for TRL harness rollouts.

• TRL rollouts run bundle tasks with their own agents and capture the solver’s calls.
• Bundle runs can route a solver to its own model endpoint.
• OpenEnv and training share the same task reward rules.
• Rollouts warn once when a server is restarting.

Since the previous review, the reward helpers moved into agent_env.openenv.reward without changing their behavior. No new actionable issues were found.

Diagram

sequenceDiagram
    participant Trainer as TRL
    participant Session as AgentEnvSession
    participant Agent as Task solver
    participant Proxy as Capture proxy
    participant Model as vLLM
    Trainer->>Session: Create rollout
    Session->>Agent: Run task with session endpoint
    Agent->>Proxy: Model calls with session key
    Proxy->>Model: Request pinned policy
    Model-->>Proxy: Tokens and log probabilities
    Proxy-->>Agent: Model response
    Session->>Session: Record scores and tear down
    Session-->>Trainer: Training trace and reward

Loading

_{Reviews (7) · Last reviewed commit: "refactor(openenv): keep the reward rules..." · Reviewed by Greptile}