AsyncGRPO: train loop-owning agents on the engine's own token ids, with a Harbor/OpenEnv example - #6947
AsyncGRPO: train loop-owning agents on the engine's own token ids, with a Harbor/OpenEnv example#6947adithya-s-k wants to merge 21 commits into
Conversation
…hrough OpenEnv Trains against a Harbor task through mini-swe-agent running in an E2B sandbox. The agent owns its own loop; TRL stands up an endpoint, lets it drive, and reads back the captured token ids and logprobs. That is what makes an installed harness trainable without reimplementing it, and it is the difference from examples/grpo_harbor, which runs Harbor tasks against harnesses written inside TRL with TRL owning the loop. mini-swe-agent is the default on measured grounds rather than taste: across a 15-harness sweep on the same 50 tasks it was the most accurate and the most turn-efficient, its prompt re-render is byte-exact against the engine's prompt_token_ids, and it is the only harness that can express a step limit. The re-render matters because TRL rebuilds each prompt locally, and for three of twelve harnesses measured that drifts (claude-code +2 tokens, gemini-cli +2, kimi-cli -10 per tool call) -- invisible for eval, forking the trajectory every turn when training. The step limit is not a cost control. Every turn re-sends the whole conversation, so a rollout's packed length grows with the SQUARE of its turn count; unbounded 58-turn rollouts were enough to OOM the loss step on an 80 GiB card. The docstring states one caveat rather than hiding it: on this path HarnessRolloutOutcome carries a single verifier scalar, not the component dict, so a 'submission' term giving partial credit is unavailable and the reward is all-or-nothing. On a suite the model solves ~16% of the time that means most groups score identically and those steps teach nothing, so the example says to shape component rewards where the suite emits them and otherwise to pick tasks the model solves sometimes.
…d note The grpo_harbor reference would dangle once that example is deprecated, and a docstring should not point at something being removed. The reward paragraph is cut to what it is -- correctness plus a correctness-gated efficiency term, with --reward-key for suites that emit a dict -- rather than a discussion of what the single-scalar path cannot express.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 025f3309ee
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| processing_class=tokenizer, | ||
| # Must match how the engine was served, or every prompt is re-rendered under a different template | ||
| # than the rollout was generated with — silent skew, not an error. | ||
| chat_template_kwargs={"enable_thinking": False}, |
There was a problem hiding this comment.
In the loop-owning mode selected here, HarnessRolloutWorker rebuilds prompts through openenv_harness._turns_from_trace, which calls tokenizer.apply_chat_template without the worker's chat_template_kwargs. Consequently this argument is a no-op: the documented vLLM command generates with enable_thinking=false, while the captured turns are locally re-rendered using the tokenizer default, silently pairing generated tokens and old logprobs with different prompt IDs. Pass these kwargs through the loop-owning trace reconstruction before relying on this setting.
Useful? React with 👍 / 👎.
| max_inflight_tasks=args.max_inflight, | ||
| vllm_server_url=args.vllm_url, | ||
| max_tokens=args.max_completion_length, | ||
| temperature=args.temperature, |
There was a problem hiding this comment.
When harness_adapter=None, the external agent submits its own requests and this worker never calls _generate_one_turn; the worker temperature is only included in HarnessRunLimits, which _run_session passes to the white-box branch. Thus --temperature does not control these rollouts, while AsyncGRPOTrainer still recomputes policy logprobs using that value. For any harness whose request temperature differs, the captured processed logprobs and trainer logprobs represent different distributions, corrupting the importance ratios even before weights become stale.
Useful? React with 👍 / 👎.
| # Trackio keys a run by name inside a project, so two relaunches of the same config land on top of | ||
| # each other and the earlier metrics read as part of the later run's history — worst exactly when | ||
| # relaunching after a crash. Stamping the name keeps them apart. | ||
| stamp = os.environ.get("SLURM_JOB_ID", "local") |
There was a problem hiding this comment.
For every non-SLURM invocation that omits --run-name, this always produces the same ...-local run name for a given configuration. That is exactly the relaunch case the preceding comment intends to prevent: Trackio histories are merged and the default output directory is reused after a local crash or repeated experiment. Use a per-invocation timestamp or unique identifier when SLURM_JOB_ID is absent.
Useful? React with 👍 / 👎.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Brings the whole stack up inside one container: openenv harbor serve on CPU, vllm serve on GPU 0, the trainer on GPU 1. Self-contained rather than importing its sibling, because hf jobs uv run uploads a single script. Everything lives in one job because AsyncGRPO syncs weights into vLLM over NCCL, which needs both on the same host's GPUs. That leaves only the OpenEnv server placeable, and keeping it here puts the proxy's hop to vLLM on localhost; hosting it on a Space via `openenv harbor push` works but adds a public hop to every model call. The tunnel is not a workaround for missing infrastructure. Jobs can publish a port at <job_id>--<port>.hf.jobs, but access needs an HF token, and the agent's Authorization header already carries its rollout session key -- that key IS how the proxy routes concurrent rollouts, so it cannot carry a second credential. An unauthenticated outbound tunnel needs no ingress at all. The readiness check verifies the tunnel SERVES THE PROXY, not that a port answers. A tunnel whose forwarding process dies keeps resolving and returns the provider's error page; agents then get HTML where an OpenAI endpoint should be, make zero model calls, and every rollout comes back unscorable while the server's own /health still reports healthy. The URL is only accepted after a request through it returns our health document. The parser was replayed over 32 real server logs and matches all 12 published URLs, gradio and cloudflare alike. The mounted bucket holds checkpoints and HF_HOME, and says so loudly when absent rather than losing them silently. Deliberately not the sandbox templates: those are keyed by image hash on the provider's side and already persist across jobs, which is the expensive warm step.
…itten
--split was required, so the reproduction command carried a placeholder nobody could paste. It now
defaults to a public Harbor suite (AdithyaSK/data_agent_rl_environment_train, verified public), which
means the whole thing is one copy-pasteable line with no <angle brackets> in it:
hf jobs uv run --flavor h200x2 --image huggingface/trl \
--secrets HF_TOKEN --secrets E2B_API_KEY \
https://raw.githubusercontent.com/.../async_grpo_harbor_hf_jobs.py
The docstring shows that form first and the bucket form second, since the bucket is optional -- without
it the job still trains, it just loses its checkpoints when the container goes away, which the script
already warns about at startup.
Also fixes a flag that could not do anything: --no-enable-thinking was declared with store_false over a
default of False, so it could only set what was already set. Thinking is off by default because the
trainer's chat_template_kwargs must agree with how the engine was served, and --enable-thinking is now
the switch that changes it.
Found by auditing the script against the Jobs environment rather than the cluster it was written on.
1. vllm was never declared. `uv run --script` builds an isolated environment, so the base image's vllm
is not guaranteed importable -- and the script's own PATH guard would then have blamed "the PEP 723
dependencies did not install", which would have been true but unhelpful.
2. No sandbox template warm. This is the one that would have looked like a harness bug. Harbor decides
whether to build from alias_exists(), which flips true when a build STARTS, so num_generations
rollouts racing their first visit to a task all see "exists" and fail against a half-built image with
404: tag 'default' does not exist. On a cold template with 8 concurrent generations that is the
common case. One serial `openenv harbor rollout` first makes the build serial and every later rollout
finds a finished image; the provider keys images by content hash, so later jobs pay nothing.
Non-fatal on failure: the warm rollout can fail for reasons that say nothing about training, and the
build it triggered still happened.
3. No GPU-count guard. On a one-GPU flavor the trainer was pointed at device 1 and would have failed
deep inside CUDA instead of at startup with a sentence naming --flavor h200x2.
4. The run name fell back to a constant ("local") when HF_JOB_ID was absent -- and that variable's name
is not something this script can verify. Trackio keys a run by name inside a project, so the constant
would fold separate jobs into one history, worst exactly when relaunching after a failure. It now
falls back to a timestamp.
sergiopaniego
left a comment
sergiopaniego
left a comment
There was a problem hiding this comment.
quick review, first pass, we need to add the example to docs/source/example_overview.md
Restructured to match the pattern the opencode recipe already uses for Hugging Face Jobs (described in sergiopaniego's TRL x OpenEnv x Harbor post, with the launcher published as a gist): a small launcher wraps the processes a Job cannot start on its own, and the training script is DOWNLOADED rather than duplicated. That last part is the reason for the change. The previous commit added a second 489-line file that restated the whole training path, so the two could drift and a reader had to diff them to see what was actually different. Now async_grpo_harbor.py is the single canonical script, byte-identical whether it runs locally or in a Job, and launcher.py holds only what is genuinely Jobs-specific. Ours needs three processes where the opencode launcher needs two, because the Harbor dataset and the capture proxy are served by openenv harbor serve. It also tunnels a different hop: opencode tunnels vLLM, since its in-sandbox proxy calls the engine directly, whereas here the sandboxed agent calls the capture proxy and the engine stays entirely private on localhost. --train-script-url exists because the equivalent opencode launcher still points at examples/scripts/openenv/opencode_hf_sandbox.py, a pre-reorg path that now 404s. A stale URL surfaces as a download failure minutes into a paid job, so this one is overridable and can be pointed at a branch before the example is merged. Carried over from the audit of the deleted file: vllm is declared, the sandbox template is warmed by one serial rollout before any group runs concurrently, there is a GPU-count guard, the tunnel readiness check verifies a request THROUGH the tunnel reaches the proxy, and the run name never falls back to a constant.
Verifying once at startup is not enough, and this is measured rather than defensive. Over 27 hours on our own cluster the published tunnel stopped serving the capture proxy 69 times -- roughly once every 24 minutes -- so a run of any length loses it mid-flight. When that happens the sandboxed agent gets the tunnel provider's error page instead of an OpenAI endpoint, makes zero model calls, and every rollout comes back unscorable, with nothing in the trainer's logs to say why. A daemon thread re-probes the published URL and restarts the Harbor server after two consecutive failures, so one flaky request cannot bounce a healthy server mid-step. The restart changes the published URL, which is fine: the trainer talks to the server over localhost and the server hands its current URL to each new sandbox, so only in-flight rollouts are lost. Which check does the work is worth recording. The failure signature originally debugged -- the provider's "no interface is running" placeholder -- was 2 of those 69. The other 137 probe failures were plain 502s. So the test is positive and generic: the URL must return OUR health document. Enumerating known failure modes would have caught almost none of them.
All four test suites on huggingface#6947 fail on one assertion, and it is not a code failure: tests/test_examples_index.py::test_examples_index_matches_folders AssertionError: Example folders missing from the Index table in example_overview.md: {'async_grpo_harbor'} `test_examples_index_matches_folders` requires every directory under `examples/` to have a matching row in the index, so adding the example without the row turns all of "latest", "dev", "minimum versions" and "without optional dependencies" red at once. One row fixes all four. Placed before `async_grpo_math` to keep the table alphabetical, matching the neighbouring `async_grpo_opencode` entry's format.
| # 1. the Harbor dataset + the capture proxy, published for the sandboxed agent | ||
| server_log = logs / "openenv-server.log" | ||
| public_proxy = start_harbor_server(args, server_log) | ||
| threading.Thread(target=supervise_tunnel, args=(args, server_log, args.tunnel_check_s), daemon=True).start() |
There was a problem hiding this comment.
Zero interval never disables supervisor
Low Severity
--tunnel-check-s is documented as disabling the supervisor at 0, but the supervisor thread is always started. A zero interval then busy-loops probes and can restart a healthy Harbor server after two immediate failures.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 3ca92ca. Configure here.
|
Nice!
Also the docstring says "a 15-harness sweep" then "three of the twelve harnesses measured". which is it, and is the sweep written up anywhere we can link? Feel free to merge |
…prompt In the loop-owning path the agent runs its own tool loop in a sandbox and TRL reads back what it did. Rebuilding each turn's prompt with apply_chat_template to recover the token ids produces a DIFFERENT string from the one the engine scored. Measured on Qwen3.5-4B, the rebuilt prompt matched the engine on 0 of 28 turns. Nothing errors when that happens, which is the problem. Turns that should chain by exact token prefix instead look like unrelated short rollouts, so one conversation fragments into several and every fragment still trains. Two production runs spent a night each in exactly that state with healthy-looking reward curves; the tell was samples_per_rollout equalling turns_mean exactly, i.e. 100% forking. _turns_from_trace now takes `prompt_token_ids` from the trace entry and raises when it is absent, naming the two vLLM flags that produce it. A hard failure rather than a fallback on purpose: the fallback IS the bug, and it is invisible. An engine without --return-tokens-as-token-ids yields rollouts that look completely normal and carry nothing to train on, so the run has to stop at the first rollout instead of discovering it at the first weight update. The tokenizer argument is gone from that path -- there is nothing left for it to do. TurnRecord gains `output_mask`, and _SampleBuilder.append_turn honours it instead of hardcoding 1 across the completion. This is not cosmetic: capture masks a turn OUT when its logprobs were rejected on ingest, while keeping its tokens as context and zero-filling the logprobs. Without the mask those positions train against logprob 0.0, i.e. p=1.0 -- a confident target for a turn we explicitly could not trust. Requires the matching OpenEnv change, which puts `prompt_token_ids` and `loss_mask` on TraceEntry. Against an older OpenEnv this raises immediately and says which flags are missing, which is the intended behaviour rather than a regression. Note for anyone tempted to prototype this as a monkeypatch: it cannot work. The rollout loop runs in a multiprocessing child created with `spawn`, which re-imports every module from scratch, so a parent-side rebind of a module function is simply lost. A patch that mutates the tokenizer OBJECT survives, because that is pickled into the child; one that rebinds a function does not. That asymmetry is why the re-rendering went unnoticed for so long -- the patch logged its install line and its counters never moved.
adithya-s-k
changed the title
# Conflicts: # docs/source/example_overview.md
…unables, fix the harness counts From @qgallouedec's review: * THIRTY flags down to twenty-three. The seven removed (`--optim`, `--no-bf16`, `--gradient-checkpointing` / `--no-gradient-checkpointing`, `--save-steps`, `--save-total-limit`, `--seed`) only re-exported `TrainingArguments`. An example is meant to be read and edited, not driven like a product, so those are literals in the config now with a note where the choice is not obvious (gradient checkpointing is on because rollout sequences here are long enough that activations dominate). Checkpointing defaults to off, since a short example run has nothing worth keeping. * The two reward tunables are plain module constants (`W_TOOL_EFFICIENCY = 0.3`, `TOOL_BUDGET = 15.0`) rather than `os.environ.get`, matching the opencode example -- one channel for tunables. The one remaining `os.environ.get` reads `SLURM_JOB_ID`, which is a fact about the environment rather than a knob. * The harness counts were not contradictory, just under-specified: FIFTEEN harnesses were probed, TWELVE produced usable rollouts, and it is three of those twelve that had to re-render. Said explicitly now. Also resolves the `docs/source/example_overview.md` conflict with upstream: upstream added an `async_grpo_math` row while this branch added `async_grpo_harbor`. Both are kept, with upstream's version of the shared row. `tests/test_examples_index.py` passes -- that gate fails all four suites at once when a row is missing, which is what @sergiopaniego's review was about.
| "a captured turn carried no `prompt_token_ids`. Serve the engine with " | ||
| "`--return-tokens-as-token-ids --logprobs-mode processed_logprobs`; without them every prompt trained " | ||
| "on here would be a local re-render, which matched the engine on 0 of 28 measured turns." | ||
| ) |
There was a problem hiding this comment.
Missing token ids fail silently
High Severity
_turns_from_trace raises when a captured turn has no prompt_token_ids, but _run_session catches every Exception and returns an unscorable empty rollout. A misconfigured engine therefore keeps the worker alive while every rollout is dropped, which is the silent multi-hour failure this raise is meant to stop.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 9c4c8fa. Configure here.
CI's code-quality job is `pre-commit run --all-files`, and its `doc-builder-style` hook rewraps docstrings to max_len 119. My `_turns_from_trace` docstring wrapped at a narrower width, so the hook modified a file and the job failed. Worth recording for next time: running ruff locally does NOT predict this job. ruff passed on all five of my changed files while the real failure was doc-builder rewrapping one line. The gate is pre-commit, so reproduce it with pre-commit (or at least the pinned doc-builder rev) rather than with ruff alone. The notebook errors ruff reported locally were a red herring -- those files are byte-identical to upstream and I touched no notebook.
…_fn shapes it `rollout_reward_fn` replaces `env_reward` wholesale, and the column it feeds is already named `rewards/harness_reward`. So the moment a caller passes a shaping term, `reward` and `rewards/harness_reward` are the SAME shaped number and the raw verifier score is logged nowhere. That makes every shaping term look exactly like progress. A reward curve that climbs because the policy got faster is indistinguishable from one that climbs because it solved more -- which is the one question the run exists to answer. It also silently breaks comparison against every earlier run, because the headline number changed units without changing name. Records `rollout/correctness_mean` from the same `_rates` channel `rollout/turns_mean` already uses, so it needs no new plumbing and drains generically. Only when a `rollout_reward_fn` is actually set: with no shaping term `reward` IS correctness and a second identical series would just be noise.
|
Holding off on merging this, one thing I'd like to settle first.
Also, since the producer isn't written yet, merging this first sets that contract "by accident" (rather than by agreement). I'd wait until what defines No rush from my side |
| # policy got faster is indistinguishable from one that climbs because it solved more, which is the | ||
| # one question the run exists to answer. | ||
| self._rates["rollout/correctness_mean"][0] += float(env_reward) | ||
| self._rates["rollout/correctness_mean"][1] += 1.0 |
There was a problem hiding this comment.
Thread-unsafe correctness rate updates
Medium Severity
rollout/correctness_mean is incremented inside _run_session on the thread pool, while _push_metrics iterates and clears _rates on the event loop. _generate_one already documents that those accumulators are not safe off-loop. A concurrent insert during items() can raise RuntimeError and abort the worker, and the same exception inside _run_session discards an otherwise valid rollout as unscorable.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 8e87edb. Configure here.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
There are 4 total unresolved issues (including 3 from previous reviews).
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit a28a3d1. Configure here.
| if value is not None | ||
| }, | ||
| } | ||
| ) |
There was a problem hiding this comment.
Sampling check rejects default policy keys
High Severity
The loop-owning consumer treats a missing captured sampling key as a mismatch. It builds the trainer policy with default repetition_penalty=1.0 (anything not None), while HarborSessionFactory is only given temperature, top_p, and top_k. Captured metadata then fails the check, _turns_from_trace raises, and _run_session swallows that as an unscorable rollout, so a Harbor run can complete with no trainable rows.
Additional Locations (2)
Reviewed by Cursor Bugbot for commit a28a3d1. Configure here.
| # A CAPTURED turn is different: the producer may mask a turn out while keeping its tokens as | ||
| # context -- e.g. its logprobs were rejected on ingest -- and that is not inferable here. | ||
| # Training those positions anyway means training against a logprob of 0.0, i.e. p = 1.0. | ||
| output_mask: list[int] | None = None |
There was a problem hiding this comment.
Is this really needed outside harness support ?
trl/trl/experimental/async_grpo/openenv_harness.py
Lines 128 to 129 in b7df6a8
these functions are basically made to expose the same mechanism. If OpenEnv changed to produce output_mask we might need to completely change harnessWorker
| ) | ||
| reward = self._rollout_reward_fn(outcome) if self._rollout_reward_fn else env_reward | ||
| if self._rollout_reward_fn is not None and env_reward is not None: | ||
| # Keep the VERIFIER's score visible on its own. `rollout_reward_fn` replaces `env_reward` wholesale, |
There was a problem hiding this comment.
I'd prefer we cleanup these LLM generated comments in general. The provide very low signal and make it hard to focus on code IMHO
4 participants
An installed coding agent can own its tool loop while AsyncGRPO trains on the exact tokens its inference engine sampled. This adds a Harbor/OpenEnv example and captured-trace ingestion.
The loop-owning path validates engine prompt/completion IDs, sampled log probabilities, masks and sampling policy. If a harness rewrites history, lossless capture starts a new row instead of overwriting sampled tokens. White-box reconciliation keeps its existing behavior.
examples/async_grpo_harbor/includes local commands and a Hugging Face Jobs launcher. The example keeps verifier scores separate from its correctness-gated efficiency reward. Weight updates drain inference before replacing weights; empty successful cache-reset responses and tool-only messages are handled. Training requires engine token IDs and processed logprobs.Validation: current main merged; all CI passes, including distributed GPU smoke. Local checks include 245 CPU tests, three HTTP control tests and pre-commit. TiTO regressions cover rewritten histories, masks, token retention, invalid logprobs and policy mismatch. Pinned-runtime Harbor/OpenCode, native OpenCode and SETA GPU smokes each passed four optimizer steps and native checkpoint save/remote restore; receipts are in adithya-s-k/HuggingEnvs#7.
Merge after OpenEnv #1036, which provides the capture API. Lossless forks can increase row/token budgets; rollout-normalized weighting and equal harness contributions are not implemented here.