diff --git a/docs/source/example_overview.md b/docs/source/example_overview.md index bf82d5ebd8b..fdd743d8bdb 100644 --- a/docs/source/example_overview.md +++ b/docs/source/example_overview.md @@ -24,9 +24,8 @@ Check for additional optional dependencies [here](https://github.com/huggingface | Example | Description | Open in Colab | | --- | --- | --- | | [`async_distillation_math`](https://github.com/huggingface/trl/tree/main/examples/async_distillation_math) | Async on-policy distillation on GSM8K with [`experimental.async_distillation.AsyncDistillationTrainer`]: the teacher is served over HTTP with vLLM, including a multi-teacher (MOPD) math + code variant. | | -| [`async_grpo_harbor`](https://github.com/huggingface/trl/tree/main/examples/async_grpo_harbor) | AsyncGRPO training on a [Harbor](harbor) task suite through any installed coding agent (`mini-swe-agent`, `opencode`, `claude-code`, ...) served over [OpenEnv](openenv): the agent owns its own loop in a sandbox and TRL trains on the token ids and logprobs recorded by the capture proxy. | | +| [`async_grpo_harbor`](https://github.com/huggingface/trl/tree/main/examples/async_grpo_harbor) | AsyncGRPO training on a [Harbor](https://huggingface.co/docs/openenv/environments/harbor) task suite through any installed coding agent (`mini-swe-agent`, `opencode`, `claude-code`, ...) served over [OpenEnv](openenv): the agent owns its own loop in a sandbox and TRL trains on the token ids and logprobs recorded by the capture proxy. | | | [`async_grpo_math`](https://github.com/huggingface/trl/tree/main/examples/async_grpo_math) | Asynchronous GRPO on GSM8K with [`experimental.async_grpo.AsyncGRPOTrainer`], decoupling generation (vLLM server) from training, including a LoRA variant that syncs only the adapter to the vLLM server. | | -| [`async_grpo_opencode`](https://github.com/huggingface/trl/tree/main/examples/async_grpo_opencode) | AsyncGRPO training of the real `opencode` coding agent on an [OpenEnv](openenv) environment (loop-owning: the external agent runs its own tool loop and TRL trains on its captured proxy trace), with a local subprocess sandbox or remote Hugging Face sandboxes. | | | [`dpo_reduce_hallucinations`](https://github.com/huggingface/trl/tree/main/examples/dpo_reduce_hallucinations) | DPO fine-tuning of a Vision Language Model to reduce hallucinations using the [openbmb/RLAIF-V-Dataset](https://huggingface.co/datasets/openbmb/RLAIF-V-Dataset) dataset. | | | [`gold_chatbot_arena`](https://github.com/huggingface/trl/tree/main/examples/gold_chatbot_arena) | General Online Logit Distillation (GOLD) of a Qwen2 teacher into a Llama 3.2 student (cross-tokenizer) on chatbot_arena_completions with [`experimental.gold.GOLDTrainer`], with full-training and LoRA variants. | | | [`gold_qwen3_vl`](https://github.com/huggingface/trl/tree/main/examples/gold_qwen3_vl) | General Online Logit Distillation (GOLD) of Qwen3-VL-8B into smaller VLM students with [`experimental.gold.GOLDTrainer`], covering same-family (JSD loss) and cross-family (ULD loss) distillation. | | diff --git a/docs/source/openenv.md b/docs/source/openenv.md index 3df7dcccb43..9d529d54997 100644 --- a/docs/source/openenv.md +++ b/docs/source/openenv.md @@ -5,7 +5,7 @@ This guide covers **how to integrate OpenEnv with TRL**. For more on OpenEnv itself, see the [OpenEnv docs](https://huggingface.co/docs/openenv). > [!NOTE] -> Ready-to-use OpenEnv examples: [`grpo_echo`](https://github.com/huggingface/trl/tree/main/examples/grpo_echo) (minimal), [`grpo_catch`](https://github.com/huggingface/trl/tree/main/examples/grpo_catch), [`grpo_wordle`](https://github.com/huggingface/trl/tree/main/examples/grpo_wordle), [`grpo_sudoku`](https://github.com/huggingface/trl/tree/main/examples/grpo_sudoku), [`grpo_multi_env`](https://github.com/huggingface/trl/tree/main/examples/grpo_multi_env), [`grpo_browsergym`](https://github.com/huggingface/trl/tree/main/examples/grpo_browsergym), [`grpo_carla`](https://github.com/huggingface/trl/tree/main/examples/grpo_carla), [`async_grpo_opencode`](https://github.com/huggingface/trl/tree/main/examples/async_grpo_opencode), and [`async_grpo_harbor`](https://github.com/huggingface/trl/tree/main/examples/async_grpo_harbor) (local and HF Jobs). +> Ready-to-use OpenEnv examples: [`grpo_echo`](https://github.com/huggingface/trl/tree/main/examples/grpo_echo) (minimal), [`grpo_catch`](https://github.com/huggingface/trl/tree/main/examples/grpo_catch), [`grpo_wordle`](https://github.com/huggingface/trl/tree/main/examples/grpo_wordle), [`grpo_sudoku`](https://github.com/huggingface/trl/tree/main/examples/grpo_sudoku), [`grpo_multi_env`](https://github.com/huggingface/trl/tree/main/examples/grpo_multi_env), [`grpo_browsergym`](https://github.com/huggingface/trl/tree/main/examples/grpo_browsergym), [`grpo_carla`](https://github.com/huggingface/trl/tree/main/examples/grpo_carla), and [`async_grpo_harbor`](https://github.com/huggingface/trl/tree/main/examples/async_grpo_harbor) (local and HF Jobs). ## When to use environments @@ -604,40 +604,49 @@ app = create_app( Use `rollout_func` when `environment_factory` doesn't fit your use case. For example, **external agent servers** where an external server owns the generation loop and manages its own agent-environment interaction protocol. -## Training on harnesses: training a real coding agent (opencode) +## Training on harnesses: training real coding agents (Harbor) The integrations above are **white-box**: TRL drives the multi-turn loop itself. It samples each turn, parses the tool calls, runs them, and feeds the results back. -Some agents cannot be driven this way because they own their own loop. A production coding agent harness like [`opencode`](https://opencode.ai) has its own planner, tool set, context management, and stop condition. You want to train that exact agent, not a reimplementation of it. +Some agents cannot be driven this way because they own their own loop. A production coding agent harness like [`opencode`](https://opencode.ai) or [`claude-code`](https://github.com/anthropics/claude-code) has its own planner, tool set, context management, and stop condition. You want to train that exact agent, not a reimplementation of it. -For this, TRL provides an experimental **black box (loop-owning)** path built on [`experimental.async_grpo.AsyncGRPOTrainer`] and a `HarnessRolloutWorker` specific for OpenEnv that drives an [OpenEnv `ResourceSessionFactory`](https://huggingface.co/docs/openenv). See [`examples/async_grpo_opencode/async_grpo_opencode.py`](https://github.com/huggingface/trl/blob/main/examples/async_grpo_opencode/async_grpo_opencode.py) for a complete, self-contained example. To scale rollouts beyond a single node, [`examples/async_grpo_opencode/opencode_hf_sandbox.py`](https://github.com/huggingface/trl/blob/main/examples/async_grpo_opencode/opencode_hf_sandbox.py) runs each rollout in its own remote Hugging Face sandbox instead of a local subprocess. - -The [Harbor example](https://github.com/huggingface/trl/tree/main/examples/async_grpo_harbor) uses the same worker with OpenEnv's Harbor integration. Harbor runs tasks through coding-agent harnesses such as OpenCode, Codex and Claude Code. It includes local setup and a Hugging Face Jobs launcher. +For this, TRL provides an experimental **black-box (loop-owning)** path built on [`experimental.async_grpo.AsyncGRPOTrainer`] and a `HarnessRolloutWorker` that drives an OpenEnv session factory. The agents and their sandboxes come from [Harbor](https://huggingface.co/docs/openenv/environments/harbor), served through OpenEnv with `openenv harbor serve`, so every agent Harbor installs is trainable the same way: you pick one with `--harness`. See [`examples/async_grpo_harbor`](https://github.com/huggingface/trl/tree/main/examples/async_grpo_harbor) for a complete example, including a Hugging Face Jobs launcher. ### How it works TRL does not sample each turn here. The agent runs to completion on its own, and TRL reads back what it did: -1. The agent runs inside an OpenEnv session (a local subprocess sandbox in the example, so no container is needed) in `transparent_proxy` mode. A small proxy inside the sandbox forwards the agent's `/v1/chat/completions` calls to your vLLM server and records each turn's token ids and logprobs. -2. When the agent stops, TRL receives the validated `TrainingTrace` from OpenEnv, builds rows from its tokens and masks, and scores the final workspace with the session's `verify()` method (a held-out verifier). -3. GRPO trains on those rows. The reward is propagated to every trained token through the group-relative advantage. +1. The OpenEnv Harbor server starts the agent on a Harbor task in a sandbox (for example E2B). Every model call the agent makes goes through the server's capture proxy to your vLLM server, which records the engine's own prompt and completion token ids, the sampled logprobs, and a per-token loss mask. +2. When the agent stops, the task's own verifier scores the result, and TRL receives the validated `TrainingTrace` from OpenEnv and builds the training rows from its token ids and masks. TRL does not re-tokenize the captured prompts. +3. GRPO trains on those rows. The reward is propagated to every trained token through the group-relative advantage, and the new weights are synced back into the same vLLM server. -Each rollout runs in its own isolated session. In the example that means one sandbox directory, one proxy on its own port, and one agent process per rollout. The isolation matters for two reasons: the proxy has to capture exactly that rollout's tokens, and one rollout must not interfere with another. The `max_inflight_tasks` setting controls how many rollouts run at the same time. +Each rollout runs in its own isolated session, with its own sandbox and its own capture session. The isolation matters for two reasons: the proxy has to capture exactly that rollout's tokens, and one rollout must not interfere with another. The `max_inflight_tasks` setting controls how many rollouts run at the same time, and the server's `MAX_CONCURRENT_ENVS` must be at least `max_inflight_tasks + 1` (the factory keeps one extra connection for task metadata). ### Wiring -Pass a `HarnessRolloutWorker` to [`experimental.async_grpo.AsyncGRPOTrainer`] with `harness_adapter=None` for a loop-owning agent. The factory constructor receives the trainer's sampling policy. Each session returns an OpenEnv `TrainingTrace` containing exact engine tokens, behavior logprobs and prefilled masks. `rollout_reward_fn` is optional; without it, TRL uses the session's verifier reward. +You pass a `HarnessRolloutWorker` to [`experimental.async_grpo.AsyncGRPOTrainer`] with `harness_adapter=None` to select loop-owning mode. The worker builds the session factory with its own sampling policy, so you pass it a callable such as `partial(HarborSessionFactory, ...)` rather than a built factory. Each session returns an OpenEnv `TrainingTrace` containing exact engine tokens, behavior logprobs and prefilled masks. `rollout_reward_fn` is optional; without it, TRL uses the session's verifier reward. ```python from functools import partial +from harbor_env.harness import HarborSessionFactory + from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer from trl.experimental.async_grpo.openenv_harness import HarnessRolloutWorker +factory = partial( + HarborSessionFactory, + "http://localhost:8200", # the OpenEnv Harbor server + split="", # a Harbor task dataset the server serves + harness="mini-swe-agent", # or "opencode", "claude-code", ... + sandbox="e2b", + llm_url=vllm_url, + model=model, +) worker = HarnessRolloutWorker( - harness_session_factory=partial(build_factory, ...), # accepts sampling= and returns a ResourceSessionFactory - harness_adapter=None, # loop-owning: the agent runs its own loop - rollout_reward_fn=my_reward, # outcome -> float | None + harness_session_factory=factory, # called by the worker with its sampling policy + harness_adapter=None, # loop-owning: the agent runs its own loop + rollout_reward_fn=my_reward, # outcome -> float | None model_name=model, dataset=dataset, reward_funcs=[], @@ -660,17 +669,17 @@ OpenEnv owns capture validation and token selection. TRL converts the producer's TRL does not know what success means for your task, so you turn each finished rollout into a scalar reward. The function receives a `HarnessRolloutOutcome` that describes what the agent did: -- `env_reward` (`float | None`): the reward from the session's `verify()` (for opencode, the fraction of held-out tests that passed), or `None` when the rollout could not be scored. +- `env_reward` (`float | None`): the reward from the task's verifier, or `None` when the rollout could not be scored. - `completion` (`list[dict]`): the final message transcript. - `trace` (`list[TraceEntry]`): the selected OpenEnv capture records, including zero-masked turns. - `tool_call_count` and `tool_failure_count` (`int`): how many tool calls the agent made, and how many looked like failures. - `tool_calls_by_name` (`dict[str, int]`): calls per tool, for example `{"bash": 3, "edit": 2}`. - `timed_out` (`bool`): whether the agent ran out of its time budget. -Return a `float`, or return `None` to mark the rollout unscorable so it is dropped from the group baseline instead of being counted as a zero. If you do not pass this function, the raw `env_reward` is used as is. The opencode example turns the dense pass fraction into a binary pass or fail and subtracts small penalties for degenerate behavior, such as never running the code or looping for far too many steps. +Return a `float`, or return `None` to mark the rollout unscorable so it is dropped from the group baseline instead of being counted as a zero. If you do not pass this function, the raw `env_reward` is used as is. The Harbor example keeps the verifier's correctness as the main signal and adds a small tool-efficiency bonus only when the task is solved. > [!NOTE] -> **Why not just use `verify()`?** `verify()` is the environment's job and answers one question, "how correct was the outcome," which keeps it clean and reusable for evaluation. The reward you train on is a separate, training-time decision (binarize the score, penalize degenerate behavior, drop unscorable rollouts). It also needs signals `verify()` never sees, since `verify()` only inspects the final workspace, while `rollout_reward_fn` also gets the trajectory (tool counts, `timed_out`, the trace). For example, a rollout can pass some tests yet never run `bash`; only `rollout_reward_fn` can see that and penalize it. +> **Why not just use the verifier?** The verifier is the environment's job and answers one question, "how correct was the outcome," which keeps it clean and reusable for evaluation. The reward you train on is a separate, training-time decision (binarize the score, penalize degenerate behavior, drop unscorable rollouts). It also needs signals the verifier never sees, since the verifier only inspects the final result, while `rollout_reward_fn` also gets the trajectory (tool counts, `timed_out`, the trace). For example, a rollout can pass some tests yet never run `bash`; only `rollout_reward_fn` can see that and penalize it. #### Token eligibility comes from OpenEnv @@ -678,21 +687,27 @@ Return a `float`, or return `None` to mark the rollout unscorable so it is dropp For example, prompt `[10, 11]`, completion `[20, 21, 22]` and mask `[0, 0, 1, 0, 1]` train on tokens 20 and 22. Token 21 remains context. A wholly masked completion is retained for context and usage accounting. -The producer excludes auxiliary calls and discarded retries. There are no `train_turn_fn` or `agent_turn_fn` parameters in this API. Existing integrations must move selection into their OpenEnv masks. The native OpenCode examples now use OpenEnv's agent selection and train final answers as well as tool-call turns. `verify()` remains independent of token eligibility. +The producer excludes auxiliary calls and discarded retries. There are no `train_turn_fn` or `agent_turn_fn` parameters in this API. Existing integrations must move selection into their OpenEnv masks. The task's verifier remains independent of token eligibility. Malformed captures fail at the OpenEnv boundary and stop the worker. Transport failures remain unscorable. Rewritten histories start separate rows by default, preserving the sampled tokens and logprobs. Multiple rows still share the rollout reward; this API does not change GRPO's weighting. ### Requirements -The vLLM server must expose tool-calling and real token ids, and enable NCCL weight sync so the agent always hits the current policy: +The OpenEnv Harbor server owns the dataset and the sandboxes. It needs the credential of the sandbox backend (for example `E2B_API_KEY`): + +```bash +MAX_CONCURRENT_ENVS=9 openenv harbor serve --dataset --port 8200 --capture-port 8300 --expose gradio +``` + +The vLLM server must expose tool calling and real token ids, and enable NCCL weight sync so the agent always hits the current policy. Rollouts and weight updates must use the same vLLM server: ```bash vllm serve \ - --enable-auto-tool-choice --tool-call-parser hermes \ + --enable-auto-tool-choice --tool-call-parser \ --logprobs-mode processed_logprobs \ --return-tokens-as-token-ids \ --weight-transfer-config '{"backend":"nccl"}' ``` > [!NOTE] -> Loop-owning training lives under `trl.experimental` and its API may change. The example installs the `opencode` CLI into a sandbox template on first run (needs internet once) and uses [`agentica-org/DeepCoder-Preview-Dataset`](https://huggingface.co/datasets/agentica-org/DeepCoder-Preview-Dataset) with a held-out stdin/stdout verifier. +> Loop-owning training lives under `trl.experimental` and its API may change. The example pins the OpenEnv revision it was tested with; see its [README](https://github.com/huggingface/trl/tree/main/examples/async_grpo_harbor) for the exact setup. diff --git a/examples/async_grpo_opencode/async_grpo_opencode.py b/examples/async_grpo_opencode/async_grpo_opencode.py deleted file mode 100644 index dc70118cace..00000000000 --- a/examples/async_grpo_opencode/async_grpo_opencode.py +++ /dev/null @@ -1,601 +0,0 @@ -# Copyright 2020-2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# /// script -# dependencies = [ -# "trl", -# "trackio", -# "datasets", -# "openenv @ git+https://github.com/huggingface/OpenEnv.git@86a180ede21e044f7929b9a7783ad83aa67d83a3", -# "openenv-opencode-env @ git+https://github.com/huggingface/OpenEnv.git@86a180ede21e044f7929b9a7783ad83aa67d83a3#subdirectory=envs/opencode_env", -# ] -# /// - -"""AsyncGRPO training of the real `opencode` coding agent (loop-owning) with a local subprocess sandbox. - -`opencode` is a genuine external coding agent that owns its own tool loop. Each rollout runs it as a local process -(no container) via the `LocalSubprocessSandboxBackend` defined below, in `transparent_proxy` mode: an in-sandbox -proxy forwards the agent's calls to your vLLM server and captures per-turn `(token_ids, logprobs)`. TRL reads that -proxy trace, rebuilds training rows, scores the workspace with a held-out verifier, and trains with GRPO. - -Task: competitive-coding problems from `agentica-org/DeepCoder-Preview-Dataset`. The agent writes `solution.py` -(reads stdin, prints stdout); the verifier runs it against the problem's HELD-OUT tests (never shown to the agent) -and returns a DENSE reward = fraction passed. `opencode_reward` then binarizes it and adds small degeneracy -penalties. This whole file is self-contained and every training-facing object is module-level (picklable), so the -rollout worker can pickle the factory + verifier into its spawned child process. - -Requirements: - - An OpenAI-compatible vLLM server (see below) reachable at `--vllm-url`. - - Internet on this node the first time: `warmup()` installs the `opencode` CLI into a template dir once. - - Install this script's pinned `openenv` and `openenv-opencode-env` dependencies, or use `uv run`. - -Run (2 GPUs: vLLM on one, trainer on the other): - -```sh -# Terminal 1 - serve the policy. Tool-calling + token-ids + NCCL weight-sync are all required. -CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-4B-Instruct-2507 \ - --host 0.0.0.0 --port 8000 \ - --enable-auto-tool-choice --tool-call-parser hermes \ - --logprobs-mode processed_logprobs \ - --return-tokens-as-token-ids \ - --weight-transfer-config '{"backend":"nccl"}' - -# Terminal 2 - train. -CUDA_VISIBLE_DEVICES=1 python examples/async_grpo_opencode/async_grpo_opencode.py \ - --model Qwen/Qwen3-4B-Instruct-2507 --vllm-url http://localhost:8000 -``` -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import os -import random -import shlex -import shutil -import signal -import socket -import subprocess -import tempfile -import time -import uuid -from functools import partial -from pathlib import Path -from typing import Any - -from datasets import Dataset, load_dataset -from opencode_env import harness as oc_harness -from opencode_env.config import OpenCodeConfig -from opencode_env.harness import OpenCodeSessionFactory -from opencode_env.sandbox.base import ExecResult, SandboxHandle -from opencode_env.task import OpenCodeTask -from openenv.core.harness import ResourceSession, ResourceSessionFactory, VerifyResult -from transformers import AutoTokenizer - -from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer -from trl.experimental.async_grpo.openenv_harness import ( - HarnessRolloutOutcome, - HarnessRolloutWorker, -) - - -# ============================================================================================================ -# Local subprocess sandbox backend -# ------------------------------------------------------------------------------------------------------------ -# OpenEnv's opencode harness only ships an E2B (cloud) backend, and a cloud sandbox can't reach a local vLLM. -# The `SandboxBackend` protocol is small, so we run opencode + its proxy as local processes on this node. The -# harness bakes the prefix `/home/user` into several paths, so each sandbox REMAPS that prefix to its own dir and -# callers pass `OpenCodeConfig(sandbox_home="/home/user")` so config-driven paths funnel through the same remap. -# ============================================================================================================ - -_OPENCODE_INSTALL = "curl -fsSL https://opencode.ai/install | bash -s -- --no-modify-path" - - -class LocalBgJob: - """A background process (the opencode agent or its proxy) running directly on the node.""" - - def __init__(self, popen: subprocess.Popen): - self._p = popen - - @property - def pid(self) -> int: - return self._p.pid - - def wait(self, timeout: float | None = None) -> int: - try: - return self._p.wait(timeout=timeout) - except subprocess.TimeoutExpired as e: - raise TimeoutError(str(e)) from e - - def kill(self) -> None: - # Kill the whole process GROUP: opencode spawns a tree (node -> bash -> python); SIGTERM to only the parent - # orphans the children, which pile up across rollouts. `start_bg` launches each job in its own session. - if self._p.poll() is not None: - return - try: - pgid = os.getpgid(self._p.pid) - except ProcessLookupError: - return - try: - os.killpg(pgid, signal.SIGTERM) - self._p.wait(timeout=5) - except (subprocess.TimeoutExpired, Exception): - try: - os.killpg(pgid, signal.SIGKILL) - except ProcessLookupError: - pass - - -class LocalSandboxHandle: - """One local 'sandbox' = a real directory on the node. The harness's hardcoded `/home/user` prefix is remapped - to this directory in every command and path, and `$HOME` points at it. `kill()` removes the directory.""" - - def __init__( - self, - root: str, - *, - home_alias: str = "/home/user", - base_env: dict[str, str] | None = None, - cleanup: bool = False, - ): - self._root = root - self._alias = home_alias - self._cleanup = cleanup - self._env = {**os.environ, "HOME": root, **(base_env or {})} - self._bg: list[LocalBgJob] = [] - - @property - def sandbox_id(self) -> str: - return self._root - - def _remap(self, s: str | None) -> str | None: - return s if s is None else s.replace(self._alias, self._root) - - def _run_env(self, envs: dict[str, str] | None) -> dict[str, str]: - return {**self._env, **(envs or {})} - - def exec(self, cmd: str, *, envs=None, cwd=None, timeout: float | None = 60) -> ExecResult: - try: - p = subprocess.run( - ["bash", "-lc", self._remap(cmd)], - cwd=self._remap(cwd) or self._root, - env=self._run_env(envs), - capture_output=True, - text=True, - timeout=timeout, - stdin=subprocess.DEVNULL, - ) - return ExecResult(exit_code=p.returncode, stdout=p.stdout, stderr=p.stderr) - except subprocess.TimeoutExpired as e: - return ExecResult(exit_code=124, stdout=e.stdout or "", stderr=f"timeout after {timeout}s") - - def start_bg(self, cmd: str, *, envs=None, cwd=None) -> LocalBgJob: - # stdin=/dev/null so the agent (and any `python solution.py` it runs) reads EOF instead of blocking forever. - p = subprocess.Popen( - ["bash", "-lc", self._remap(cmd)], - cwd=self._remap(cwd) or self._root, - env=self._run_env(envs), - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - stdin=subprocess.DEVNULL, - start_new_session=True, # own process group so kill() reaps the whole opencode tree - ) - job = LocalBgJob(p) - self._bg.append(job) - return job - - def write_text(self, path: str, content: str) -> None: - path = self._remap(path) - Path(path).parent.mkdir(parents=True, exist_ok=True) - Path(path).write_text(content) - - def read_text(self, path: str) -> str: - return Path(self._remap(path)).read_text() - - def exists(self, path: str) -> bool: - return Path(self._remap(path)).exists() - - def kill(self) -> None: - for job in self._bg: - try: - job.kill() - except Exception: - pass - self._bg.clear() - if self._cleanup: - shutil.rmtree(self._root, ignore_errors=True) - - -class LocalSubprocessSandboxBackend: - """Produces per-rollout `LocalSandboxHandle`s, each in its own `uuid` dir hardlink-cloned from a template that - has opencode pre-installed (`warmup()`), so concurrent sandboxes never share state and never re-install.""" - - def __init__(self, root: str, *, home_alias: str = "/home/user"): - self._root = root - self._alias = home_alias - self._template = os.path.join(root, "_template") - - def warmup(self) -> None: - """Install opencode ONCE into the template dir (run in the parent, before rollouts spawn).""" - marker = os.path.join(self._template, ".opencode", "bin", "opencode") - if os.path.exists(marker): - return - os.makedirs(self._template, exist_ok=True) - subprocess.run( - ["bash", "-lc", _OPENCODE_INSTALL], - env={**os.environ, "HOME": self._template}, - check=True, - timeout=400, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - - def create(self, *, timeout_s: int = 900, envs=None, metadata=None) -> LocalSandboxHandle: - name = (metadata or {}).get("episode_id") or uuid.uuid4().hex - sdir = os.path.join(self._root, name) - shutil.rmtree(sdir, ignore_errors=True) - os.makedirs(sdir, exist_ok=True) - if os.path.isdir(self._template): - subprocess.run(["cp", "-al", f"{self._template}/.", f"{sdir}/"], check=True) # hardlink-clone - for sub in ("workdir", "task", "logs/agent", "logs/verifier", ".config/opencode"): - 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) - - -# ============================================================================================================ -# Dataset + held-out verifier -# ============================================================================================================ - -DATASET = "agentica-org/DeepCoder-Preview-Dataset" -DATASET_CONFIG = "primeintellect" # carries 11-93 stdin/stdout tests per problem -N_TESTS_EVAL = 12 # cap held-out tests scored per rollout (bounds verify latency); dense reward stays fine-grained -PER_TEST_TIMEOUT = 6 - - -def _instruction_id(instruction: str) -> str: - return hashlib.sha1(instruction.encode()).hexdigest() - - -def _clean_problem(problem: str) -> str: - # DeepCoder wraps problems in boilerplate that fights our tool-writing instruction ("...return the code."), which - # nudges a weak model to dump code in its reply instead of writing solution.py. Strip it so our instruction leads. - text = problem.strip() - prefix = "Solve the following coding problem using the programming language python:" - if text.startswith(prefix): - text = text[len(prefix) :].strip() - cut = len(text) - for marker in ("The input will be", "Now solve the problem", "Now solve this problem"): - idx = text.rfind(marker) - if idx != -1: - cut = min(cut, idx) - return text[:cut].strip() - - -def _coding_instruction(problem: str) -> str: - # Multi-turn loop: write -> run on the example cases -> read feedback -> fix -> repeat. The examples in the problem - # are the feedback signal; the held-out tests (the reward) stay hidden. - problem = _clean_problem(problem) - return ( - "Solve the problem below by writing a Python program to `solution.py` in the current working directory " - "(it does not exist yet, so create it with the `write` tool). The program must read its input from standard " - "input (stdin) and write ONLY the required answer to standard output (stdout).\n" - "Then TEST it against the example cases shown in the problem: run your program on each example input with " - "bash, e.g. `printf '' | python3 solution.py`, and compare its output to the expected output. " - "If an example fails or the program errors, use the `edit` tool to fix `solution.py`, then run it again. " - "Keep going until your program passes all the provided examples, then stop.\n\n" - f"PROBLEM:\n{problem}" - ) - - -# Runner executed INSIDE the sandbox: run solution.py against each held-out test with a hard time/mem/output cap, -# compare normalized stdout, print `SCORE: `. Whitespace-insensitive match (competitive judges usually are). -_RUNNER_SRC = r""" -import json, subprocess, sys - -PER_TEST_TIMEOUT = {per_test_timeout} - -def norm(s): - lines = [ln.rstrip() for ln in (s or "").replace("\r\n", "\n").split("\n")] - while lines and lines[-1] == "": - lines.pop() - return "\n".join(lines) - -def ok(got, exp): - if norm(got) == norm(exp): - return True - return (got or "").split() == (exp or "").split() - -tests = json.load(open("_tests.json")) -passed = 0 -for t in tests: - cmd = "ulimit -v 2000000 2>/dev/null; timeout %d python3 solution.py 2>/dev/null | head -c 2000000" % PER_TEST_TIMEOUT - try: - p = subprocess.run(["bash", "-c", cmd], input=t.get("input", ""), capture_output=True, text=True, - timeout=PER_TEST_TIMEOUT + 5) - if ok(p.stdout, t.get("output", "")): - passed += 1 - except Exception: - pass -print("SCORE: %.6f" % (passed / len(tests) if tests else 0.0)) -""" - - -def _run_dense_tests(sandbox: SandboxHandle, tests: list[dict[str, str]]) -> float: - """Run the sandbox's `solution.py` against `tests`; return the fraction passed.""" - if not tests or not sandbox.exists("/home/user/workdir/solution.py"): - return 0.0 - tests = tests[:N_TESTS_EVAL] - sandbox.write_text("/home/user/workdir/_tests.json", json.dumps(tests)) - sandbox.write_text("/home/user/workdir/_run_tests.py", _RUNNER_SRC.format(per_test_timeout=PER_TEST_TIMEOUT)) - r = sandbox.exec("cd /home/user/workdir && python3 _run_tests.py", timeout=PER_TEST_TIMEOUT * len(tests) + 30) - for line in (r.stdout or "").splitlines(): - if line.startswith("SCORE:"): - return float(line.split(":", 1)[1].strip()) - return 0.0 - - -class DeepCoderStdinVerifier: - """Dense stdin/stdout verifier. Holds the held-out test map (keyed by `sha1(instruction)`) so it survives the - pickle into the rollout child. `session.verify(...)` calls it as `verifier(sandbox, task)`.""" - - def __init__(self, tests_by_id: dict[str, list[dict[str, str]]]): - self._tests_by_id = tests_by_id - - def __call__(self, sandbox: SandboxHandle, task: OpenCodeTask) -> VerifyResult: - tests = self._tests_by_id.get(_instruction_id(task.instruction)) - return VerifyResult(env_reward=_run_dense_tests(sandbox, tests or []), done=True) - - -def build_dataset(n_prompts: int, seed: int) -> tuple[list[dict], dict[str, list[dict[str, str]]]]: - """Return `(rows, tests_by_id)`: the prompt rows (problem statement only) and the held-out test map for the - verifier. Tests are NOT put in the prompt - the agent only sees the statement (which includes sample cases).""" - rows = list(load_dataset(DATASET, DATASET_CONFIG, split="train")) - random.Random(seed).shuffle(rows) - out: list[dict] = [] - tests_by_id: dict[str, list[dict[str, str]]] = {} - for r in rows: - if len(out) >= n_prompts: - break - raw = r["tests"] - if isinstance(raw, str): - raw = json.loads(raw) - tests = [ - {"input": t.get("input", ""), "output": t.get("output", "")} - for t in raw - if t.get("type", "stdin_stdout") == "stdin_stdout" and t.get("output") is not None - ] - if len(tests) < 3: # need enough tests for a meaningful dense fraction - continue - instruction = _coding_instruction(r["problem"]) - tests_by_id[_instruction_id(instruction)] = tests[:N_TESTS_EVAL] - out.append({"prompt": [{"role": "user", "content": instruction}]}) - return out, tests_by_id - - -# ============================================================================================================ -# opencode session factory (local sandbox + per-session proxy port) -# ============================================================================================================ - - -class OpencodeTaskFactory(ResourceSessionFactory): - """Adapts the worker's `create(prompt=messages, seed, episode_id)` onto `OpenCodeSessionFactory`, which wants an - `OpenCodeTask`: pull the instruction out of the last user message.""" - - def __init__(self, inner: OpenCodeSessionFactory): - self._inner = inner - - def create(self, task: Any, seed: int | None = None, episode_id: str | None = None) -> ResourceSession: - instruction = task[-1]["content"] if isinstance(task, list) and task else str(task) - return self._inner.create(instruction, seed=seed, episode_id=episode_id) - - -def _free_port() -> int: - with socket.socket() as s: - s.bind(("127.0.0.1", 0)) - return s.getsockname()[1] - - -class FreePortOpenCodeSessionFactory(OpenCodeSessionFactory): - """Same as `OpenCodeSessionFactory` but binds the in-sandbox proxy to a free port per session instead of the - hardcoded `_PROXY_PORT = 7000`, so several opencode sandboxes can run at once on one node. Mirrors OpenEnv's - `_start_proxy` exactly except for the port.""" - - def _start_proxy(self, sandbox): - port = _free_port() - trace_path = oc_harness.proxy_trace_path(self._config) - log_path = oc_harness.proxy_log_path(self._config) - if not sandbox.exists("/home/user/proxy/interception.py"): - self._exec_with_retry( - sandbox, - "pip install --quiet 'fastapi>=0.104' 'uvicorn[standard]>=0.24' 'httpx>=0.27' 2>&1 | tail -20", - timeout=180, - attempts=3, - backoff_s=2.0, - label="proxy deps install", - ) - sandbox.write_text("/home/user/proxy/interception.py", oc_harness._PROXY_SOURCE_PATH.read_text()) - sandbox.write_text("/home/user/proxy/__init__.py", "") - - if self.sampling is not None: - sandbox.write_text("/home/user/proxy/interception.py", oc_harness._PROXY_SOURCE_PATH.read_text()) - - proxy_args = [ - "python", "interception.py", "--upstream-url", self._config.base_url, - "--trace", trace_path, "--port", str(port), - "--top-logprobs", str(self._config.proxy_top_logprobs), - ] # fmt: skip - if self.sampling is not None: - proxy_args += ["--sampling", json.dumps(self.sampling)] - if self._config.proxy_max_tokens_cap is not None: - proxy_args += ["--max-tokens-cap", str(self._config.proxy_max_tokens_cap)] - if self._config.proxy_disable_thinking: - proxy_args.append("--disable-thinking") - if self._config.model: - proxy_args += ["--model-override", self._config.model] - - quoted = " ".join(shlex.quote(a) for a in proxy_args) - proxy_cmd = f"cd /home/user/proxy && {quoted} > {shlex.quote(log_path)} 2>&1" - proxy_job = sandbox.start_bg(proxy_cmd, envs={"OPENCODE_UPSTREAM_API_KEY": self._config.api_key}) - - for _ in range(120): - if sandbox.exec(f"curl -sf http://127.0.0.1:{port}/healthz", timeout=5).exit_code == 0: - break - time.sleep(0.5) - else: - log = "" - try: - log = sandbox.read_text(log_path) - except Exception: - pass - proxy_job.kill() - raise RuntimeError(f"proxy did not start on :{port}\n{log[-2000:]}") - - return proxy_job, f"http://127.0.0.1:{port}/v1", trace_path - - -def build_factory( - sandbox_root: str, vllm_url: str, model: str, tests_by_id: dict, *, sampling: dict -) -> OpencodeTaskFactory: - config = OpenCodeConfig( - extra_opencode_json={"agent": {"build": {"temperature": sampling["temperature"], "top_p": sampling["top_p"]}}}, - provider="openai_compatible", - base_url=f"{vllm_url}/v1", - model=model, # proxy --model-override forces this exact id on upstream requests - sandbox_home="/home/user", # remapped to each sandbox's real dir by LocalSandboxHandle - agent_timeout_s=180.0, # bounds edit/bash-loop blowups; legit solves finish in <90s - disabled_tools=["webfetch", "question", "task"], # no web, no user, no sub-agents - run_format="json", - ) - backend = LocalSubprocessSandboxBackend(sandbox_root) - inner = FreePortOpenCodeSessionFactory( - config=config, - sandbox_backend=backend, - mode="transparent_proxy", - sampling=sampling, - verifier=DeepCoderStdinVerifier(tests_by_id), - ) - return OpencodeTaskFactory(inner) - - -# ============================================================================================================ -# Rollout reward -# ============================================================================================================ - - -def opencode_reward(outcome: HarnessRolloutOutcome) -> float | None: - """Binary terminal verifier + degeneracy penalties. Long-horizon credit is carried by the terminal reward, - propagated to every trained token through the group-relative advantage. - - - unscorable rollout -> None (dropped from the group baseline) - - never ran its code (no `bash`) -> -0.1 (kills blind-write / prose-dump / give-up) - - else BINARY base: all held-out tests pass -> 1.0; timed out or failed -> 0.0 - - minus a step penalty for tool calls beyond a budget (bounds runaway edit/bash loops), capped at 0.5 - """ - step_budget, step_penalty, step_penalty_cap = 20, 0.03, 0.5 - frac = outcome.env_reward - if frac is None: - return None - if outcome.tool_calls_by_name.get("bash", 0) == 0: - return -0.1 - base = 0.0 if outcome.timed_out else (1.0 if frac >= 1.0 - 1e-9 else 0.0) - over = max(0, outcome.tool_call_count - step_budget) - return base - min(step_penalty_cap, step_penalty * over) - - -# ============================================================================================================ -# Training -# ============================================================================================================ - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--model", default="Qwen/Qwen3-4B-Instruct-2507") - p.add_argument("--vllm-url", default="http://localhost:8000") - p.add_argument( - "--num-generations", type=int, default=8 - ) # >1 gives within-group pass/fail split -> nonzero advantage - p.add_argument("--max-inflight", type=int, default=8) # concurrent rollouts (each its own sandbox + proxy port) - p.add_argument("--max-completion-length", type=int, default=16384) - p.add_argument("--max-steps", type=int, default=100) - p.add_argument("--n-prompts", type=int, default=64) - p.add_argument("--learning-rate", type=float, default=1e-5) - p.add_argument("--temperature", type=float, default=1.0) - p.add_argument("--max-staleness", type=int, default=4) # lower -> fresher rollouts -> ratios near 1 (more stable) - p.add_argument("--seed", type=int, default=0) - p.add_argument("--output-dir", default="async_grpo_opencode") - p.add_argument("--project", default="opencode") - p.add_argument("--trackio-space-id", default=None) # optional: host the trackio dashboard on a HF Space - p.add_argument("--sandbox-root", default=None) # where per-rollout sandbox dirs live (default: a fresh tempdir) - args = p.parse_args() - - sandbox_root = args.sandbox_root or tempfile.mkdtemp(prefix="trl_opencode_") - LocalSubprocessSandboxBackend(sandbox_root).warmup() - - tokenizer = AutoTokenizer.from_pretrained(args.model) - rows, tests_by_id = build_dataset(n_prompts=args.n_prompts, seed=args.seed) - dataset = Dataset.from_list(rows) - - config = AsyncGRPOConfig( - output_dir=args.output_dir, - save_strategy="no", - per_device_train_batch_size=4, - num_generations=args.num_generations, - max_completion_length=args.max_completion_length, - max_steps=args.max_steps, - learning_rate=args.learning_rate, - temperature=args.temperature, - max_staleness=args.max_staleness, - vllm_server_base_url=args.vllm_url, - report_to="trackio", - project=args.project, - trackio_space_id=args.trackio_space_id, - log_completions=True, - ) - - worker = HarnessRolloutWorker( - harness_session_factory=partial(build_factory, sandbox_root, args.vllm_url, args.model, tests_by_id), - harness_adapter=None, # loop-owning: opencode runs its own loop; TRL reads the proxy trace - rollout_reward_fn=opencode_reward, # reward policy (binary verifier + degeneracy penalties) - model_name=args.model, - dataset=dataset, - reward_funcs=[], # reward comes from the harness verifier via rollout_reward_fn, not 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=args.temperature, - fork_threshold_tokens=1024, - log_completions=True, - num_completions_to_print=2, - ) - - trainer = AsyncGRPOTrainer( - model=args.model, - args=config, - train_dataset=dataset, - processing_class=tokenizer, - rollout_worker=worker, - ) - trainer.train() - - trainer.end() - - -if __name__ == "__main__": - main() diff --git a/examples/async_grpo_opencode/opencode_hf_sandbox.py b/examples/async_grpo_opencode/opencode_hf_sandbox.py deleted file mode 100644 index 3533c4a4798..00000000000 --- a/examples/async_grpo_opencode/opencode_hf_sandbox.py +++ /dev/null @@ -1,403 +0,0 @@ -# Copyright 2020-2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# /// script -# dependencies = [ -# "trl", -# "trackio", -# "datasets", -# "huggingface_hub>=1.22", -# "openenv @ git+https://github.com/huggingface/OpenEnv.git@86a180ede21e044f7929b9a7783ad83aa67d83a3", -# "openenv-opencode-env @ git+https://github.com/huggingface/OpenEnv.git@86a180ede21e044f7929b9a7783ad83aa67d83a3#subdirectory=envs/opencode_env", -# ] -# /// - -"""AsyncGRPO training of the real `opencode` coding agent (loop-owning) in REMOTE Hugging Face sandboxes. - -Same training path as `async_grpo_opencode.py`, but each rollout runs opencode in its own remote Hugging Face sandbox -(`HFSandboxBackend`) instead of a local subprocess, so rollouts scale out beyond a single node. opencode owns its -own tool loop; an in-sandbox proxy (`transparent_proxy` mode) forwards its calls to your vLLM server and captures -per-turn `(token_ids, logprobs)`. TRL reads that proxy trace, rebuilds training rows, scores the workspace with a -held-out verifier, and trains with GRPO. - -Two vLLM URLs, on purpose: - - `--vllm-url` (default `http://localhost:8000`): the TRAINER <-> vLLM link. Stays local for NCCL weight-sync. - - `--sandbox-vllm-url`: a url the remote sandboxes use to reach that same vLLM (the in-sandbox proxy forwards - there). Remote sandboxes cannot see `localhost`, so this must be reachable from outside: a public vLLM - endpoint, or a tunnel to your local one (see below). Not tied to any tunnel provider. - -Where opencode lives: nothing is installed per rollout. The default sandbox image -`ghcr.io/huggingface/openenv-opencode-sandbox:latest` pre-bakes the opencode CLI + the proxy (deps and -`interception.py`) under `/root`, so the harness skips the cold install. Pass `--sandbox-image python:3.12` to fall -back to cold-installing opencode + proxy deps per rollout. - -Task: competitive-coding problems from `agentica-org/DeepCoder-Preview-Dataset`. The agent writes `solution.py` -(reads stdin, prints stdout); the verifier runs it against the problem's HELD-OUT tests (never shown to the agent) -and returns a DENSE reward = fraction passed. `opencode_reward` keeps that dense signal and adds small degeneracy -penalties. This whole file is self-contained and every training-facing object is module-level (picklable), so the -rollout worker can pickle the factory + verifier into its spawned child process. - -Requirements: - - An OpenAI-compatible vLLM server (see below), reachable locally by the trainer and publicly by the sandboxes. - - An HF token with Jobs + Sandbox access in the environment (`HF_TOKEN`); each rollout is one HF sandbox. - -Run (2 GPUs: vLLM on one, trainer on the other; a tunnel exposes vLLM to the sandboxes): - -```sh -# Terminal 1 - serve the policy. Tool-calling + token-ids + NCCL weight-sync are all required. -CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-4B-Instruct-2507 \ - --host 0.0.0.0 --port 8000 \ - --enable-auto-tool-choice --tool-call-parser hermes \ - --logprobs-mode processed_logprobs \ - --return-tokens-as-token-ids \ - --max-model-len 98304 \ - --weight-transfer-config '{"backend":"nccl"}' - -# Terminal 2 - expose that vLLM publicly for the remote sandboxes. -cloudflared tunnel --no-autoupdate --url http://localhost:8000 # prints https://.trycloudflare.com - -# Terminal 3 - train. Trainer talks localhost (NCCL); sandboxes reach vLLM through the tunnel. -CUDA_VISIBLE_DEVICES=1 python examples/async_grpo_opencode/opencode_hf_sandbox.py \ - --model Qwen/Qwen3-4B-Instruct-2507 \ - --vllm-url http://localhost:8000 \ - --sandbox-vllm-url https://.trycloudflare.com -``` -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import random -from functools import partial -from typing import Any - -from datasets import Dataset, load_dataset -from opencode_env.config import OpenCodeConfig -from opencode_env.harness import OpenCodeSessionFactory -from opencode_env.sandbox import HFSandboxBackend, SandboxHandle -from opencode_env.task import OpenCodeTask -from openenv.core.harness import ResourceSession, ResourceSessionFactory, VerifyResult -from transformers import AutoTokenizer - -from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer -from trl.experimental.async_grpo.openenv_harness import ( - HarnessRolloutOutcome, - HarnessRolloutWorker, -) - - -# The HF sandbox image bakes opencode + the proxy under `/root`, and the sandbox execs as root, so `$HOME` and every -# harness path (`workdir`, `.opencode/bin`, `proxy/`) hang off `/root`. This is the only path difference from the -# local `async_grpo_opencode.py`, whose subprocess sandbox uses `/home/user`. -SANDBOX_IMAGE = "ghcr.io/huggingface/openenv-opencode-sandbox:latest" -SANDBOX_HOME = "/root" -WORKDIR = f"{SANDBOX_HOME}/workdir" - - -# ============================================================================================================ -# Dataset + held-out verifier -# ============================================================================================================ - -DATASET = "agentica-org/DeepCoder-Preview-Dataset" -DATASET_CONFIG = "primeintellect" # carries 11-93 stdin/stdout tests per problem -N_TESTS_EVAL = 12 # cap held-out tests scored per rollout (bounds verify latency); dense reward stays fine-grained -PER_TEST_TIMEOUT = 6 - - -def _instruction_id(instruction: str) -> str: - return hashlib.sha1(instruction.encode()).hexdigest() - - -def _clean_problem(problem: str) -> str: - # DeepCoder wraps problems in boilerplate that fights our tool-writing instruction ("...return the code."), which - # nudges a weak model to dump code in its reply instead of writing solution.py. Strip it so our instruction leads. - text = problem.strip() - prefix = "Solve the following coding problem using the programming language python:" - if text.startswith(prefix): - text = text[len(prefix) :].strip() - cut = len(text) - for marker in ("The input will be", "Now solve the problem", "Now solve this problem"): - idx = text.rfind(marker) - if idx != -1: - cut = min(cut, idx) - return text[:cut].strip() - - -def _coding_instruction(problem: str) -> str: - # Multi-turn loop: write -> run on the example cases -> read feedback -> fix -> repeat. The examples in the problem - # are the feedback signal; the held-out tests (the reward) stay hidden. - problem = _clean_problem(problem) - return ( - "Solve the problem below by writing a Python program to `solution.py` in the current working directory " - "(it does not exist yet, so create it with the `write` tool). The program must read its input from standard " - "input (stdin) and write ONLY the required answer to standard output (stdout).\n" - "Then TEST it against the example cases shown in the problem: run your program on each example input with " - "bash, e.g. `printf '' | python3 solution.py`, and compare its output to the expected output. " - "If an example fails or the program errors, use the `edit` tool to fix `solution.py`, then run it again. " - "Keep going until your program passes all the provided examples, then stop.\n\n" - f"PROBLEM:\n{problem}" - ) - - -# Runner executed INSIDE the sandbox: run solution.py against each held-out test with a hard time/mem/output cap, -# compare normalized stdout, print `SCORE: `. Whitespace-insensitive match (competitive judges usually are). -_RUNNER_SRC = r""" -import json, subprocess, sys - -PER_TEST_TIMEOUT = {per_test_timeout} - -def norm(s): - lines = [ln.rstrip() for ln in (s or "").replace("\r\n", "\n").split("\n")] - while lines and lines[-1] == "": - lines.pop() - return "\n".join(lines) - -def ok(got, exp): - if norm(got) == norm(exp): - return True - return (got or "").split() == (exp or "").split() - -tests = json.load(open("_tests.json")) -passed = 0 -for t in tests: - cmd = "ulimit -v 2000000 2>/dev/null; timeout %d python3 solution.py 2>/dev/null | head -c 2000000" % PER_TEST_TIMEOUT - try: - p = subprocess.run(["bash", "-c", cmd], input=t.get("input", ""), capture_output=True, text=True, - timeout=PER_TEST_TIMEOUT + 5) - if ok(p.stdout, t.get("output", "")): - passed += 1 - except Exception: - pass -print("SCORE: %.6f" % (passed / len(tests) if tests else 0.0)) -""" - - -def _run_dense_tests(sandbox: SandboxHandle, tests: list[dict[str, str]]) -> float: - """Run the sandbox's `solution.py` against `tests`; return the fraction passed.""" - if not tests or not sandbox.exists(f"{WORKDIR}/solution.py"): - return 0.0 - tests = tests[:N_TESTS_EVAL] - sandbox.write_text(f"{WORKDIR}/_tests.json", json.dumps(tests)) - sandbox.write_text(f"{WORKDIR}/_run_tests.py", _RUNNER_SRC.format(per_test_timeout=PER_TEST_TIMEOUT)) - r = sandbox.exec(f"cd {WORKDIR} && python3 _run_tests.py", timeout=PER_TEST_TIMEOUT * len(tests) + 30) - for line in (r.stdout or "").splitlines(): - if line.startswith("SCORE:"): - return float(line.split(":", 1)[1].strip()) - return 0.0 - - -class DeepCoderStdinVerifier: - """Dense stdin/stdout verifier. Holds the held-out test map (keyed by `sha1(instruction)`) so it survives the - pickle into the rollout child. `session.verify(...)` calls it as `verifier(sandbox, task)`.""" - - def __init__(self, tests_by_id: dict[str, list[dict[str, str]]]): - self._tests_by_id = tests_by_id - - def __call__(self, sandbox: SandboxHandle, task: OpenCodeTask) -> VerifyResult: - tests = self._tests_by_id.get(_instruction_id(task.instruction)) - return VerifyResult(env_reward=_run_dense_tests(sandbox, tests or []), done=True) - - -def build_dataset(n_prompts: int, seed: int) -> tuple[list[dict], dict[str, list[dict[str, str]]]]: - """Return `(rows, tests_by_id)`: the prompt rows (problem statement only) and the held-out test map for the - verifier. Tests are NOT put in the prompt - the agent only sees the statement (which includes sample cases).""" - rows = list(load_dataset(DATASET, DATASET_CONFIG, split="train")) - random.Random(seed).shuffle(rows) - out: list[dict] = [] - tests_by_id: dict[str, list[dict[str, str]]] = {} - for r in rows: - if len(out) >= n_prompts: - break - raw = r["tests"] - if isinstance(raw, str): - raw = json.loads(raw) - tests = [ - {"input": t.get("input", ""), "output": t.get("output", "")} - for t in raw - if t.get("type", "stdin_stdout") == "stdin_stdout" and t.get("output") is not None - ] - if len(tests) < 3: # need enough tests for a meaningful dense fraction - continue - instruction = _coding_instruction(r["problem"]) - tests_by_id[_instruction_id(instruction)] = tests[:N_TESTS_EVAL] - out.append({"prompt": [{"role": "user", "content": instruction}]}) - return out, tests_by_id - - -# ============================================================================================================ -# opencode session factory (remote HF sandbox + in-sandbox proxy) -# ============================================================================================================ - - -class OpencodeTaskFactory(ResourceSessionFactory): - """Adapts the worker's `create(prompt=messages, seed, episode_id)` onto `OpenCodeSessionFactory`, which wants an - `OpenCodeTask`: pull the instruction out of the last user message.""" - - def __init__(self, inner: OpenCodeSessionFactory): - self._inner = inner - - def create(self, task: Any, seed: int | None = None, episode_id: str | None = None) -> ResourceSession: - instruction = task[-1]["content"] if isinstance(task, list) and task else str(task) - return self._inner.create(instruction, seed=seed, episode_id=episode_id) - - -def build_factory( - sandbox_vllm_url: str, model: str, tests_by_id: dict, image: str, flavor: str, *, sampling: dict -) -> OpencodeTaskFactory: - config = OpenCodeConfig( - extra_opencode_json={"agent": {"build": {"temperature": sampling["temperature"], "top_p": sampling["top_p"]}}}, - provider="openai_compatible", - base_url=f"{sandbox_vllm_url}/v1", # the in-sandbox proxy forwards here; remote, so a public url (tunnel) - model=model, # proxy --model-override forces this exact id on upstream requests - sandbox_home=SANDBOX_HOME, # the HF sandbox execs as root; opencode + proxy are baked under /root - agent_timeout_s=600.0, # remote hop adds latency vs the local backend; give the edit/bash loop more room - disabled_tools=["webfetch", "question", "task"], # no web, no user, no sub-agents - run_format="json", - proxy_max_tokens_cap=8192, # keep each turn's completion + the growing multi-turn prompt under --max-model-len - ) - inner = OpenCodeSessionFactory( - config=config, - sandbox_backend=HFSandboxBackend(image=image, flavor=flavor), # each rollout = its own remote HF sandbox - mode="transparent_proxy", - sampling=sampling, - verifier=DeepCoderStdinVerifier(tests_by_id), - ) - return OpencodeTaskFactory(inner) - - -# ============================================================================================================ -# Rollout reward -# ============================================================================================================ - - -def opencode_reward(outcome: HarnessRolloutOutcome) -> float | None: - """Dense terminal verifier + degeneracy penalties. Long-horizon credit is carried by the terminal reward, - propagated to every trained token through the group-relative advantage. - - - unscorable rollout (no verifier score) -> None (dropped from the group baseline) - - never ran its code (no `bash`) -> -0.1 (kills blind-write / prose-dump / give-up) - - else DENSE base: the fraction of held-out tests passed (partial credit); timed out -> 0.0 - - minus a step penalty for tool calls beyond a budget (bounds runaway edit/bash loops), capped at 0.5 - """ - step_budget, step_penalty, step_penalty_cap = 30, 0.03, 0.5 - frac = outcome.env_reward - bash = outcome.tool_calls_by_name.get("bash", 0) - if frac is None: - return None - if bash == 0: - return -0.1 - base = 0.0 if outcome.timed_out else frac # dense: fraction of held-out tests passed (partial credit) - over = max(0, outcome.tool_call_count - step_budget) - return base - min(step_penalty_cap, step_penalty * over) - - -# ============================================================================================================ -# Training -# ============================================================================================================ - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--model", default="Qwen/Qwen3-4B-Instruct-2507") - p.add_argument("--vllm-url", default="http://localhost:8000") # trainer <-> vLLM (weight sync, NCCL): local - p.add_argument("--sandbox-vllm-url", required=True) # public url the remote sandboxes reach vLLM through (tunnel) - p.add_argument( - "--sandbox-image", default=SANDBOX_IMAGE - ) # pre-baked opencode+proxy; use python:3.12 to cold-install - p.add_argument("--sandbox-flavor", default="cpu-basic") # the agent only runs python/bash, no GPU needed - p.add_argument( - "--num-generations", type=int, default=8 - ) # >1 gives within-group pass/fail split -> nonzero advantage - p.add_argument("--max-inflight", type=int, default=8) # concurrent rollouts (each its own remote sandbox) - p.add_argument("--max-completion-length", type=int, default=16384) - p.add_argument("--max-steps", type=int, default=100) - p.add_argument("--n-prompts", type=int, default=64) - p.add_argument("--learning-rate", type=float, default=1e-5) - p.add_argument("--temperature", type=float, default=1.0) - p.add_argument("--max-staleness", type=int, default=4) # lower -> fresher rollouts -> ratios near 1 (more stable) - p.add_argument("--seed", type=int, default=0) - p.add_argument("--output-dir", default="async_grpo_opencode_hf_sandbox") - p.add_argument("--project", default="opencode-hf-sandbox") - p.add_argument("--trackio-space-id", default=None) # optional: host the trackio dashboard on a HF Space - p.add_argument("--push-to-hub", action="store_true") # push the trained policy to the Hub at --hub-model-id - p.add_argument("--hub-model-id", default=None) - p.add_argument("--optim", default="adamw_torch") # e.g. paged_adamw_8bit to fit a larger policy on one GPU - p.add_argument("--gradient-checkpointing", action="store_true") # trade compute for memory on a larger policy - p.add_argument("--gradient-accumulation-steps", type=int, default=1) # more prompts per step -> smoother reward - args = p.parse_args() - - tokenizer = AutoTokenizer.from_pretrained(args.model) - rows, tests_by_id = build_dataset(n_prompts=args.n_prompts, seed=args.seed) - dataset = Dataset.from_list(rows) - - config = AsyncGRPOConfig( - output_dir=args.output_dir, - save_strategy="no", - per_device_train_batch_size=4, - gradient_accumulation_steps=args.gradient_accumulation_steps, - num_generations=args.num_generations, - max_completion_length=args.max_completion_length, - max_steps=args.max_steps, - learning_rate=args.learning_rate, - temperature=args.temperature, - max_staleness=args.max_staleness, - vllm_server_base_url=args.vllm_url, - report_to="trackio", - project=args.project, - trackio_space_id=args.trackio_space_id, - log_completions=True, - optim=args.optim, - gradient_checkpointing=args.gradient_checkpointing, - push_to_hub=args.push_to_hub, - hub_model_id=args.hub_model_id, - ) - - worker = HarnessRolloutWorker( - harness_session_factory=partial( - build_factory, args.sandbox_vllm_url, args.model, tests_by_id, args.sandbox_image, args.sandbox_flavor - ), - harness_adapter=None, # loop-owning: opencode runs its own loop; TRL reads the proxy trace - rollout_reward_fn=opencode_reward, # reward policy (binary verifier + degeneracy penalties) - model_name=args.model, - dataset=dataset, - reward_funcs=[], # reward comes from the harness verifier via rollout_reward_fn, not 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=args.temperature, - fork_threshold_tokens=1024, - log_completions=True, - num_completions_to_print=2, - ) - - trainer = AsyncGRPOTrainer( - model=args.model, - args=config, - train_dataset=dataset, - processing_class=tokenizer, - rollout_worker=worker, - ) - trainer.train() - if args.push_to_hub: - trainer.push_to_hub() - - trainer.end() - - -if __name__ == "__main__": - main()