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):
@@ -560,13 +560,13 @@ async def _generate_one_turn(prompt_ids):
     loop._generate_one_turn = _generate_one_turn
     loop._execute_tool_calls = lambda tool_calls, tool_dict: ([{"role": "tool", "name": "t", "content": "ok"}], 1, 0)
 
-    # _generate_one returns (completion, completion_ids, sequences, n_calls, n_failures).
+    # _generate_one returns (completion, completion_ids, sequences, n_calls, n_failures, rollout_reward).
     return asyncio.run(loop._generate_one([{"role": "user", "content": "hi"}], {}, []))
 
 
 class TestRolloutLoop(TrlTestCase):
     def test_single_turn_no_tool_call(self, monkeypatch):
-        completion, completion_ids, sequences, n_calls, n_failures = _run(
+        completion, completion_ids, sequences, n_calls, n_failures, _ = _run(
             monkeypatch,
             prompt_ids=[[1, 2, 3]],
             turns=[([10, 11], [-0.1, -0.2])],
@@ -582,7 +582,7 @@ def test_single_turn_no_tool_call(self, monkeypatch):
 
     def test_clean_two_turns_stay_one_row(self, monkeypatch):
         # Turn 2's re-tokenized prompt starts with what we held (gen tokens + tool tokens) -> CLEAN.
-        completion, completion_ids, sequences, n_calls, n_failures = _run(
+        completion, completion_ids, sequences, n_calls, n_failures, _ = _run(
             monkeypatch,
             prompt_ids=[[1, 2, 3], [1, 2, 3, 10, 11, 20, 21]],
             turns=[([10, 11], [-0.1, -0.2]), ([30, 31], [-0.3, -0.4])],
@@ -598,7 +598,7 @@ def test_clean_two_turns_stay_one_row(self, monkeypatch):
 
     def test_history_rewrite_forks_into_two_rows(self, monkeypatch):
         # Turn 2's prompt diverges inside turn 1's answer and the new turn is >= fork_threshold -> FORK.
-        _, _, sequences, _, _ = _run(
+        _, _, sequences, _, _, _ = _run(
             monkeypatch,
             prompt_ids=[[1, 2, 3], [1, 2, 3, 99, 88, 77]],
             turns=[([10, 11, 12, 13], [-0.1] * 4), ([30, 31, 32], [-0.2] * 3)],
@@ -615,7 +615,7 @@ def test_history_rewrite_forks_into_two_rows(self, monkeypatch):
 
     def test_max_tool_calling_iterations_caps_turns(self, monkeypatch):
         # max_iters=0: even though turn 1 is a tool call, the loop breaks before executing it.
-        completion, _, sequences, n_calls, _ = _run(
+        completion, _, sequences, n_calls, _, _ = _run(
             monkeypatch,
             prompt_ids=[[1, 2, 3]],
             turns=[([10, 11], [-0.1, -0.2])],
@@ -654,6 +654,7 @@ def _group(completions_sequences, completions_ids):
         tool_failure_counts=[0] * n,
         model_version=7,
         env_rewards=[None] * n,
+        rollout_rewards=[None] * n,
     )
 
 

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."
@@ -544,6 +549,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
 
@@ -572,6 +578,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)
@@ -580,6 +587,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
@@ -696,16 +704,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
@@ -740,7 +751,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]
@@ -802,6 +813,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.
@@ -906,6 +920,8 @@ class AsyncRolloutWorker:
     execute on CPU.
     """
 
+    _loop_cls: "type[_AsyncRolloutLoop]" = _AsyncRolloutLoop
+
     def __init__(
         self,
         *,
@@ -965,6 +981,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,