Async grpo OpenEnv harness rollout

[AmineDiro]

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 the rollout worker spawn path and scoring contract (new reward column and _generate_one arity) while the default AsyncRolloutWorker behavior stays the same; harness mode runs synchronous OpenEnv sessions on a thread pool and depends on correct proxy traces and verify rewards for GRPO signal.

Overview
Adds HarnessRolloutWorker / openenv_harness.py, an Async GRPO rollout backend that runs OpenEnv ResourceSessionFactory sessions instead of TRL’s built-in multi-turn tool loop. Rollouts can use white-box mode (adapter drives turns via vLLM through _sample_turn) or loop-owning mode (agent uses a proxy; TRL rebuilds TurnRecords from the trace and still runs _chain_to_sequences for training rows). session.verify supplies harness_reward, so training can run with empty reward_funcs.

The async rollout stack is extended to support this: spawned children take a configurable _loop_cls, RolloutGroup carries rollout_rewards, _generate_one returns an optional sixth reward value, and _provides_rollout_reward lets subclasses register an extra reward column in _score_group. RolloutWorkerProtocol now allows multiprocessing.Queue and renames update_model_version’s argument; tests follow the new return tuple and stub rollout_rewards.

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

[cursor]

Final tool calls dropped

Medium Severity

_messages_from_trace rebuilds only the final assistant content and omits tool_calls (and treats JSON null content as None). Episodes that end on a tool-call turn therefore hand session.verify an incomplete transcript, so harness rewards can be wrong for truncated or tool-heavy OpenCode runs.

 

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

[qgallouedec]

I guess this is valid

[cursor]

Session pool never stops

Medium Severity

_HarnessRolloutLoop creates a long-lived ThreadPoolExecutor for harness sessions but never shuts it down when the asyncio loop stops. Those workers are non-daemon and can keep running wait_for_completion, so the child process often fails to exit cleanly and is force-terminated after the parent's join timeout.

Additional Locations (1)

• trl/experimental/async_grpo/openenv_harness.py#L88-L92

 

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

[qgallouedec]

yes, this needs a shutdown() on teardown.

[cursor]

Empty traces skew advantages

High Severity

When the proxy trace is empty or _trace_output_ids recovers no token ids, _chain_to_sequences yields no training rows, but session.verify still supplies a rollout_reward that enters the group's advantage normalization. Sibling generations are then trained against rewards from conversations that contribute no samples.

Additional Locations (1)

• trl/experimental/async_grpo/openenv_harness.py#L144-L156

 

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

[qgallouedec]

When the proxy trace is empty or _trace_output_ids recovers no token ids

no idea in what case this is possible

[AmineDiro]

I dont think thats real case

[bot-ci-comment]

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.

[cursor]

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.}

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

[qgallouedec]

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