Add process advantage weighting for agent rollouts

Author: Harold-lkk

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

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__}()"

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}"

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

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",

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,13 +205,17 @@ 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.")
         super().__init__(rollout_ctl, sample_params, hf_checkpoint, judger, logger)
         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:

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"]

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

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]: