Skip to content

feat(training): train a bundle's own agents with AReaL's agent workflows - #438

Open
earakely-scale wants to merge 6 commits into
mainfrom
feat/areal
Open

earakely-scale wants to merge 6 commits into
mainfrom
feat/areal

Conversation

@earakely-scale

@earakely-scale earakely-scale commented Nov 2, 2026 •

Copy link
Copy Markdown
Collaborator

What

agent_env.training.areal trains a bundle's own agents with AReaL's agent workflows (rollout.agent.mode: inline), the way agent_env.training.harness does for TRL's harness rollouts.

  • AgentEnvWorkflow is the inline agent AReaL's OpenAIProxyWorkflow runs for each sample (workflow: agent_env.training.areal.AgentEnvWorkflow, with the bundle and selection as workflow_kwargs). It plays the row's task as agent-env run does, on local Docker. The solver's deploy alone gets AReaL's proxy address and the sample's session key as its model endpoint, so the proxy records the tokens AReaL trains on. The task's other agents and judges keep the configured endpoint. run returns the task's reward as a float.
  • task_rows(bundle, ...) is the trainer's dataset: a row per trainable task, with its id and first prompt as messages.
  • A run that fails, or records no score its reward reads, raises.
    • On AReaL after 2.1.0, should_reject_failed_sample leaves that one sample out of its group, and the group's other samples still train. The exception is when AReaL itself found the model at fault (context overflow, invalid tool JSON); AReaL keeps those as reward 0.
    • AReaL 2.1.0 has no such hook and rejects the whole group.
  • When AReaL cancels a run, run stops the task and passes the cancel on only once the run is torn down, so no sandbox outlives a cancelled sample.
  • Runs play on the event loop the training sessions share, on a thread of its own, so a run's blocking calls never stall AReaL's loop.
  • AReaL may build the workflow again on every submit, so a process prepares the bundle only once per set of arguments.
  • No extra and no AReaL import: AReaL calls the workflow by its methods alone. It installs into AReaL's own environment with a core pip install agentenv-framework.

Shared with the harness helper

  • _solver.py holds what the two helpers share, moved from harness.py with no behaviour change: which tasks can train their solver, their preparation for local Docker, and the overrides that point the solver alone at the trainer's endpoint.
  • The reward rules move from agent_env.openenv.reward to agent_env.training.reward. agent_env.openenv's package init imports the OpenEnv server, so nothing under it imports without the openenv extra, which this helper doesn't need. build_app imports the rules from the new place. (agent_env.openenv.reward first shipped in 0.9.1408.)
  • The session loop moves from _openenv.py, which imports openenv, to _loop.py. The trl, skyrl and harness helpers import it from there.
  • The training conftest now shares the harness tests' stand-in bundle, steps and model. Their registries fixture is no longer autouse: the bundle fixture requests it.

Tests

  • tst/unit/training/areal_test.py checks these behaviours against the harness tests' stand-in steps and model:
    • the solver alone is pointed at the proxy, under the session key;
    • the reward comes back as a float;
    • agent= swaps the solver;
    • a failed run raises, and so does an unscored one;
    • a cancelled run is torn down before the cancel returns;
    • a row naming no task is refused;
    • the bundle is prepared once however often AReaL builds the workflow;
    • the reject hook gives the right answer.
  • New tier tst/areal, run by .github/workflows/areal.yml on AReaL 2.1.0 and on main with CPU torch over the dev environment. It is not a required check. AReaL's own rollout engine resolves the workflow from its import path into its grouped proxy workflow. Each sample plays under its own session, its reward goes on its last completion, and a failing sample is rejected (alone on main, with its group on 2.1.0).
  • tst/integration/training/areal_test.py: opencode, in local Docker, plays its smoke task against a stand-in model at a proxy-style address without /v1, and the run returns reward 1.0. The stand-in now listens on every interface, as AReaL's proxy does, since a container reaches the host only through the bridge gateway. It no longer needs openenv.
  • make areal-test, and the AGENTS.md rows for the tier, its CI job and the training package.

Run on a Linux box:

  • the unit tier: 12,490 passed and 2 failed. The two failures are test_mcp_cli_builder's generated CLIs, which import httpx under the box's system python3; they are unrelated to this change.
  • the AReaL tier: 2/2 on 2.1.0 and 2/2 on main.
  • the AReaL integration test, the harness integration tests (opencode and claude-code) and tst/unit/training again on the final stand-in: 100 passed.

A GPU proof (GRPO through AReaL on a bundle's own opencode agent) follows in a comment.

🤖 Generated with Claude Code

RetriggerView in GreptileConfidence Score: 5/5

The PR appears safe to merge; both previous findings are addressed and no new blocking issues were found.

Summary

Adds AReaL training for a bundle’s own agents, with shared task preparation, rewards, and a background event loop.

  • AReaL can train a bundle’s agents by running its tasks as inline workflows.
  • TRL and AReaL share task setup and session-loop helpers.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[AReaL sample] --> B[AgentEnvWorkflow]
    B --> C[Run task on shared loop]
    C --> D[Solver calls AReaL proxy]
    D --> C
    C --> E[Reject unrecovered failures]
    E --> F[Read scores]
    F --> G[Tear down run]
    G --> H[Return reward]
Loading

Reviews (2) · Last reviewed commit: "fix(training): leave out a run with a st..." · Reviewed by Greptile

earakely-scale and others added 5 commits November 1, 2026 21:19
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ches it

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…aL's proxy does

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@earakely-scale
earakely-scale requested a review from a team as a code owner November 2, 2026 05:38
Comment thread src/agent_env/training/areal.py Resolved
Comment thread src/agent_env/training/areal.py Resolved
… bundle without the reward

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@earakely-scale

Copy link
Copy Markdown
Collaborator Author

GPU proof setup: one g6.12xlarge (4×L4), AReaL main at d25f292 (SGLang 0.5.10, torch 2.9.1), Qwen3-4B-Instruct-2507 with LoRA, SGLang on one GPU and the FSDP actor on two, the framework at 62d9196.

Plumbing: works end to end.

  • AReaL resolved agent_env.training.areal.AgentEnvWorkflow from its import path and ran it in inline mode.
  • The shipped opencode and Claude Code smoke tasks went through AReaL's real proxy: 12 of 12 runs scored 1.0 (6 each, 90–111 s a run).

Training: no learning shown in this run. It shows rollouts, rewards, rejections and updates working, not a gain.

  • The base model is nearly bimodal on the arithmetic bundle at AReaL's sampling (temperature 1, no top-p or top-k). Product and sha score 1.0, factorial-digits and modpow score 1 in 8, and the other four score 0.
  • Five GRPO steps on the two in-between tasks played 104 rollouts:
    • 63 were scored;
    • 41 failed and were each left out of their group through should_reject_failed_sample, the rest of the group still training.
  • Every trained batch scored 0, so the updates had nothing to learn from.
  • The sixth weight sync (LoRA from disk into SGLang) timed out after 600 s.

What AReaL needed along the way, for the docs:

  • Use AReaL main for opencode-style agents. 2.1.0's proxy serves OpenAI chat completions only at /chat/completions, so agents that append /v1 to their base URL get a 404. opencode is one. AReaL main also serves /v1/chat/completions. Claude Code's /v1/messages works on both.
  • Reinstall FA2 after uv sync --extra cuda. The sync installs flash-attn-4 and removes the FA2 wheel AReaL's install guide preinstalls. Reinstall FA2 with --no-deps afterwards.
  • Fit the actor on 24 GB GPUs with actor.optimizer_dtype: bfloat16 and actor.fsdp.memory_efficient_load: true. By default the FSDP actor builds the model in float32 and runs out of memory.
  • Give the proxy an admin key. It listens on the host's address, where the agents' containers reach it, and refuses the default key: ++rollout.agent.admin_api_key=$(openssl rand -hex 24).
  • Keep rollout.agent.engine_max_tokens below sglang.context_length. SGLang refuses a request totalling exactly its context length. Claude Code's ~15k-token prompt hit this at the defaults (both 32768).
  • Possibly turn the decoupled loss off. The behaviour log-probs recorded through the proxy came out far below the actor's own (behave_imp_weight around 1e-5), which zeroes the decoupled loss. The training run used use_decoupled_loss: false.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant