# Copyright (c) ModelScope Contributors. All rights reserved. # Multi-turn Rollout Schedulers for GRPO training. import asyncio import json from abc import ABC from copy import deepcopy from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from swift.infer_engine.protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, RequestConfig, RolloutInferRequest, RolloutOutput) from swift.template import Messages from swift.utils import remove_response from .gym_env import Env, envs if TYPE_CHECKING: # Imported only for type hints; importing it at runtime pulls in vllm, which would make # `swift.rollout` (and thus GRPO trainer init) hard-require vllm even when use_vllm=False. from swift.infer_engine import GRPOVllmEngine class RolloutScheduler(ABC): # Single Turn Rollout Scheduler def __init__(self, infer_engine: Optional['GRPOVllmEngine'] = None, max_turns: Optional[int] = None, *args, **kwargs): self.infer_engine = infer_engine # Tokenizer can be passed explicitly (e.g., in colocate mode where infer_engine may be None) self._tokenizer = kwargs.get('tokenizer', None) self.max_turns = max_turns # ------------------------------------------------------------------ # Universal hooks — called by BOTH ``run()`` (server mode) and # ``run_multi_turn()`` (colocate mode). Override these to inject # environment lifecycle logic (e.g. gym env.reset / env.step) without # overriding the full ``run()`` method. # # Hooks are async so that gym environments (whose reset/step are async) # can be awaited directly. In server mode ``run()`` awaits them natively; # in colocate mode ``run_multi_turn()`` drives them via a dedicated loop. # ------------------------------------------------------------------ async def on_trajectory_start(self, requests: List['RolloutInferRequest']) -> None: """Called before the first inference turn to initialize per-trajectory state. Mutate ``requests`` in place (e.g. inject env initial observation). Default: no-op. """ pass async def on_turn_end(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> Dict[str, Any]: """Called after assistant message is appended, before ``check_finished``. Use this to advance environment state (e.g. ``env.step``) and surface per-turn metadata. Returns: Dict with optional keys: - 'done' (bool): if present, overrides ``check_finished`` result - 'rollout_infos' (dict): merged into the trajectory's accumulated infos Default: empty dict (no-op). """ return {} async def async_infer(self, infer_requests: List[Union['RolloutInferRequest', Dict[str, Any]]], request_config: 'RequestConfig', *, use_tqdm: Optional[bool] = None, **kwargs) -> List['RolloutOutput']: """ Perform asynchronous batched inference for multiple rollout requests. This method serves as the main entry point for multi-round training inference. It executes the `run` method for each inference request concurrently and aggregates the results into a single flattened list. Each inference request can be either a `RolloutInferRequest` instance or a dictionary that can be converted into one. The results from all requests are collected asynchronously using the underlying inference engine. Args: infer_requests (List[Union[RolloutInferRequest, Dict[str, Any]]]): A list of inference requests. Each request can be either: - A `RolloutInferRequest` object. - A dictionary containing the fields required to initialize a `RolloutInferRequest`. request_config (RequestConfig): Configuration object specifying inference settings. Must satisfy `request_config.n == 1`, as only single-response generation is supported. use_tqdm (Optional[bool], optional): Whether to display a progress bar during batch inference. If `None`, it defaults to `True` when there are multiple requests, otherwise `False`. **kwargs: Additional arguments forwarded to the underlying `run` method. Returns: List[RolloutOutput]: A list of RolloutOutput objects corresponding to the provided inference requests. Raises: AssertionError: If `request_config.n` is not equal to `1`. Notes: - Internally, this method converts dict-based requests into `RolloutInferRequest` instances. - Uses `infer_engine._batch_infer_stream` to perform concurrent execution. - The returned list is guaranteed to be flattened, even if individual tasks return lists of responses. """ assert request_config.n == 1 async def _infer_async_single(infer_request: Union['RolloutInferRequest', Dict[str, Any]], request_config: 'RequestConfig', **kwargs): if isinstance(infer_request, Dict): infer_request = RolloutInferRequest(**infer_request) return await self.run(infer_request, request_config, **kwargs) tasks = [_infer_async_single(infer_request, request_config, **kwargs) for infer_request in infer_requests] if use_tqdm is None: use_tqdm = len(infer_requests) > 1 # Execute all tasks and flatten the results results = await self.infer_engine._batch_infer_stream(tasks, request_config.stream, use_tqdm, None) # Flatten the results since each task may return a list flattened_results = [] for result in results: if isinstance(result, list): flattened_results.extend(result) else: flattened_results.append(result) return flattened_results async def run(self, infer_request: 'RolloutInferRequest', request_config: 'RequestConfig', **kwargs) -> 'RolloutOutput': response: 'ChatCompletionResponse' = await self.infer_engine.infer_async(infer_request, request_config, **kwargs) response_token_ids = response.choices[0].token_ids response_loss_mask = [1] * len(response_token_ids) return RolloutOutput( response=response, messages=infer_request.messages, response_token_ids=[response_token_ids], response_loss_mask=[response_loss_mask], rollout_infos={'num_turns': 1}) def __getattr__(self, key: str): try: return object.__getattribute__(self, key) except AttributeError: pass try: infer_engine = object.__getattribute__(self, 'infer_engine') if hasattr(infer_engine, key): return getattr(infer_engine, key) if hasattr(infer_engine.engine, key): return getattr(infer_engine.engine, key) except AttributeError: raise AttributeError(f'{type(self).__name__} object has no attribute {key}') @property def engine(self): return self.infer_engine @property def tokenizer(self): """Get tokenizer, prioritizing explicitly passed tokenizer over infer_engine's tokenizer.""" if self._tokenizer is not None: return self._tokenizer if self.infer_engine is not None: return self.infer_engine.tokenizer return None class MultiTurnScheduler(RolloutScheduler, ABC): """ Abstract base class for multi-turn rollout scheduling. Provides default implementation for multi-turn conversation management with two customization approaches: 1. FULL CUSTOMIZATION: Override the `run()` method to implement completely custom multi-turn logic. - Gives full control over the rollout process - Must handle all turn management and termination logic 2. PARTIAL CUSTOMIZATION: Implement the required `step()` method and optionally override `check_finished()` - Uses MultiTurnScheduler's run() method infrastructure - Only need to implement turn transition logic in step() - Optionally customize termination conditions Note: You must implement at least one of these approaches in your subclass. Options: - If each round's response_token_ids are included in the RolloutOutput, the Trainer can skip encoding the completion text into token_ids when calculating loss. This avoids potential training inconsistencies due to asymmetric encode/decode behavior. See: https://github.com/0russwest0/Agent-R1/issues/30#issuecomment-2826155367 - If both response_token_ids and response_loss_mask are returned in the RolloutOutput, you can manually control the loss mask for each token. The Trainer will use the provided loss_mask values directly when computing the loss. Note: Returning response_loss_mask requires that response_token_ids are also returned, as the two must be aligned in length for correct loss computation. You can refer to MathTipsScheduler as an example of how to use response_token_ids and response_loss_mask. Loss mask configuration: During rollout, some parts of the completion (e.g., environment observations embedded in completion) may need to be masked out from loss computation. There are two supported strategies: 1. Use the built-in `loss_scale` parameter in ms-swift and do not return response token ids. 2. Return response_token_ids along with a corresponding response_loss_mask (of equal length) to indicate the loss mask for each token. # noqa """ async def run(self, infer_request: 'RolloutInferRequest', request_config: 'RequestConfig', **kwargs) -> Union['RolloutOutput', List['RolloutOutput']]: """Execute multi-turn conversation rollout with built-in turn management logic. This implements the default multi-turn interaction flow that can be overridden to customize conversation handling behavior. The default logic provides: 1. Automatic conversation turn management and stopping conditions 2. Seamless message accumulation across multiple turns 3. Response token tracking and loss mask management 4. Configurable early stopping mechanisms Args: infer_request: The initial inference request containing conversation messages request_config: Configuration parameters for the inference request **kwargs: Additional inference parameters passed to the engine Returns: RolloutOutput containing the complete conversation history and metadata, or a list of outputs for batched requests Customization Approaches: - Override check_finished() to implement custom stopping criteria - Override step() to customize turn-to-turn transition logic - Override this entire run() method for completely custom multi-turn behavior Important Notes: - Method overriding is only supported when using server mode (swift rollout) with vllm_use_async_engine=True - Custom implementations must maintain async/await compatibility - Ensure proper handling of conversation state across turns Example: class CustomScheduler(MultiTurnScheduler): async def run(self, infer_request, request_config, **kwargs): # Implement custom multi-turn conversation logic # Must return RolloutOutput or List[RolloutOutput] ... """ current_request = infer_request await self.on_trajectory_start([current_request]) current_turn = 1 rollout_infos = {} total_response_ids = [] total_response_loss_mask = [] total_rollout_logprobs = [] while True: messages = current_request.messages if current_turn == 1: # If it's the first turn, remove the response # Keep the original logic, but I think this step is unnecessary here. remove_response(messages) # Get model response response: 'ChatCompletionResponse' = await self.infer_engine.infer_async( current_request, request_config, **kwargs) response_choice: 'ChatCompletionResponseChoice' = response.choices[0] if current_turn > 1 and not messages[-1]['content']: # The dummy assistant message was intentionally kept during `infer_async` # to ensure correct history processing by the template. # It is now removed before appending the new completion. # otherwise, a syntax error would occur when executing messages[-1]['content'] += completion. remove_response(messages) # Update conversation history completion = response_choice.message.content is_continuation = False if messages[-1]['role'] == 'assistant': messages[-1]['content'] += completion is_continuation = True else: messages.append({'role': 'assistant', 'content': completion}) # Check stopping conditions turn_result = await self.on_turn_end(current_request, response_choice, current_turn) if turn_result.get('rollout_infos'): rollout_infos.update(turn_result['rollout_infos']) should_stop = self.check_finished(current_request, response_choice, current_turn) if 'done' in turn_result: should_stop = turn_result['done'] # double-check if user forget to judge the max_turns if self.max_turns: should_stop = should_stop or (current_turn >= self.max_turns) if should_stop: # Collect final turn's data current_logprobs = self._extract_logprobs_from_choice(response_choice) final_token_ids = response_choice.token_ids if is_continuation and total_response_ids: # For continuation, extend the last turn's data total_response_ids[-1].extend(final_token_ids) if total_response_loss_mask: total_response_loss_mask[-1].extend([1] * len(final_token_ids)) if total_rollout_logprobs and current_logprobs: total_rollout_logprobs[-1].extend(current_logprobs) elif not total_response_ids: # First turn stopped immediately - need to initialize with final response data if final_token_ids: total_response_ids = [list(final_token_ids)] total_response_loss_mask = [[1] * len(final_token_ids)] if current_logprobs: total_rollout_logprobs = [current_logprobs] # Validate rollout_logprobs completeness: if logprobs are incomplete (missing for some turns), # clear them to disable rollout importance sampling correction (which requires complete logprobs) # Note: rollout_logprobs should match the number of loss_mask=1 tokens, not total response tokens # because completion_mask in grpo_trainer is based on labels != -100, which corresponds to loss_mask=1 final_rollout_logprobs = total_rollout_logprobs if total_rollout_logprobs: total_logprob_count = sum(len(turn_lps) for turn_lps in total_rollout_logprobs) if total_response_loss_mask: # Check if the number of logprobs matches the number of loss_mask=1 tokens total_loss_mask_1_count = sum(sum(mask) for mask in total_response_loss_mask) if total_loss_mask_1_count != total_logprob_count: # Incomplete logprobs, clear them final_rollout_logprobs = [] else: if total_response_ids: total_response_id_count = sum(len(turn_ids) for turn_ids in total_response_ids) if total_response_id_count != total_logprob_count: final_rollout_logprobs = [] else: final_rollout_logprobs = [] return RolloutOutput( response=response, messages=messages, response_token_ids=total_response_ids, response_loss_mask=total_response_loss_mask, rollout_infos={ **rollout_infos, 'num_turns': current_turn }, rollout_logprobs=final_rollout_logprobs, ) # Prepare next turn ret = self.step(current_request, response_choice, current_turn) current_request: 'RolloutInferRequest' = ret['infer_request'] # Track response tokens and masks return_token_id = False if 'response_token_ids' in ret: if is_continuation and total_response_ids: total_response_ids[-1].extend(ret['response_token_ids']) else: total_response_ids.append(ret['response_token_ids']) return_token_id = True if 'response_loss_mask' in ret: assert return_token_id, 'You must return response_token_ids if you want to return response_loss_mask' assert len(ret['response_loss_mask']) == len(ret['response_token_ids']), \ 'response_loss_mask must have the same length as response_token_ids' if is_continuation and total_response_loss_mask: total_response_loss_mask[-1].extend(ret['response_loss_mask']) else: total_response_loss_mask.append(ret['response_loss_mask']) if 'rollout_infos' in ret: # Always overwrite the rollout info for this step. # If you need to keep all step-wise details, switch to append or merge instead. rollout_infos.update(ret['rollout_infos']) # Track rollout_logprobs for rollout importance sampling correction # Prefer step's returned logprobs (which may be modified/truncated) over raw response_choice logprobs if 'rollout_logprobs' in ret and ret['rollout_logprobs']: current_logprobs = ret['rollout_logprobs'] else: current_logprobs = self._extract_logprobs_from_choice(response_choice) if current_logprobs: if is_continuation and total_rollout_logprobs: total_rollout_logprobs[-1].extend(current_logprobs) else: total_rollout_logprobs.append(current_logprobs) current_turn += 1 def step(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> Dict: """ Handles transition between conversation turns. Args: infer_request: Current inference request response_choice: Response from current turn current_turn: Current turn number Returns: Dict[str, Any]: A dictionary containing inference results with the following structure: - infer_request (required): Main inference request object - response_token_ids (Optional[List[int]]): Token IDs of response for current rollout turn - response_loss_mask (Optional[List[int]]): Loss mask for response tokens (same length as response_token_ids) # noqa - rollout_logprobs (Optional[List[float]]): Log probabilities for response tokens. If not provided, will be extracted from response_choice.logprobs as fallback. Useful when modifying response content (e.g., adding prompts) to avoid logprob misalignment. - rollout_infos (Optional[Dict[str, Any]]): Additional metadata (must be serializable) """ raise NotImplementedError( 'Please implement the `step` method in your MultiTurnScheduler subclass, or override the `run` method.') def check_finished(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> bool: """ Default termination logic for checking if a multi-turn rollout should end. This method is invoked by: - The base class MultiTurnScheduler.run() method, OR - Custom run() methods when explicitly called Note: This is the default implementation that can be overridden by subclasses for custom termination logic. Termination Conditions: 1. When response hits length limit (finish_reason == 'length') 2. When conversation reaches max_turns (if max_turns is set) Args: infer_request: The inference request object response_choice: Contains generation results including finish_reason current_turn: Current conversation turn count Returns: bool: True to terminate conversation, False to continue """ if response_choice.finish_reason == 'length': return True if self.max_turns and current_turn >= self.max_turns: return True return False @staticmethod def _extract_logprobs_from_choice(response_choice: 'ChatCompletionResponseChoice') -> List[float]: """Extract logprobs list from response choice for rollout importance sampling. Args: response_choice: The response choice containing logprobs Returns: List of logprob values, or empty list if not available """ if response_choice.logprobs is None: return [] if 'content' in response_choice.logprobs: return [item['logprob'] for item in response_choice.logprobs['content']] return [] class ThinkingModelTipsScheduler(MultiTurnScheduler): """ Scheduler for multi-turn reasoning with Thinking class models. Key Features: 1. Parses both "think" and "answer" content from each assistant response. 2. For each round, only the "think" content from the last round is retained in the message history. 3. Each round's conversation history is processed independently. 4. Returns a list of RolloutOutput objects, one for each round. 5. Please set `--loss_scale last_round` for training last round response. The scheduler will automatically inject a tip prompt if the answer is incorrect, encouraging the model to recheck its reasoning. # noqa """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) acc_func = kwargs.get('acc_function', None) if acc_func is None: from swift.rewards.orm import MathAccuracy acc_func = MathAccuracy() self.acc_func = acc_func self.tips_prompt = 'The answer is not correct, It seems You made a mistake, you need to recheck very carefully.' async def run(self, infer_request: 'RolloutInferRequest', request_config: 'RequestConfig', **kwargs) -> List['RolloutOutput']: """ Execute multi-turn inference for Thinking models. Args: infer_request (RolloutInferRequest): The initial inference request containing the conversation history. request_config (RequestConfig): Configuration for the inference request. **kwargs: Additional arguments for the inference engine. Returns: List[RolloutOutput]: A list of RolloutOutput objects, one for each reasoning round. """ current_request = infer_request current_turn = 1 rollout_outputs = [] while True: messages = current_request.messages # Obtain model response for the current turn response: 'ChatCompletionResponse' = await self.infer_engine.infer_async( current_request, request_config, **kwargs) response_choice: 'ChatCompletionResponseChoice' = response.choices[0] completion = response_choice.message.content # Append the assistant's response to the message history messages.append({'role': 'assistant', 'content': completion}) # Construct the message history for this round, keeping only the last "think" content messages_with_last_think = self._build_messages(messages) # Create a RolloutOutput for the current round round_output = RolloutOutput( response=response, messages=messages_with_last_think, response_token_ids=response_choice.token_ids, rollout_infos={'num_turns': current_turn}) # Store the output for this round rollout_outputs.append(round_output) # Determine whether to stop the multi-turn reasoning should_stop = self.check_finished(current_request, response_choice, current_turn) if should_stop: break # Prepare for the next turn by updating the inference request ret = self.step(current_request, response_choice, current_turn) current_request: 'RolloutInferRequest' = ret['infer_request'] current_turn += 1 return rollout_outputs def check_finished(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> bool: last_query = infer_request.messages[-2]['content'] # tips once if self.tips_prompt in last_query: return True completion = response_choice.message.content solution = infer_request.data_dict['solution'] acc = self.acc_func([completion], [solution])[0] if acc == 1: return True return super().check_finished(infer_request, response_choice, current_turn) def step(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> Dict: infer_request.messages.append({'role': 'user', 'content': self.tips_prompt}) return {'infer_request': infer_request} def _is_thinking_template(self) -> bool: if not hasattr(self.infer_engine, 'template'): return False template = self.infer_engine.template return template.template_meta.is_thinking def _build_messages(self, original_messages: Messages) -> Messages: """ Build history for a specific round, keeping only the think content from the last round. Args: original_messages: Original conversation messages Returns: Messages: History for this specific round """ from copy import deepcopy # If this is a thinking template, use the template's method to prepare messages if self._is_thinking_template(): # Create a mock inputs object to use the template's _swift_prepare_inputs method class MockInputs: def __init__(self, messages): self.messages = deepcopy(messages) mock_inputs = MockInputs(original_messages) # Set up the template for inference mode template = self.infer_engine.template # _swift_prepare_inputs will remove historical thinking content when in train mode, patch the mode here original_mode = template.mode template.mode = 'train' # Use the template's method to prepare messages template._swift_prepare_inputs(mock_inputs) # Restore original mode template.mode = original_mode return mock_inputs.messages else: # Fallback to manual processing for non-thinking templates round_messages = [] # Process messages in original order for i, msg in enumerate(original_messages): if msg['role'] == 'assistant' and isinstance(msg['content'], str) and i != len(original_messages) - 1: # For assistant messages assistant_no_think = msg['content'].split('')[-1].strip() round_messages.append(assistant_no_think) else: round_messages.append(deepcopy(msg)) return round_messages class MathTipsScheduler(MultiTurnScheduler): tips_prompt = 'But wait... It seems I made a mistake,' def __init__(self, *args, **kwargs): from swift.rewards.orm import MathAccuracy super().__init__(*args, **kwargs) self.acc_func = kwargs.get('acc_function', MathAccuracy()) # Cache the tokenized tips_prompt length for loss mask computation self._tips_token_ids = None def _get_tips_token_ids(self, tokenizer) -> List[int]: """Get tokenized tips_prompt (cached for efficiency).""" if self._tips_token_ids is None: # Tokenize without special tokens to get the raw token ids self._tips_token_ids = tokenizer.encode(self.tips_prompt, add_special_tokens=False) return self._tips_token_ids def check_finished(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> bool: last_completion = infer_request.messages[-1]['content'] # we only give tips once if self.tips_prompt in last_completion: return True solution = infer_request.data_dict['solution'] acc = self.acc_func([last_completion], [solution])[0] if acc == 1: return True return super().check_finished(infer_request, response_choice, current_turn) def step(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> Dict: completion = response_choice.message.content response_token_ids = list(response_choice.token_ids) if response_choice.token_ids else [] # Extract logprobs from response_choice before any truncation rollout_logprobs = self._extract_logprobs_from_choice(response_choice) # Truncate completion at or tags truncate_idx = len(completion) if '' in completion: truncate_idx = min(truncate_idx, completion.index('')) if '' in completion: truncate_idx = min(truncate_idx, completion.index('')) if truncate_idx < len(completion): # Need to truncate token_ids and logprobs as well truncated_completion = completion[:truncate_idx] if response_token_ids and self.tokenizer is not None: # Find the token index corresponding to the truncation point # by decoding progressively until we reach or exceed the truncation point token_truncate_idx = len(response_token_ids) for i in range(1, len(response_token_ids) + 1): decoded = self.tokenizer.decode(response_token_ids[:i], skip_special_tokens=False) if len(decoded) >= truncate_idx: token_truncate_idx = i break response_token_ids = response_token_ids[:token_truncate_idx] # Truncate logprobs to match if rollout_logprobs: rollout_logprobs = rollout_logprobs[:token_truncate_idx] completion = truncated_completion # Add tips_prompt completion += self.tips_prompt # Compute loss_mask for tips tokens # Note: rollout_logprobs should NOT include tips tokens because: # 1. Tips tokens have loss_mask=0, so their labels will be -100 # 2. completion_mask = (labels != -100), so tips tokens won't be in completion_mask # 3. rollout_logprobs must align with completion_mask, not response_token_ids if response_token_ids and self.tokenizer is not None: tips_token_ids = self._get_tips_token_ids(self.tokenizer) # Loss mask: original tokens = 1, tips tokens = 0 response_loss_mask = [1] * len(response_token_ids) + [0] * len(tips_token_ids) # Append tips token ids to response response_token_ids = response_token_ids + tips_token_ids # Do NOT extend rollout_logprobs for tips tokens - they are masked out in completion_mask else: response_loss_mask = [] # Update messages if infer_request.messages[-1]['role'] == 'assistant': if not infer_request.messages[-1]['content']: # Multi-turn continuation: pop the dummy input we add in last turn infer_request.messages.pop(-1) infer_request.messages[-1]['content'] = completion else: infer_request.messages.append({'role': 'assistant', 'content': completion}) result = {'infer_request': infer_request} if response_token_ids: result['response_token_ids'] = response_token_ids result['response_loss_mask'] = response_loss_mask if rollout_logprobs: result['rollout_logprobs'] = rollout_logprobs return result class GYMScheduler(MultiTurnScheduler): """Gym environment-driven scheduler using universal hooks. Implements ``on_trajectory_start`` (env.reset) and ``on_turn_end`` (env.step) to integrate gym environments into the multi-turn protocol. Works in both server mode (``run()``) and colocate mode (``run_multi_turn()``). """ def __init__(self, infer_engine: Optional['GRPOVllmEngine'] = None, max_turns: Optional[int] = None, **kwargs): super().__init__(infer_engine, max_turns, **kwargs) self.gym_env_name = kwargs.get('gym_env', None) # Per-trajectory state (keyed by uuid) self._envs: Dict[str, Env] = {} self._total_rewards: Dict[str, float] = {} self._step_rewards: Dict[str, List[float]] = {} self._pending_obs: Dict[str, Optional[str]] = {} async def _close_and_remove(self, uuid: str) -> None: """Close env for a given uuid and remove all associated state.""" env = self._envs.pop(uuid, None) if env is not None: try: await env.close() except Exception: pass self._total_rewards.pop(uuid, None) self._step_rewards.pop(uuid, None) self._pending_obs.pop(uuid, None) # ------------------------------------------------------------------ # Universal async hooks (called by both run() and run_multi_turn()) # ------------------------------------------------------------------ async def on_trajectory_start(self, requests: List['RolloutInferRequest']) -> None: """Create one env per request and seed messages with initial observation.""" async def _init_single(req: 'RolloutInferRequest') -> None: uuid = req.uuid if uuid in self._envs: await self._close_and_remove(uuid) env_config = (req.data_dict or {}).get('env_config', {}) if hasattr(req, 'data_dict') else {} env = self._create_env(env_config) observation, info, system_message = await env.reset(req) messages: Messages = [] if system_message: messages.append({'role': 'system', 'content': system_message}) messages.append({'role': 'user', 'content': observation}) req.messages = messages self._envs[uuid] = env self._total_rewards[uuid] = 0.0 self._step_rewards[uuid] = [] self._pending_obs[uuid] = None await asyncio.gather(*[_init_single(req) for req in requests]) async def on_turn_end(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> Dict[str, Any]: """Advance the gym env, accumulate reward, and return done + rollout_infos.""" uuid = infer_request.uuid env = self._envs.get(uuid) if env is None: return {'done': True, 'rollout_infos': {}} next_obs, reward, done, info = await env.step(deepcopy(infer_request.messages)) self._total_rewards[uuid] = self._total_rewards.get(uuid, 0.0) + float(reward) self._step_rewards.setdefault(uuid, []).append(float(reward)) self._pending_obs[uuid] = None if done else next_obs rollout_infos: Dict[str, Any] = { 'total_reward': self._total_rewards[uuid], 'step_rewards': list(self._step_rewards.get(uuid, [])), 'gym_done': done, } if done: await self._close_and_remove(uuid) return {'done': done, 'rollout_infos': rollout_infos} # ------------------------------------------------------------------ # Step hook (injects next observation for the next turn) # ------------------------------------------------------------------ def step(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> Dict[str, Any]: uuid = infer_request.uuid next_obs = self._pending_obs.get(uuid) if next_obs is not None: infer_request.messages.append({'role': 'user', 'content': next_obs}) self._pending_obs[uuid] = None return {'infer_request': infer_request} # ------------------------------------------------------------------ # Env helpers # ------------------------------------------------------------------ def _create_env(self, env_config: Dict) -> Env: env_name = env_config.get('name', self.gym_env_name) if env_name not in envs: raise ValueError(f"Environment '{env_name}' not found. Available: {list(envs.keys())}") return envs[env_name](env_config) class OpenEnvScheduler(GYMScheduler): """GYMScheduler specialised for OpenEnv environments. Unlike GYMScheduler which uses async ``Env`` instances, OpenEnvScheduler uses :class:`OpenEnvWrapper` whose ``reset()`` / ``step()`` / ``close()`` are **synchronous** (blocking WebSocket I/O). Subclasses that override ``on_trajectory_start`` / ``on_turn_end`` should wrap sync wrapper calls with ``asyncio.to_thread()`` to avoid blocking the event loop. Action parsing (LLM text → dict) and observation formatting (dict → str) are handled by overridable :meth:`parse_action` and :meth:`format_observation` methods, eliminating the need for ``openenv_*`` command-line parameters. All OpenEnv configuration (``base_url``, ``system_message``, ``reset_kwargs`` …) comes from the dataset's per-row ``env_config``. """ def _create_env(self, env_config: Dict) -> Any: """Create an :class:`OpenEnvWrapper` (not an ``Env`` subclass).""" from .openenv_wrapper import OpenEnvWrapper return OpenEnvWrapper(env_config) async def _close_and_remove(self, uuid: str) -> None: """Close wrapper for a given uuid and remove all associated state. Wrapper.close() is synchronous; use ``asyncio.to_thread`` to avoid blocking the event loop. """ import asyncio wrapper = self._envs.pop(uuid, None) if wrapper is not None: try: await asyncio.to_thread(wrapper.close) except Exception: pass self._total_rewards.pop(uuid, None) self._step_rewards.pop(uuid, None) self._pending_obs.pop(uuid, None) async def on_trajectory_start(self, requests: List['RolloutInferRequest']) -> None: """Create one wrapper per request, call ``reset()``, seed messages. Uses a semaphore to limit concurrent environment creations (default 4) to avoid overwhelming the OpenEnv server with simultaneous WebSocket connections. """ semaphore = asyncio.Semaphore(getattr(self, 'max_concurrent_envs', 4)) async def _init_single(req: 'RolloutInferRequest') -> None: async with semaphore: uuid = req.uuid if uuid in self._envs: await self._close_and_remove(uuid) row_env_config = (req.data_dict or {}).get('env_config', {}) if hasattr(req, 'data_dict') else {} env_config = {**getattr(self, 'env_config_defaults', {}), **row_env_config} wrapper = self._create_env(env_config) obs, metadata = wrapper.reset() system_message = env_config.get('system_message', '') messages: Messages = [] if system_message: messages.append({'role': 'system', 'content': system_message}) messages.append({'role': 'user', 'content': self.format_observation(obs)}) req.messages = messages self._envs[uuid] = wrapper self._total_rewards[uuid] = 0.0 self._step_rewards[uuid] = [] self._pending_obs[uuid] = None await asyncio.gather(*[_init_single(req) for req in requests]) async def on_turn_end(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', current_turn: int) -> Dict[str, Any]: """Parse LLM response, call ``wrapper.step()``, accumulate reward.""" uuid = infer_request.uuid wrapper = self._envs.get(uuid) if wrapper is None: return {'done': True, 'rollout_infos': {}} action_text = response_choice.message.content action_dict = self.parse_action(action_text) obs, reward, done, metadata = wrapper.step(action_dict) self._total_rewards[uuid] = self._total_rewards.get(uuid, 0.0) + float(reward) self._step_rewards.setdefault(uuid, []).append(float(reward)) next_obs = None if done else self.format_observation(obs) self._pending_obs[uuid] = next_obs rollout_infos: Dict[str, Any] = { 'total_reward': self._total_rewards[uuid], 'step_rewards': list(self._step_rewards.get(uuid, [])), 'gym_done': done, } if done: await self._close_and_remove(uuid) return {'done': done, 'rollout_infos': rollout_infos} def parse_action(self, text: str) -> Dict[str, Any]: """Parse LLM response text into an OpenEnv action dict. Default: try ``json.loads``, fall back to ``{"message": text}``. """ text = text.strip() # Strip markdown code blocks (e.g. ```json ... ```) if text.startswith('```'): lines = text.splitlines() if len(lines) >= 2 and lines[0].startswith('```') and lines[-1].strip().startswith('```'): text = '\n'.join(lines[1:-1]).strip() try: parsed = json.loads(text) if isinstance(parsed, dict): return parsed return {'message': str(parsed)} except (json.JSONDecodeError, ValueError): return {'message': text} def format_observation(self, observation: Any) -> str: """Format OpenEnv observation into a string for the LLM. Default: ``json.dumps``. Override in subclasses for environment-specific formatting (e.g. extract a ``"question"`` field). """ try: return json.dumps(observation, ensure_ascii=False, default=str) except (TypeError, ValueError): return str(observation) multi_turns = { 'math_tip_trick': MathTipsScheduler, 'gym_scheduler': GYMScheduler, 'openenv_scheduler': OpenEnvScheduler, 'thinking_tips_scheduler': ThinkingModelTipsScheduler, }