From 0767d9e073f238809bec55ceeaa2a96f33218893 Mon Sep 17 00:00:00 2001 From: liukuikun <641417025@qq.com> Date: Tue, 30 Jun 2026 06:43:33 +0000 Subject: [PATCH] Add process advantage weighting for agent rollouts --- xtuner/v1/data_proto/rl_data.py | 5 ++ xtuner/v1/rl/advantage/base.py | 18 +++++ .../agent_in_localhost_loop.py | 17 ++++ .../localhost_agent_loop/compose.py | 3 +- .../agent_loop/sandbox_agent_loop/__init__.py | 2 + .../agent_in_sandbox_loop.py | 17 ++++ .../agent_loop/sandbox_agent_loop/compose.py | 78 +++++++++++++++++++ xtuner/v1/rl/rollout/chat_template.py | 3 + xtuner/v1/rl/rollout/trace_store.py | 21 ++++- xtuner/v1/train/rl_trainer.py | 52 ++++++++++++- 10 files changed, 212 insertions(+), 4 deletions(-) create mode 100644 xtuner/v1/rl/agent_loop/sandbox_agent_loop/compose.py diff --git a/xtuner/v1/data_proto/rl_data.py b/xtuner/v1/data_proto/rl_data.py index 7db20f2940..52317e1acb 100644 --- a/xtuner/v1/data_proto/rl_data.py +++ b/xtuner/v1/data_proto/rl_data.py @@ -120,6 +120,10 @@ class RolloutState(BaseModel): input_ids: list[int] | None = None labels: list[int] | None = None + # Per-token multiplier applied to positive advantages after outcome reward + # advantage estimation. Coordinates match input_ids / labels; trainer uses + # advantage_weight[1:] to align with shifted_labels. + advantage_weight: list[float] | None = None # --- Judger 输出 --- reward: dict[str, Any] | None = None @@ -248,6 +252,7 @@ def reset_rollout_response(rollout_state: RolloutState) -> RolloutState: rollout_state.finish_reason = None rollout_state.response_mask = [] rollout_state.response_model_steps = [] + rollout_state.advantage_weight = None rollout_state.reward = None rollout_state.error_msg = None return rollout_state diff --git a/xtuner/v1/rl/advantage/base.py b/xtuner/v1/rl/advantage/base.py index 44d279abbe..79693aa2f6 100644 --- a/xtuner/v1/rl/advantage/base.py +++ b/xtuner/v1/rl/advantage/base.py @@ -58,5 +58,23 @@ def compute(self, rewards: torch.Tensor, group: list[Any]) -> torch.Tensor: """ ... + def expand_to_token_advantages( + self, + *, + base_advantage: float, + rollout_state: Any, + shifted_labels: list[int], + shifted_advantage_weight: list[float] | None = None, + ) -> tuple[list[float], dict[str, Any]]: + """Expand a sample-level advantage to token-level advantages. + + ``compute`` intentionally stays sample/session-level. This hook lets + downstream projects shape token credit after labels and optional + per-token weights are known by the trainer. + """ + + del rollout_state, shifted_advantage_weight + return [0.0 if label == -100 else base_advantage for label in shifted_labels], {} + def __repr__(self) -> str: return f"{self.__class__.__name__}()" diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py index b70edb888a..86334f406a 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py @@ -85,6 +85,7 @@ class AgentInLocalhostLoopConfig(AgentLoopConfig): sample_timeout_s: float | None = None mode: Literal["train", "eval"] = "train" requires_rollout_proxy: bool = True + process_advantage_builder: str | None = None def build_local( self, @@ -101,6 +102,7 @@ def build_local( max_concurrent_samples=self.max_concurrent_samples, sample_timeout_s=self.sample_timeout_s, mode=self.mode, + process_advantage_builder=self.process_advantage_builder, ) @@ -117,6 +119,7 @@ def __init__( max_concurrent_samples: int | None = None, sample_timeout_s: float | None = None, mode: Literal["train", "eval"] = "train", + process_advantage_builder: str | None = None, ): if hf_checkpoint is None: raise ValueError("hf_checkpoint must be provided for AgentInLocalhostLoop.") @@ -125,6 +128,9 @@ def __init__( self.sample_timeout_s = sample_timeout_s self._sample_semaphore = asyncio.Semaphore(max_concurrent_samples) if max_concurrent_samples else None self.mode = mode + self.process_advantage_builder = ( + _import_from_path(process_advantage_builder) if process_advantage_builder is not None else None + ) async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: async def generate_one(state: RolloutState) -> RolloutState: @@ -246,6 +252,16 @@ async def _fill_rollout_state(self, rollout_state: RolloutState, item: AgentRoll rollout_state.input_ids = data["input_ids"] rollout_state.labels = data["labels"] + rollout_state.extra_fields["agent_trace_segments"] = data.get("segments", []) + if self.process_advantage_builder is not None: + rollout_state.advantage_weight, process_adv_summary = self.process_advantage_builder( + segment["messages"], + data["labels"], + data.get("segments"), + ) + rollout_state.extra_fields["process_adv"] = process_adv_summary + else: + rollout_state.advantage_weight = None rollout_state.response_ids = [ token_id for token_id, label in zip(data["input_ids"][1:], data["labels"][1:]) if label != -100 ] @@ -267,6 +283,7 @@ def _fill_eval_rollout_state(self, rollout_state: RolloutState, item: AgentRollo rollout_state.routed_experts = None rollout_state.response_mask = None rollout_state.response_model_steps = None + rollout_state.advantage_weight = None rollout_state.extra_fields["agent_status"] = item.status.value if item.error is not None: rollout_state.error_msg = f"{item.error.stage}/{item.error.category}: {item.error.message}" diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/compose.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/compose.py index c09d30590f..0b223485f9 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/compose.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/compose.py @@ -35,12 +35,13 @@ def __init__( async def run(self, item: AgentRolloutItem, record: StageRecord) -> float: record.status = StageStatus.RUNNING record.started_at = record.started_at or time.monotonic() + record.judger_name = self.name try: weighted_score = 0.0 total_weight = 0.0 for stage in self.stages: name = getattr(stage, "name", stage.__class__.__name__) - child_record = item.judgers.setdefault(name, StageRecord()) + child_record = item.judgers.setdefault(name, StageRecord(judger_name=name)) score = float(await stage.run(item, child_record)) stage_weight = max(float(getattr(stage, "weight", 1.0)), 0.0) weighted_score += score * stage_weight diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/__init__.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/__init__.py index 00b5ef1f65..5b250df865 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/__init__.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/__init__.py @@ -11,6 +11,7 @@ AgentInSandboxLoop, AgentInSandboxLoopConfig, ) +from xtuner.v1.rl.agent_loop.sandbox_agent_loop.compose import SandboxComposeStage from xtuner.v1.rl.agent_loop.sandbox_agent_loop.hooks import ( DownloadHook, ExecHook, @@ -71,6 +72,7 @@ "RunAgentInstallDeps", "Runner", "SandboxPool", + "SandboxComposeStage", "SandboxSpec", "SandboxStage", "ShellEntry", diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py index fcb2294221..bb6cd6e735 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/agent_in_sandbox_loop.py @@ -178,6 +178,7 @@ class AgentInSandboxLoopConfig(AgentLoopConfig): max_concurrent_samples: int | None = None mode: Literal["train", "eval"] = "train" requires_rollout_proxy: bool = True + process_advantage_builder: str | None = None def build_local( self, rollout_controller: RolloutController | None = None, judger: Judger | None = None, logger=None @@ -190,6 +191,7 @@ def build_local( logger=logger, max_concurrent_samples=self.max_concurrent_samples, mode=self.mode, + process_advantage_builder=self.process_advantage_builder, ) @@ -203,6 +205,7 @@ def __init__( logger=None, max_concurrent_samples: int | None = None, mode: Literal["train", "eval"] = "train", + process_advantage_builder: str | None = None, ): if hf_checkpoint is None: raise ValueError("hf_checkpoint must be provided for AgentInSandboxLoop.") @@ -210,6 +213,9 @@ def __init__( self.max_concurrent_samples = max_concurrent_samples self._sample_semaphore = asyncio.Semaphore(max_concurrent_samples) if max_concurrent_samples else None self.mode = mode + self.process_advantage_builder = ( + _import_from_path(process_advantage_builder) if process_advantage_builder is not None else None + ) async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: async def generate_one(state: RolloutState) -> list[RolloutState]: @@ -313,6 +319,16 @@ async def _build_rollout_states(self, rollout_state: RolloutState, item: AgentRo data = await trace_store.export_training_trace.remote(str(rollout_state.session_id), prompt_text) segment_state.input_ids = data["input_ids"] segment_state.labels = data["labels"] + segment_state.extra_fields["agent_trace_segments"] = data.get("segments", []) + if self.process_advantage_builder is not None: + segment_state.advantage_weight, process_adv_summary = self.process_advantage_builder( + messages, + data["labels"], + data.get("segments"), + ) + segment_state.extra_fields["process_adv"] = process_adv_summary + else: + segment_state.advantage_weight = None # Agentic training consumes input_ids/labels directly. response_ids is # filled here only so rollout throughput logging can print rollout_tgs. segment_state.response_ids = [ @@ -341,6 +357,7 @@ def _fill_eval_rollout_state(self, rollout_state: RolloutState, item: AgentRollo rollout_state.routed_experts = None rollout_state.response_mask = None rollout_state.response_model_steps = None + rollout_state.advantage_weight = None rollout_state.extra_fields["agent_status"] = item.status.value selected_agent = _selected_agent(item) if selected_agent is not None: diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/compose.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/compose.py new file mode 100644 index 0000000000..2cf0778895 --- /dev/null +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/compose.py @@ -0,0 +1,78 @@ +"""Composable sandbox validation stages.""" + +from __future__ import annotations + +import time +from typing import Any + +from lagent.utils import create_object + +from xtuner.v1.rl.agent_loop.sandbox_agent_loop.sandbox import SandboxPool +from xtuner.v1.rl.agent_loop.sandbox_agent_loop.schemas import ( + AgentRolloutItem, + RolloutError, + StageRecord, + StageStatus, +) + + +class SandboxComposeStage: + """Compose multiple sandbox validation stages behind ``run(...) -> float``. + + Stages with ``weight=0`` still run, but do not contribute to the returned + score. This is used for process-adv annotators that mutate rollout + artifacts without changing outcome reward. + """ + + def __init__( + self, + stages: list[Any], + *, + name: str = "validate", + weight: float = 1.0, + ): + if not stages: + raise ValueError("SandboxComposeStage.stages is empty") + self.name = name + self.stages = [create_object(stage) for stage in stages] + self.weight = weight + + async def run(self, item: AgentRolloutItem, pool: SandboxPool, record: StageRecord) -> float: + record.status = StageStatus.RUNNING + record.started_at = record.started_at or time.monotonic() + record.judger_name = self.name + try: + weighted_score = 0.0 + total_weight = 0.0 + for stage in self.stages: + name = getattr(stage, "name", stage.__class__.__name__) + child_record = item.judgers.setdefault(name, StageRecord(judger_name=name)) + score = float(await stage.run(item, pool, child_record)) + stage_weight = max(float(getattr(stage, "weight", 1.0)), 0.0) + weighted_score += score * stage_weight + total_weight += stage_weight + record.score = weighted_score / total_weight if total_weight > 0 else 0.0 + record.status = StageStatus.COMPLETED + return record.score + except Exception as exc: + record.status = StageStatus.FAILED + child_error = next( + (child.error for child in item.judgers.values() if child.error is not None), + None, + ) + record.error = ( + record.error + or child_error + or RolloutError( + stage=self.name, + category="validate_failed", + type=type(exc).__name__, + message=str(exc), + ) + ) + raise + finally: + record.finished_at = time.monotonic() + + +__all__ = ["SandboxComposeStage"] diff --git a/xtuner/v1/rl/rollout/chat_template.py b/xtuner/v1/rl/rollout/chat_template.py index afe5c1915c..ed7282684a 100644 --- a/xtuner/v1/rl/rollout/chat_template.py +++ b/xtuner/v1/rl/rollout/chat_template.py @@ -4,6 +4,7 @@ _RAW_ARGUMENTS_KEY = "__xtuner_raw_arguments__" +_PROCESS_ONLY_MESSAGE_KEYS = ("finish_reason", "metainfo") def canonicalize_messages_for_chat_template(messages: list[dict]) -> list[dict]: @@ -19,6 +20,8 @@ def canonicalize_messages_for_chat_template(messages: list[dict]) -> list[dict]: messages = copy.deepcopy(messages) for message in messages: + for key in _PROCESS_ONLY_MESSAGE_KEYS: + message.pop(key, None) tool_calls = message.get("tool_calls") if not isinstance(tool_calls, list): continue diff --git a/xtuner/v1/rl/rollout/trace_store.py b/xtuner/v1/rl/rollout/trace_store.py index 719590bd80..20336d0819 100644 --- a/xtuner/v1/rl/rollout/trace_store.py +++ b/xtuner/v1/rl/rollout/trace_store.py @@ -323,7 +323,7 @@ def export_training_trace(self, session_id: str, prompt_text: str) -> dict: Returns: dict: The trace dictionary containing `input_ids`, `labels`, `logprobs`, - and `routed_experts`. + `routed_experts`, and per-segment token spans. Raises: ValueError: If the prompt_text does not completely match the trace keys in the session. @@ -353,17 +353,34 @@ def export_training_trace(self, session_id: str, prompt_text: str) -> dict: f"prompt_len={len(prompt_text)} matched_len={len(key)} key_count={len(session_keys)}. " "See the logged '[TraceStore] prompt mismatch' report for the full diff." ) - trace: dict[str, list[Any]] = {"input_ids": [], "labels": [], "logprobs": [], "routed_experts": []} + trace: dict[str, list[Any]] = { + "input_ids": [], + "labels": [], + "logprobs": [], + "routed_experts": [], + "segments": [], + } for node in nodes: node_val = node.value if not isinstance(node_val, TokenizedSegment): raise TypeError(f"Unexpected trace node value type: {type(node_val)!r}") assert node_val.labels is not None assert node_val.logprobs is not None + start = len(trace["input_ids"]) + end = start + len(node_val.token_ids) + trainable = any(label != -100 for label in node_val.labels) trace["input_ids"].extend(node_val.token_ids) trace["labels"].extend(node_val.labels) trace["logprobs"].extend(node_val.logprobs) trace["routed_experts"].append(node_val.expert_key) + trace["segments"].append( + { + "start": start, + "end": end, + "trainable": trainable, + "kind": "assistant_response" if trainable else "context_delta", + } + ) return trace def get_objects(self, keys: list[str]) -> list[ray.ObjectRef]: diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index eecb796ddd..e124c21441 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -1010,6 +1010,11 @@ def _prepare_train_data( response_len_list = [] tool_turns_list: list[int] = [] training_tokens = 0 + process_adv_affected_samples = 0 + process_adv_affected_tokens = 0 + process_adv_affected_positive_tokens = 0 + process_adv_weight_sum = 0.0 + process_adv_alignment_mismatch = 0 data_batches = [] @@ -1092,7 +1097,45 @@ def _prepare_train_data( response_len_list.append(response_len) advatnages_val = sample_advantages[i] - actual_advantages = [0.0 if label == -100 else advatnages_val for label in shifted_labels] + advantage_weight = group[i].advantage_weight + if advantage_weight is None: + shifted_advantage_weight = [1.0] * len(shifted_labels) + else: + assert len(advantage_weight) == len(raw_input_ids), ( + f"{len(advantage_weight)} vs {len(raw_input_ids)}, data: {group[i]}" + ) + shifted_advantage_weight = [float(weight) for weight in advantage_weight[1:]] + + process_adv_summary = group[i].extra_fields.get("process_adv") + if isinstance(process_adv_summary, dict) and process_adv_summary.get("alignment_mismatch"): + process_adv_alignment_mismatch += 1 + + expand_to_token_advantages = getattr( + self._advantage_estimator, + "expand_to_token_advantages", + None, + ) + if expand_to_token_advantages is None: + actual_advantages = [0.0 if label == -100 else advatnages_val for label in shifted_labels] + token_advantage_info = {} + else: + actual_advantages, token_advantage_info = expand_to_token_advantages( + base_advantage=advatnages_val, + rollout_state=group[i], + shifted_labels=shifted_labels, + shifted_advantage_weight=shifted_advantage_weight, + ) + assert len(actual_advantages) == len(shifted_labels), ( + f"{len(actual_advantages)} vs {len(shifted_labels)}, data: {group[i]}" + ) + affected_tokens = int(token_advantage_info.get("affected_tokens", 0) or 0) + if affected_tokens: + process_adv_affected_samples += int(token_advantage_info.get("affected_samples", 1) or 1) + process_adv_affected_tokens += affected_tokens + process_adv_weight_sum += float(token_advantage_info.get("weight_sum", 0.0) or 0.0) + process_adv_affected_positive_tokens += int( + token_advantage_info.get("affected_positive_adv_tokens", 0) or 0 + ) advantages_list.extend(actual_advantages) assert len(input_ids) <= pack_max_length, f"{len(input_ids)} vs {pack_max_length}" @@ -1230,6 +1273,13 @@ def _prepare_train_data( "prompt_len/min": prompt_len_t.min().item(), "prompt_len/max": prompt_len_t.max().item(), } + if process_adv_affected_tokens or process_adv_alignment_mismatch: + info_dict["process_adv/affected_samples"] = process_adv_affected_samples + info_dict["process_adv/affected_tokens"] = process_adv_affected_tokens + info_dict["process_adv/affected_positive_adv_tokens"] = process_adv_affected_positive_tokens + info_dict["process_adv/alignment_mismatch"] = process_adv_alignment_mismatch + if process_adv_affected_tokens: + info_dict["process_adv/weight_mean"] = process_adv_weight_sum / process_adv_affected_tokens if tool_turns_list: tool_turns_t = torch.tensor(tool_turns_list, dtype=torch.float32) info_dict["tool_turns/mean"] = tool_turns_t.mean().item()