Async grpo OpenEnv harness rollout

tests/experimental/test_async_grpo_trainer.py

@@ -97,8 +97,8 @@ def _fill_queue(self):
     def start(self):
         self._fill_queue()
 
-    def update_model_version(self, version):
-        self._model_version = version
+    def update_model_version(self, model_version):
+        self._model_version = model_version
         self._fill_queue()
 
     def stop(self):

trl/experimental/async_grpo/async_grpo_trainer.py

@@ -20,6 +20,7 @@
 from collections import defaultdict
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass
+from multiprocessing.queues import Queue as MPQueue
 from typing import Any, Protocol
 
 import torch
@@ -61,11 +62,14 @@ class RolloutWorkerProtocol(Protocol):
     runs reward models on their own GPUs.
 
     Attributes:
-        rollout_buffer (`queue.Queue`):
-            Queue the trainer drains; the worker pushes scored `RolloutSample`s onto it.
+        rollout_buffer (`queue.Queue` or `multiprocessing.queues.Queue`):
+            Queue the trainer drains; the worker pushes scored `RolloutSample`s onto it. The two queue types are
+            structurally identical (`get` / `put_nowait` / `qsize`) but nominally unrelated, so both are allowed: the
+            default [`AsyncRolloutWorker`] runs its loop in a spawned process and uses `multiprocessing.Queue`, while
+            an in-process worker uses `queue.Queue`.
     """
 
-    rollout_buffer: queue.Queue
+    rollout_buffer: queue.Queue | MPQueue
 
     def start(self) -> None:
         """Begin producing rollouts. Called once on train begin, after the initial weight sync."""
@@ -75,7 +79,7 @@ def stop(self) -> None:
         """Stop the worker and release its resources. Called on train end."""
         ...
 
-    def update_model_version(self, version: int) -> None:
+    def update_model_version(self, model_version: int) -> None:
         """Tell the worker which policy version is now live, so it can tag or discard stale samples."""
         ...
 

trl/experimental/async_grpo/async_rollout_worker.py

@@ -179,6 +179,7 @@ class RolloutGroup:
     tool_failure_counts: list[int]
     model_version: int
     env_rewards: list[tuple[type, float] | None]
+    rollout_rewards: list[float | None]
     queued_at: float = 0.0
 
 
@@ -237,6 +238,7 @@ def _watch():
 
 
 def _child_main(
+    loop_cls: "type[_AsyncRolloutLoop]",
     loop_kwargs: dict[str, Any],
     samples_queue: MPQueue,
     model_version_value: MPValue,
@@ -252,7 +254,7 @@ def _child_main(
 
     PartialState()
 
-    rollout_loop = _AsyncRolloutLoop(
+    rollout_loop = loop_cls(
         **loop_kwargs,
         rollout_buffer=samples_queue,
         model_version_value=model_version_value,
@@ -277,6 +279,9 @@ class _AsyncRolloutLoop:
     policy version from the shared `mp.Value` (`model_version_value`).
     """
 
+    # Class attribute  declares that this loop produces its own per-rollout reward
+    _provides_rollout_reward: bool = False
+
     def __init__(
         self,
         *,
@@ -397,7 +402,7 @@ def __init__(
 
         # At least one reward source is required: either `reward_funcs`, or an environment that owns the reward via a
         # `get_reward` method.
-        if not self.reward_funcs and not self._env_reward_types:
+        if not self.reward_funcs and not self._env_reward_types and not self._provides_rollout_reward:
             raise ValueError(
                 "No reward source provided. Pass `reward_funcs`, or an `environment_factory` whose environment "
                 "defines a `get_reward` method."
@@ -539,6 +544,7 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
                             tool_failure_counts=[],
                             model_version=self.model_version,
                             env_rewards=[],
+                            rollout_rewards=[],
                         )
                         pending_completed[group_id] = 0
 
@@ -567,6 +573,7 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
                         sequences,
                         tool_call_count,
                         tool_failure_count,
+                        rollout_reward,
                     ) = task.result()
                     group = pending_groups[group_id]
                     group.prompts.append(prompt)
@@ -575,6 +582,7 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
                     group.completions_sequences.append(sequences)
                     group.tool_call_counts.append(tool_call_count)
                     group.tool_failure_counts.append(tool_failure_count)
+                    group.rollout_rewards.append(rollout_reward)
                     # The environment owns the reward: score it now, while this rollout's environment still holds its
                     # final state and before returning it to the pool. `get_reward` may be async awaiting yields to
                     # inflight requests instead of halting them. The env is returned to the pool only after scoring, so
@@ -691,16 +699,19 @@ def _repeat_iterator(self) -> Iterator[tuple[int, dict[str, Any]]]:
 
     async def _generate_one(
         self, prompt: Messages, tool_dict: dict[str, Callable], tools: list[Callable]
-    ) -> tuple[list[dict[str, str]], list[int], list[TrainingSequence], int, int]:
+    ) -> tuple[list[dict[str, str]], list[int], list[TrainingSequence], int, int, float | None]:
         """Roll out one conversation, re-tokenizing the whole message list each turn and reconciling drift.
 
         Every turn renders the full conversation through the chat template, generates, records the turn, and feeds the
         tool result back as a message (re-tokenized on the next whole-conversation pass). At the end,
         `_chain_to_sequences` reconciles re-tokenization drift into one or more training rows: a clean append stays one
         row, a rewrite (dropped reasoning, summarized history) forks a new row. Rebuilding `prompt_ids` from the
         message list each turn (instead of gluing tokens on the end and never looking back) is what lets the reconciler
-        catch rewrites. Returns the completion messages, their token ids, and the reconciled training rows (each row
-        carries its own `input_ids`).
+        catch rewrites.
+
+        Returns `(completion messages, their token ids, reconciled training rows (each carries its own `input_ids`),
+        tool-call count, tool-failure count, rollout reward)`; the built-in worker's rollout reward is always `None`
+        (it scores via `reward_funcs`/env), a trailing slot loop subclasses override.
         """
         messages = list(prompt)  # a MESSAGE list, not a token list
         rollout_id = uuid.uuid4().hex
@@ -735,7 +746,7 @@ async def _generate_one(
             messages.extend(tool_messages)  # tool result goes back as a MESSAGE, re-tokenized next turn
             iteration_num += 1
         sequences = _chain_to_sequences(turns, rollout_id, self._fork_threshold_tokens)  # >= 1 row per conversation
-        return completion, completion_ids, sequences, tool_call_count, tool_failure_count
+        return completion, completion_ids, sequences, tool_call_count, tool_failure_count, None
 
     def _execute_tool_calls(
         self, tool_calls: list[dict[str, Any]], tool_dict: dict[str, Callable]
@@ -797,6 +808,9 @@ async def _score_group(self, group: RolloutGroup) -> list[RolloutSample]:
             column = [r[1] if (r is not None and r[0] is env_type) else None for r in group.env_rewards]
             all_rewards = [*all_rewards, column]
 
+        if self._provides_rollout_reward:
+            all_rewards = [*all_rewards, list(group.rollout_rewards)]
+
         # Reward funcs may return None per-sample (unparseable gold). Convert to NaN. A completion
         # for which every func returned None is unscorable: nansum would give 0 and the row would
         # pull the policy away from actually-correct answers.
@@ -901,6 +915,8 @@ class AsyncRolloutWorker:
     execute on CPU.
     """
 
+    _loop_cls: "type[_AsyncRolloutLoop]" = _AsyncRolloutLoop
+
     def __init__(
         self,
         *,
@@ -960,6 +976,7 @@ def start(self) -> None:
         self._process = self._mp_ctx.Process(
             target=_child_main,
             args=(
+                self._loop_cls,
                 self._loop_kwargs,
                 self.rollout_buffer,
                 self._model_version_value,

trl/experimental/async_grpo/openenv_harness.py

@@ -0,0 +1,193 @@
+# Copyright 2020-2025 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.
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import hashlib
+import json
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+from openenv.core.harness import HarnessRunLimits, ModelStepResult
+from openenv.core.llm_client import LLMResponse, ToolCall
+
+from ...chat_template_utils import parse_response
+from .async_rollout_worker import (
+    AsyncRolloutWorker,
+    TurnRecord,
+    _AsyncRolloutLoop,
+    _chain_to_sequences,
+)
+
+
+Message = dict[str, Any]
+
+
+def _tools_to_schema(tools: list) -> list[dict] | None:
+    """Convert OpenEnv `Tool`s (MCP spec) into the OpenAI function schema `apply_chat_template` expects."""
+    if not tools:
+        return None
+    return [
+        {
+            "type": "function",
+            "function": {"name": t.name, "description": t.description, "parameters": t.input_schema},
+        }
+        for t in tools
+    ]
+
+
+def _msg_to_llm_response(msg: Message) -> LLMResponse:
+    """Convert a `parse_response` assistant message into the `LLMResponse` the harness adapter expects."""
+    tool_calls = []
+    for i, tc in enumerate(msg.get("tool_calls") or []):
+        fn = tc.get("function", tc)
+        args = fn.get("arguments", {})
+        if isinstance(args, str):
+            try:
+                args = json.loads(args)
+            except json.JSONDecodeError:
+                args = {}
+        tool_calls.append(ToolCall(id=tc.get("id") or f"call_{i}", name=fn["name"], args=args))
+    return LLMResponse(content=msg.get("content") or "", tool_calls=tool_calls)
+
+
+class _HarnessRolloutLoop(_AsyncRolloutLoop):
+    """`_AsyncRolloutLoop` whose `_generate_one` drives an OpenEnv session instead of TRL's own turn loop."""
+
+    _provides_rollout_reward = True
+
+    def __init__(self, *, harness_session_factory, harness_adapter=None, **loop_kwargs):
+        super().__init__(**loop_kwargs)
+        self._factory = harness_session_factory
+        # An adapter (e.g. MCPHarnessAdapter) selects white-box (TRL samples each turn); `None` selects loop-owning
+        # (the agent runs its own loop and we read its proxy trace).
+        self._adapter = harness_adapter
+        self._limits = HarnessRunLimits(
+            max_turns=self.max_tool_calling_iterations if self.max_tool_calling_iterations is not None else 8,
+            sampling={"temperature": self.temperature, "max_tokens": self.max_tokens},
+        )
+        # Sized to max_inflight so sessions run concurrently without queuing or per-rollout thread churn.
+        self._session_pool = ThreadPoolExecutor(
+            max_workers=max(1, self.max_inflight_tasks), thread_name_prefix="harness-session"
+        )
+        self.reward_func_names.append("harness_reward")
+
+    async def _generate_one(self, prompt, tool_dict, tools):
+        # OpenEnv's harness layer is synchronous, so run the whole session on the pool.
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(self._session_pool, self._run_session, prompt)
+
+    def _run_session(self, prompt):
+        """Drive one OpenEnv session to completion and return the `_generate_one` tuple + the verify reward."""
+        rollout_id = uuid.uuid4().hex
+        # Stable per-group seed so seed-driven factories hand every generation of a prompt the same task.
+        seed = int(hashlib.sha256(repr(prompt).encode()).hexdigest(), 16) % (2**31)
+        session = self._factory.create(prompt, seed=seed, episode_id=rollout_id)
+        try:
+            if self._adapter is not None:
+                # white-box: the adapter runs the tool loop, calling `_sample_turn` each turn.
+                turns: list[TurnRecord] = []
+                result = self._adapter.run_white_box(
+                    functools.partial(self._sample_turn, turns), session, self._limits
+                )
+                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)
+            else:
+                # loop-owning: the agent hits vLLM through its own proxy; read the per-turn trace.
+                session.wait_for_completion()
+                trace = session.fetch_proxy_trace()
+                turns = _turns_from_trace(trace, self.tokenizer)
+                completion = _messages_from_trace(trace)
+                tool_call_count = tool_failure_count = 0
+            verify = session.verify(completion)
+            reward = float(verify.env_reward) if verify.env_reward is not None else None
+            sequences = _chain_to_sequences(turns, rollout_id, self._fork_threshold_tokens)
+            completion_ids = [tid for turn in turns for tid in turn.output_ids]
+            return completion, completion_ids, sequences, tool_call_count, tool_failure_count, reward
+        finally:
+            session.close()
+
+    def _sample_turn(self, turns: list[TurnRecord], messages, tools, sampling) -> ModelStepResult:
+        """OpenEnv `ModelStep`: sample one assistant turn against vLLM and record a `TurnRecord` into `turns`."""
+        prompt_ids = self.tokenizer.apply_chat_template(
+            messages,
+            tools=_tools_to_schema(tools),
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=False,
+            chat_template=self.chat_template,
+            **self.chat_template_kwargs,
+        )
+        # ModelStep is sync on a pool thread; bridge the async vLLM POST onto the loop's event loop.
+        turn_ids, logprobs = asyncio.run_coroutine_threadsafe(self._generate_one_turn(prompt_ids), self._loop).result()
+        turns.append(TurnRecord(prompt_ids, turn_ids, logprobs))
+        message = parse_response(self.tokenizer, turn_ids, prefix=prompt_ids)
+        return ModelStepResult(
+            response=_msg_to_llm_response(message), prompt_ids=prompt_ids, completion_ids=turn_ids, logprobs=logprobs
+        )
+
+
+def _trace_output_ids(entry: dict) -> list[int]:
+    """Generated token ids for one proxy-trace turn.
+
+    Prefer `completion_token_ids`; if empty, recover them from `completion_tokens`, which vLLM renders as
+    `"token_id:{id}"` when launched with `--return-tokens-as-token-ids`, avoiding a re-encode of the decoded text.
+    """
+    ids = entry.get("completion_token_ids") or []
+    if ids:
+        return list(ids)
+    return [
+        int(t[len("token_id:") :]) for t in (entry.get("completion_tokens") or []) if str(t).startswith("token_id:")
+    ]
+
+
+def _turns_from_trace(trace: list[dict], tokenizer) -> list[TurnRecord]:
+    """Loop-owning path: rebuild per-turn `TurnRecord`s from a proxy trace. Re-tokenize each request's messages
+    (passing its `tools` so the prompt matches what the upstream rendered); ids + logprobs come from the capture."""
+    turns = []
+    for entry in trace:
+        request = entry["request"]
+        prompt_ids = tokenizer.apply_chat_template(
+            request["messages"],
+            tools=request.get("tools"),
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=False,
+        )
+        turns.append(TurnRecord(prompt_ids, _trace_output_ids(entry), entry["per_token_logps"]))
+    return turns
+
+
+def _messages_from_trace(trace: list[dict]) -> list[Message]:
+    """Loop-owning path: the transcript is the last request's messages plus the final assistant reply."""
+    if not trace:
+        return []
+    last = trace[-1]
+    content = last["response"]["choices"][0]["message"].get("content", "")
+    return list(last["request"]["messages"]) + [{"role": "assistant", "content": content}]
+
+
+class HarnessRolloutWorker(AsyncRolloutWorker):
+    """AsyncGRPO rollout worker that drives an OpenEnv `ResourceSessionFactory`.
+
+    Construct it with the usual `AsyncRolloutWorker` kwargs plus `harness_session_factory` (and optionally
+    `harness_adapter`), then inject it via `AsyncGRPOTrainer(rollout_worker=...)`. Only the spawned child's loop class
+    differs.
+    """
+
+    _loop_cls = _HarnessRolloutLoop