diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py index bff753de76a..a7c08e91774 100644 --- a/tests/experimental/test_async_grpo_trainer.py +++ b/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, ) diff --git a/trl/experimental/async_grpo/async_grpo_trainer.py b/trl/experimental/async_grpo/async_grpo_trainer.py index a213b8be91c..bf562af5963 100644 --- a/trl/experimental/async_grpo/async_grpo_trainer.py +++ b/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.""" ... diff --git a/trl/experimental/async_grpo/async_rollout_worker.py b/trl/experimental/async_grpo/async_rollout_worker.py index 90817c608e3..8d17c0f0960 100644 --- a/trl/experimental/async_grpo/async_rollout_worker.py +++ b/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,7 +704,7 @@ 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 @@ -704,8 +712,11 @@ async def _generate_one( `_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, diff --git a/trl/experimental/async_grpo/openenv_harness.py b/trl/experimental/async_grpo/openenv_harness.py new file mode 100644 index 00000000000..faf7605cbed --- /dev/null +++ b/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