Async grpo OpenEnv harness rollout by AmineDiro · Pull Request #6420 · huggingface/trl

AmineDiro · 2026-07-16T13:39:34Z

WIP

local sandbox: https://gist.github.com/AmineDiro/efc1b8f5b654690311f6f7ae6df7ad67
opencode harness: https://gist.github.com/AmineDiro/4c236418255c8cc6c697eacddf886981

from __future__ import annotations

import argparse

from datasets import Dataset
from transformers import AutoTokenizer

from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer
from trl.experimental.async_grpo.openenv_harness import HarnessRolloutWorker

from opencode_harness import MODEL, build_dataset, build_factory


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default=MODEL)
    p.add_argument("--vllm-url", default="http://localhost:8000")
    p.add_argument("--num-generations", type=int, default=4)
    p.add_argument("--max-inflight", type=int, default=4)  # per-rollout uuid sandbox + free proxy port -> concurrent
    p.add_argument("--max-completion-length", type=int, default=16384)
    p.add_argument("--max-steps", type=int, default=10)
    p.add_argument("--n-prompts", type=int, default=32)
    p.add_argument("--project", default="async-grpo-opencode-mbpp")
    p.add_argument("--output-dir", default="async_grpo_opencode_mbpp")
    args = p.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    dataset = Dataset.from_list(build_dataset(n_prompts=args.n_prompts, seed=0))

    worker = HarnessRolloutWorker(
        harness_session_factory=build_factory(),
        harness_adapter=None,  # loop-owning: opencode runs its own loop; TRL reads the proxy trace
        model_name=args.model,
        dataset=dataset,
        reward_funcs=[],
        processing_class=tokenizer,
        num_generations=args.num_generations,
        max_inflight_tasks=args.max_inflight,
        vllm_server_url=args.vllm_url,
        max_tokens=args.max_completion_length,
        temperature=1.0,
        fork_threshold_tokens=1024,
        log_completions=True,
        num_completions_to_print=2,
    )

    config = AsyncGRPOConfig(
        output_dir=args.output_dir,
        save_strategy="no",
        per_device_train_batch_size=4,
        num_generations=args.num_generations,
        gradient_accumulation_steps=1,
        max_completion_length=args.max_completion_length,
        max_steps=args.max_steps,
        learning_rate=1e-5,
        vllm_server_base_url=args.vllm_url,
        report_to="trackio",
        project=args.project,
        log_completions=True,
    )

    trainer = AsyncGRPOTrainer(
        model=args.model,
        args=config,
        train_dataset=dataset,
        processing_class=tokenizer,
        rollout_worker=worker,
    )
    trainer.train()


if __name__ == "__main__":
    main()

Note

Medium Risk
Changes touch the experimental async GRPO rollout/scoring path and spawn external agent processes; default worker behavior stays backward compatible, but harness mode adds new failure modes (session create, trace parsing) that are mostly swallowed as unscorable rollouts.

Overview
AsyncGRPO can now train on rollouts produced by OpenEnv sessions (external agents that own their tool loop) instead of only TRL’s built-in vLLM tool loop.

A new HarnessRolloutWorker / _HarnessRolloutLoop runs each rollout via a ResourceSessionFactory: either white-box (OpenEnv HarnessAdapter + TRL sampling per turn) or loop-owning (wait for the agent, read the in-sandbox proxy trace for token ids/logprobs, then session.verify). Rewards can come from a rollout_reward_fn on HarnessRolloutOutcome without reward_funcs; hooks train_turn_fn / agent_turn_fn filter which trace turns become training rows.

The default AsyncRolloutWorker is extended so _generate_one returns a sixth rollout_reward (always None today), RolloutGroup carries rollout_rewards, and _score_group sums them when _provides_rollout_reward is set. The child process accepts a pluggable _loop_cls. RolloutWorkerProtocol allows multiprocessing.Queue as well as queue.Queue.

examples/scripts/openenv/opencode.py demonstrates end-to-end GRPO on DeepCoder with the real opencode CLI, a local subprocess sandbox (remapped /home/user), per-session proxy ports, held-out stdin/stdout verification, and custom reward/turn filtering.

Tests are updated for the new return tuple and rollout_rewards on groups.

^{Reviewed by Cursor Bugbot for commit 444ca7d. Bugbot is set up for automated code reviews on this repo. Configure here.}

bot-ci-comment · 2026-07-16T15:33:44Z

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.

qgallouedec · 2026-07-20T04:22:28Z

Thanks I think it looks good overall. I didn't try myself, but from reading the edits, nothing much to say.

Check the cursor review, it's likely the reason of the CI failure. Once the CI green I approve and we merge

sergiopaniego

thanks!!
approving but with 2 ideas:

we could add a short harness-rollout section to openenv.md in this PR: white-box vs loop-owning, --return-tokens-as-token-ids, reward via session.verify() → harness_reward.
we should add an example script. I think this could be added once we add support for HF sandbox + harnesses. wdyt? it's in the roadmap here -> huggingface/OpenEnv#940

sergiopaniego · 2026-07-22T09:27:02Z

Ran this branch loop-owning with OpenCode on HF Jobs; two things surfaced (both line up with the Cursor comments above, now with a concrete repro):

1. _messages_from_trace crashes on empty choices. The last proxy-trace entry can have choices == [] (an aborted/blank final request from the agent); choices[0] then raises IndexError, which kills the rollout child and stops training:

File ".../trl/experimental/async_grpo/openenv_harness.py", in _messages_from_trace
    content = last["response"]["choices"][0]["message"].get("content", "")
IndexError: list index out of range

Minimal guard:

choices = last["response"].get("choices") or []
content = choices[0]["message"].get("content", "") if choices else ""

(plus rebuilding the final assistant tool_calls, per the line-182 comment). With this guard applied, the same run trains to completion.

2. Rollout child doesn't exit cleanly. Confirmed the [RANK 0] Child did not exit within 15s; terminating path (the session-pool shutdown() already flagged). Side effect with a remote sandbox backend: in-flight sandboxes then leak, because session.close() never runs when the child is force-terminated.

sergiopaniego · 2026-07-22T14:53:01Z

Heads-up on another rough edge hit while training a loop-owning harness (Pi coding agent) on this branch, separate from the empty-choices one:

print_prompt_completions_sample → format_entry (trl/trainer/utils.py) does t.append(msg["content"]), which assumes content is a str. Agents that emit structured content (a list of parts, e.g. [{"type": "text", "text": "..."}] — Pi does this) trip rich's Text.append with TypeError: Only str or Text can be appended to Text, crashing the rollout worker when log_completions=True.

Repro-side workaround: coerce structured content to text before append, e.g.

c = msg["content"]
t.append(c if isinstance(c, str) else "".join(
    p.get("text", "") if isinstance(p, dict) else str(p) for p in c))

Happy to open a small PR if useful.

qgallouedec · 2026-07-22T15:50:18Z

@sergiopaniego can you check #5309, would it solve the issue? I can revamp it you think it makes sense

AmineDiro · 2026-07-22T16:15:06Z

print_prompt_completions_sample → format_entry (trl/trainer/utils.py) does t.append(msg["content"]), which assumes content is a str. Agents that emit structured content (a list of parts, e.g. [{"type": "text", "text": "..."}] — Pi does this) trip rich's Text.append with TypeError: Only str or Text can be appended to Text, crashing the rollout worker when log_completions=True.

@sergiopaniego thanks for catching this, but print_prompt_completions_sample might not be useful in harness training. I've opened this PR for tracking traces in trackio directly 👍🏼

…ssion cleanup

cursor · 2026-07-22T16:32:29Z

+    last = entries[-1]
+    choices = (last.get("response") or {}).get("choices") or []
+    content = choices[0]["message"].get("content", "") if choices else ""
+    return list(last["request"].get("messages") or []) + [{"role": "assistant", "content": content}]


Structured content crashes logging

High Severity

Loop-owning completions keep agent message content as-is, including structured lists. With log_completions=True, print_prompt_completions_sample appends that value to rich Text and raises TypeError, which kills the score loop and stops the rollout worker.

^{Reviewed by Cursor Bugbot for commit 06b6083. Configure here.}

AmineDiro · 2026-07-22T16:38:09Z

In this commit : 06b6083

fixed @sergiopaniego you can rerun if possible :

Sandbox leaks fixed: Added live-session tracking and explicit pool teardown to ensure remote sandboxes are cleanly killed on shutdown.
Tool counts fixed: Removed double-counting bugs; there is now a single source of truth for tool metrics.

OpenEnv Workarounds @burtenshaw @adithya-s-k @sergiopaniego

I had to work around some OpenEnv fundamental type /API limitation. I noted them as TODO(@openenv) in this PR:

Missing Types & Methods: Declared local types (TraceEntry, loop-owning protocols, generic factories) to patch gaps in OpenEnv's API.
Heuristic Call Filtering: Manually filtering out background LLM calls (like summarization) because OpenEnv doesn't tag trace purposes.

So what I just pushed is

The harness is now completely framework-agnostic. Opencode-specific logic is moved to the app layer via three injectable hooks: rollout_reward_fn, train_turn_fn, and agent_turn_fn.
- rollout_reward_fn(HarnessRolloutOutcome) -> float | None — reward shaping (default: env reward, unshaped).
- train_turn_fn(HarnessTurn) -> bool — which turns to reinforce ie to actually train on.
- agent_turn_fn(list[TraceEntry]) -> list[TraceEntry] — which trace entries are real agent turns
Strict Typing: Replaced raw dicts with strict, documented data models (e.g., HarnessRolloutOutcome, HarnessTurn).

I also added this : cf #6349 @qgallouedec

Deterministic Rollouts: Tasks are now seeded by group_id (not the prompt), ensuring all generations within a group get the exact same task.

qgallouedec

lgtm, just a bug (not sure) on group id

qgallouedec · 2026-07-22T16:54:24Z

+def _tool_failure_count(entries: list[TraceEntry]) -> int:
+    """Best-effort tool-failure count across the real agent turns (`entries`). The agent (not TRL) executed the tools,
+    so we only see their free-text RESULTS, which come back as `role="tool"` messages in a later request; a failure can
+    only be inferred from error-looking result text. The last agent request holds the fullest message list; dedupe by
+    (name, content) so a result echoed across requests is not counted twice."""
+    seen, failures = set(), 0
+    for msg in (entries[-1]["request"] if entries else {}).get("messages") or []:
+        if msg.get("role") != "tool":
+            continue
+        key = (msg.get("name"), str(msg.get("content"))[:200])
+        if key in seen:
+            continue
+        seen.add(key)
+        text = str(msg.get("content") or "").lower()
+        if any(w in text for w in ("error", "failed", "traceback", "exception")):
+            failures += 1
+    return failures


quite fragile, consider dropping it.
if you want to keep it, let's keep in mind that we should re-visit later when openenv exposes real tool-result status

100% agree, it was needed for tracking and reward, but this needs to be exposed from the trace directly

cursor · 2026-07-22T17:05:58Z

+        return []
+    last = entries[-1]
+    choices = (last.get("response") or {}).get("choices") or []
+    content = choices[0]["message"].get("content", "") if choices else ""


Fragile final-choice message access

Medium Severity

Empty choices is handled, but a non-empty choice missing message still does choices[0]["message"] and raises. That happens after turns are already built, so the outer handler discards an otherwise trainable rollout as unscorable.

Additional Locations (1)

trl/experimental/async_grpo/openenv_harness.py#L216-L239

^{Reviewed by Cursor Bugbot for commit 9015404. Configure here.}

Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>

cursor · 2026-07-22T18:44:59Z

+                )
+                completion = result.messages
+                tool_call_count = int(result.metrics.get("tool_calls", len(result.tool_trace)))
+                tool_failure_count = sum(1 for entry in result.tool_trace if entry.result.error is not None)


White-box hooks incomplete

Medium Severity

On the white-box path, train_turn_fn is never applied and tool_calls_by_name stays an empty dict. rollout_reward_fn therefore sees zero per-tool counts, and every sampled turn is trained even when the caller passed a turn filter.

Additional Locations (1)

trl/experimental/async_grpo/openenv_harness.py#L223-L232

^{Reviewed by Cursor Bugbot for commit 789acec. Configure here.}

cursor · 2026-07-22T18:44:59Z

+    `_start_proxy` exactly except for the port."""
+
+    def _start_proxy(self, sandbox):
+        port = _free_port()


Proxy port race

Medium Severity

_free_port binds a port, closes the socket, then later starts the proxy on that port. Concurrent rollouts can each observe the same free port between release and bind, so proxy startup fails and those rollouts become unscorable.

^{Reviewed by Cursor Bugbot for commit 789acec. Configure here.}

cursor

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

There are 7 total unresolved issues (including 5 from previous reviews).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 444ca7d. Configure here.}

cursor · 2026-07-22T18:49:25Z

+            tokenize=True,
+            return_dict=False,
+        )
+        turns.append(TurnRecord(prompt_ids, _trace_output_ids(entry), entry.get("per_token_logps") or []))


Mismatched logprobs break sequences

Medium Severity

_turns_from_trace builds a TurnRecord from proxy completion_token_ids and per_token_logps without checking equal length. _SampleBuilder._append then extends tokens/mask by the id length and logprobs by the logprob list length, so a partial capture silently misaligns training tensors.

Additional Locations (1)

trl/experimental/async_grpo/async_rollout_worker.py#L145-L149

^{Reviewed by Cursor Bugbot for commit 444ca7d. Configure here.}

cursor · 2026-07-22T18:49:25Z

+            d = os.path.join(sdir, sub)
+            shutil.rmtree(d, ignore_errors=True)
+            os.makedirs(d, exist_ok=True)
+        return LocalSandboxHandle(sdir, home_alias=self._alias, base_env=envs, cleanup=True)


Hardlink sandboxes share install files

Medium Severity

create hardlink-clones the template with cp -al and only replaces a few subdirs. The .opencode install tree stays shared across sandboxes, so in-place writes (cache, autoupdate, or package mutation) can corrupt the template and concurrent rollouts despite the comment claiming isolation.

^{Reviewed by Cursor Bugbot for commit 444ca7d. Configure here.}

AmineDiro requested a review from qgallouedec July 16, 2026 13:39

AmineDiro added 2 commits July 16, 2026 14:04

impl openenv harness rollout worker

b458bcc

clean up openenv harness rollout worker comments

1903378

AmineDiro force-pushed the async-grpo-harness-rollout branch from dabfebb to 1903378 Compare July 16, 2026 14:07

simplify child loop-class swap via _loop_cls attribute

284da95

AmineDiro marked this pull request as ready for review July 16, 2026 14:38

AmineDiro requested review from adithya-s-k, burtenshaw and sergiopaniego July 16, 2026 14:40

cursor Bot reviewed Jul 16, 2026

View reviewed changes

Comment thread trl/experimental/async_grpo/openenv_harness.py Outdated

Comment thread trl/experimental/async_grpo/openenv_harness.py

Comment thread trl/experimental/async_grpo/openenv_harness.py

Merge branch 'main' into async-grpo-harness-rollout

4a4af23

cursor Bot reviewed Jul 16, 2026

View reviewed changes

Comment thread tests/experimental/test_async_grpo_trainer.py

AmineDiro changed the title ~~Async grpo harness rollout~~ Async grpo OpenEnv harness rollout Jul 16, 2026

AmineDiro and others added 2 commits July 20, 2026 08:40

Fix async_grpo unit tests for 6-tuple rollout API

f91730c

Merge branch 'main' into async-grpo-harness-rollout

e6d759d

sergiopaniego approved these changes Jul 20, 2026

View reviewed changes

This was referenced Jul 20, 2026

RFC 006: Agentic RL through Harness Interception — token-faithful traces for TRL huggingface/OpenEnv#941

Open

Add HFSandboxBackend for the OpenCode harness huggingface/OpenEnv#998

Open

sergiopaniego mentioned this pull request Jul 22, 2026

Add pi_env: Pi coding-agent environment (HF sandbox + transparent-proxy logprobs) huggingface/OpenEnv#999

Open

Add harness rollout worker policy hooks, typed trace, and graceful se…

06b6083

…ssion cleanup

cursor Bot reviewed Jul 22, 2026

View reviewed changes

qgallouedec approved these changes Jul 22, 2026

View reviewed changes

Add opencode loop-owning AsyncGRPO example (local sandbox)

9015404

cursor Bot reviewed Jul 22, 2026

View reviewed changes

Update trl/experimental/async_grpo/async_rollout_worker.py

789acec

Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>

cursor Bot reviewed Jul 22, 2026

View reviewed changes

Thread group_id into _generate_one so harness per-group seed varies

444ca7d

cursor Bot reviewed Jul 22, 2026

View reviewed changes

Uh oh!

Conversation

AmineDiro commented Jul 16, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bot-ci-comment Bot commented Jul 16, 2026

Uh oh!

Uh oh!

qgallouedec commented Jul 20, 2026

Uh oh!

sergiopaniego left a comment

Choose a reason for hiding this comment

Uh oh!

sergiopaniego commented Jul 22, 2026

Uh oh!

sergiopaniego commented Jul 22, 2026

Uh oh!

qgallouedec commented Jul 22, 2026

Uh oh!

AmineDiro commented Jul 22, 2026

Uh oh!

Uh oh!

cursor Bot Jul 22, 2026

Choose a reason for hiding this comment

Structured content crashes logging

Uh oh!

AmineDiro commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

OpenEnv Workarounds @burtenshaw @adithya-s-k @sergiopaniego

Uh oh!

qgallouedec left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

qgallouedec Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

AmineDiro Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cursor Bot Jul 22, 2026

Choose a reason for hiding this comment

Fragile final-choice message access

Uh oh!

Uh oh!

cursor Bot Jul 22, 2026

Choose a reason for hiding this comment

White-box hooks incomplete

Uh oh!

cursor Bot Jul 22, 2026

Choose a reason for hiding this comment

Proxy port race

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

cursor Bot Jul 22, 2026

Choose a reason for hiding this comment

Mismatched logprobs break sequences

Uh oh!

cursor Bot Jul 22, 2026

Choose a reason for hiding this comment

Hardlink sandboxes share install files

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

AmineDiro commented Jul 16, 2026 •

edited by cursor Bot

Loading

AmineDiro commented Jul 22, 2026 •

edited

Loading