[Refactor] Move train data preparation into AgentLoop and TrainingController

docs/zh_cn/rl/advanced_tutorial/agent_loop.md

@@ -84,16 +84,17 @@ AgentLoop 输入和输出都是 `RolloutState`。如果后续使用预置 `RLTra
 AgentLoop 返回的 `RolloutState` 如果要进入训练，至少需要满足：
 
 - `status == Status.COMPLETED`。预置 trainer 会跳过 `ABORTED`、`FILTERED`、`FAILED` 的样本组。
-- `response_ids` 非空。`_prepare_train_data()` 用它构造训练 token。
+- `response_ids` 非空。`canonicalize_train_fields()` 用它构造训练 token。
 - `response` 非空。Judger 和轨迹保存依赖文本 response。
-- `reward["score"]` 存在。`_prepare_train_data()` 会直接读取它计算 advantage。
+- `reward["score"]` 存在。训练 controller 会直接读取它计算 advantage。
 
 如果提供以下字段，还需要满足长度约定：
 
 - `logprobs`：长度必须等于 `len(response_ids)`。
-- `response_mask`：长度必须等于 `len(response_ids)`。mask 为 `0` 的 token 会被转成训练 label `-100`，对应 advantage 也会置为 `0.0`。
 
-这也是自定义 AgentLoop 最容易出错的地方：工具返回、环境反馈、系统插入内容等不是模型生成的 token，应该在 `response_mask` 中置为 `0`，并给对应 `logprobs` 填 `0.0`。
+labels 是唯一的监督载体。`AgentLoop.generate_group()` 会在末尾调用 `canonicalize_train_fields()`，为 prompt+response 型样本构造全序列 `input_ids`/`labels`/`logprobs`（三者等长、未 shift）。这也是自定义 AgentLoop 最容易出错的地方：工具返回、环境反馈、系统插入内容等不是模型生成的 token，不参与训练——对应 label 直接写 `-100`，`logprobs` 填 `0.0`。prompt+response 型样本不需要手动构造这些字段（基类会兜底）；自行组装全序列的多轮 loop 则必须自己把语义洞烙进 labels。训练侧只做 shift 和 advantage 计算，不处理任何掩码。
+
+`response_ids` 统一约定为 `input_ids` 去掉 prompt 前缀的连续后缀（`len(response_ids) == len(input_ids) - len(prompt_ids)`，环境插入或工具返回的 token 也包含在内）；`response_model_steps` 与之等长，记录每位 token 的来源模型版本。token 级 staleness 据此对 prompt+response 与 agentic 全序列样本统一生效：逐 token 判定新鲜度，过期 token 的 label 烙成 `-100`，语义洞位（label 为 `-100`）不受影响。
 
 ## SingleTurnAgentLoop
 
@@ -147,7 +148,7 @@ agent_loop_config = SingleTurnAgentLoopConfig(
 自定义 AgentLoop 通常需要做四件事：
 
 1. 继承 `AgentLoop`，实现 `generate_sample()`。
-2. 在 `generate_sample()` 中维护 `tokens`、`sample_params`、`response_ids`、`response`、`logprobs`、`response_mask`、`status`，不要在这里调用外部 Judger。
+2. 在 `generate_sample()` 中维护 `tokens`、`sample_params`、`response_ids`、`response`、`logprobs`、`status`，不要在这里调用外部 Judger。prompt+response 型样本的训练字段（`input_ids`/`labels`）由基类 `canonicalize_train_fields()` 兜底构造；自行组装全序列的 loop 需自己写 `input_ids`/`labels`，并把语义洞直接烙进 labels，同时把 `response_ids` 写成 `input_ids` 去掉 prompt 前缀的连续后缀（含环境 token）。
 3. 若覆盖 `generate_group()`，在其中显式编排 Teacher、Judger 和组级过滤；没有 validity check 时应让完成的样本尽早进入 Teacher，有 validity check 时应只把过滤通过的完整组发送给 Teacher。
 4. 继承 `AgentLoopConfig`，实现 `build_local()`，这样才能接入 `TaskSpecConfig.agent_loop_config`，并复用 Ray actor 构建逻辑。
 
@@ -200,7 +201,7 @@ class ToolAgentLoop(AgentLoop):
     async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState:
         final_response_ids: list[int] = []
         final_logprobs: list[float] = []
-        final_response_mask: list[int] = []
+        final_supervised: list[int] = []
 
         cur_tokens = list(rollout_state.tokens or rollout_state.prompt_ids or [])
         remaining_tokens = self.sample_params.max_tokens
@@ -221,7 +222,7 @@ class ToolAgentLoop(AgentLoop):
 
             final_response_ids.extend(response_ids)
             final_logprobs.extend(logprobs)
-            final_response_mask.extend([1] * len(response_ids))
+            final_supervised.extend([1] * len(response_ids))
             cur_tokens.extend(response_ids)
 
             tool_tokens = self._run_tool_and_encode_result(rollout_state)
@@ -230,7 +231,7 @@ class ToolAgentLoop(AgentLoop):
 
             final_response_ids.extend(tool_tokens)
             final_logprobs.extend([0.0] * len(tool_tokens))
-            final_response_mask.extend([0] * len(tool_tokens))
+            final_supervised.extend([0] * len(tool_tokens))
             cur_tokens.extend(tool_tokens)
 
             remaining_tokens = self.sample_params.max_tokens - len(final_response_ids)
@@ -239,11 +240,17 @@ class ToolAgentLoop(AgentLoop):
 
         rollout_state.response_ids = final_response_ids[: self.sample_params.max_tokens]
         rollout_state.logprobs = final_logprobs[: self.sample_params.max_tokens]
-        rollout_state.response_mask = final_response_mask[: self.sample_params.max_tokens]
+        supervised = final_supervised[: self.sample_params.max_tokens]
         rollout_state.response = self.tokenizer.decode(rollout_state.response_ids)
 
-        assert len(rollout_state.response_ids) == len(rollout_state.logprobs)
-        assert len(rollout_state.response_ids) == len(rollout_state.response_mask)
+        # 语义洞直接烙进 labels：工具/环境 token 的 label 为 -100，不参与训练。
+        prompt_ids = list(rollout_state.prompt_ids or [])
+        rollout_state.input_ids = prompt_ids + list(rollout_state.response_ids)
+        rollout_state.labels = [-100] * len(prompt_ids) + [
+            resp_id if flag else -100 for resp_id, flag in zip(rollout_state.response_ids, supervised)
+        ]
+
+        assert len(rollout_state.input_ids) == len(rollout_state.labels) == len(rollout_state.logprobs)
 
         if rollout_state.status == Status.COMPLETED and self.judger is not None and not self.enable_batch_judge:
             rollout_state = await self.run_judger(rollout_state)
@@ -252,8 +259,8 @@ class ToolAgentLoop(AgentLoop):
 
 这个例子强调两个约定：
 
-- 模型生成 token 的 `response_mask` 为 `1`。
-- 工具或环境插入 token 的 `response_mask` 为 `0`，`logprobs` 填 `0.0`。
+- 模型生成 token 的 label 保留 token id（参与训练）。
+- 工具或环境插入 token 的 label 写 `-100`，`logprobs` 填 `0.0`。
 
 ### 覆盖 generate_group
 
@@ -302,7 +309,7 @@ class CustomAgentLoop(AgentLoop):
         return rollout_state
 ```
 
-如果任务的多轮上下文、工具结果或环境状态不能用默认 handler 合并，需要自己定义续跑逻辑。核心原则是：续跑后的 `response_ids`、`response`、`logprobs`、`response_mask` 必须仍然表示完整 response，而不是只有本轮新增部分。
+如果任务的多轮上下文、工具结果或环境状态不能用默认 handler 合并，需要自己定义续跑逻辑。核心原则是：续跑后的 `response_ids`、`response`、`logprobs`（以及自组装 loop 的 `input_ids`/`labels`）必须仍然表示完整 response，而不是只有本轮新增部分。
 
 ## 在训练配置中使用
 
@@ -341,8 +348,9 @@ agent_loop_manager_cfg = AgentLoopManagerConfig(
 - `generate_sample()` 是否只处理单条 `RolloutState`。
 - 推理前是否设置了 `rollout_state.tokens`。
 - 每次调用 `rollout_ctl.generate.remote()` 前是否设置了本轮 `sample_params`。
-- 返回训练前，`response_ids`、`response`、`logprobs`、`response_mask` 是否完整且长度一致。
-- 非模型生成 token 是否在 `response_mask` 中置为 `0`。
+- 返回训练前，`response_ids`、`response`、`logprobs` 是否完整且长度一致；自组装全序列的 loop 还需保证 `input_ids`/`labels` 齐全且等长。
+- 非模型生成 token 的 label 是否已写 `-100`（prompt+response 型样本由基类 canonicalize 兜底；自组装 loop 自行烙入）。
+- `response_ids` 是否为 `input_ids` 去掉 prompt 前缀的连续后缀（`len(response_ids) == len(input_ids) - len(prompt_ids)`）。
 - 需要 Judger 时，是否通过 `self.run_judger(...)` 调用打分，以复用 pause/cancel 处理。
-- 若使用 `_prepare_train_data()`，是否保证最终有 `reward["score"]`。
+- 是否保证最终有 `reward["score"]`。
 - 若使用 async partial rollout，是否正确处理 `enable_partial_rollout` 和历史 response 合并。

docs/zh_cn/rl/advanced_tutorial/judger.md

@@ -124,11 +124,11 @@ async def batch_judge(self, rollout_states: list[RolloutState]) -> list[RolloutS
 
 在预置 `SingleTurnAgentLoop` 中，进入 Judger 前的 `RolloutState` 通常已经包含：
 
-- `prompt_ids`：prompt token，后续 `_prepare_train_data()` 会用它拼接训练输入。
+- `prompt_ids`：prompt token，后续训练 controller 会用它拼接训练输入。
 - `response_ids`：模型生成的 response token，后续训练数据直接依赖它。
 - `response`：模型生成的文本，`NativeJudger` 会用它计算 reward。
 - `logprobs`：response token 的 rollout logprob。如果存在，长度必须和 `response_ids` 一致。
-- `response_mask`：哪些 response token 参与训练。为空时，`_prepare_train_data()` 默认所有 response token 都参与训练。
+- `labels`：监督目标（labels 是唯一监督载体，语义上不参与训练的位置为 `-100`）。prompt+response 型样本由 loop 侧 `canonicalize_train_fields()` 兜底构造。
 - `reward_model["ground_truth"]`：标准答案或标签，预置 rule-based Judger 通常依赖这个字段。
 
 `NativeJudger` 传给 `reward_handler` 的字段更窄：
@@ -154,13 +154,13 @@ rollout_state.reward = {
 }
 ```
 
-如果这批数据后续要进入 `_prepare_train_data()`，必须满足：
+如果这批数据后续要进入训练 controller，必须满足：
 
 - `reward` 不能为 `None`。
-- `reward` 必须包含数值型 `score` 字段，因为 `_prepare_train_data()` 会直接读取 `data.reward["score"]` 来计算 advantage。
+- `reward` 必须包含数值型 `score` 字段，因为 controller 会直接读取 `data.reward["score"]` 来计算 advantage。
 - `status` 不能是 `ABORTED`、`FILTERED` 或 `FAILED`。
 - `response` 和 `response_ids` 都不能为空。
-- 如果提供 `response_mask`，长度必须和 `response_ids` 一致。mask 为 `0` 的 token 会在训练 label 中变成 `-100`，对应 advantage 也会置为 `0.0`。
+- labels 中 `-100` 的位置不参与训练，对应 advantage 也会置为 `0.0`。
 - 如果提供 `logprobs`，长度必须和 `response_ids` 一致。
 
 其他 reward 字段可以按任务需要扩展，例如 `acc`、`format`、`tool_ok`、`reason` 等。

recipe/verl_agent/common/agent_loop_verl_tool.py

@@ -134,11 +134,20 @@ async def generate_sample(self, rollout_state: RolloutState) -> RolloutState:
         # TODO: handle samples with corrupted tool tokens ?
 
         # convert verl_tool_agent_loop output to rollout_state
-        rollout_state.prompt_ids = output.prompt_ids
-        rollout_state.response_ids = output.response_ids
-        rollout_state.logprobs = output.response_logprobs
+        # verl 的 response_mask（工具输出位为 0）是本 loop 的语义监督，直接烙进 labels（-100）。
+        prompt_ids = list(output.prompt_ids)
+        response_ids = list(output.response_ids)
+        semantic_mask = output.response_mask if output.response_mask is not None else [1] * len(response_ids)
+        rollout_state.prompt_ids = prompt_ids
+        rollout_state.response_ids = response_ids
+        rollout_state.logprobs = (
+            None if output.response_logprobs is None else [0.0] * len(prompt_ids) + list(output.response_logprobs)
+        )
         rollout_state.routed_experts = output.routed_experts
-        rollout_state.response_mask = output.response_mask
+        rollout_state.input_ids = prompt_ids + response_ids
+        rollout_state.labels = [-100] * len(prompt_ids) + [
+            resp_id if mask_value else -100 for resp_id, mask_value in zip(response_ids, semantic_mask)
+        ]
         rollout_state.status = Status.COMPLETED
         rollout_state.extra_fields.update(output.extra_fields)
         # judger needs response in text format

tests/rl/test_multi_task_agent_loop_manager.py

@@ -312,13 +312,15 @@ def test_trainer_rejects_conflicting_legacy_and_task_validity_checks(self):
             _migrate_legacy_is_valid_sample_fn(cfg, MagicMock())
 
     async def test_take_train_batch_applies_token_staleness_mask(self):
-        # 启用 token staleness 时，公开 produce_batch 路径应返回最终 effective mask。
+        # 启用 token staleness 时，公开 produce_batch 路径应把过期 token 的监督清出 labels（只清零不恢复）。
         state = RolloutState(
             rollout_id=1,
             group_id=1,
             message=[{"role": "user", "content": "prompt"}],
             prompt_ids=[1, 2],
             response_ids=[3, 4],
+            input_ids=[1, 2, 3, 4],
+            labels=[-100, -100, 3, 4],
             response_model_steps=[0, 4],
             status=Status.COMPLETED,
         )
@@ -344,17 +346,18 @@ async def test_take_train_batch_applies_token_staleness_mask(self):
 
         result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4)
 
-        self.assertEqual(result.rollout_states[0][0].response_mask, [0, 1])
+        self.assertEqual(result.rollout_states[0][0].labels, [-100, -100, -100, 4])
         self.assertEqual(replay_buffer.task_token_stale_threshold_calls, [{"task": 4}])
 
-    async def test_take_train_batch_skips_agentic_token_staleness_mask(self):
+    async def test_take_train_batch_skips_states_without_response_ids(self):
+        # response_ids 缺失的样本没有可对齐的 response 段，不参与 token staleness 烙制。
         state = RolloutState(
             rollout_id=1,
             group_id=1,
             message=[{"role": "user", "content": "prompt"}],
             input_ids=[1, 2],
             labels=[-100, 2],
-            response_mask=[1],
+            logprobs=[0.0, -0.1],
             status=Status.COMPLETED,
         )
         strategy = _FakeProduceStrategy(token_stale_threshold=4)
@@ -378,7 +381,79 @@ async def test_take_train_batch_skips_agentic_token_staleness_mask(self):
 
         result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4)
 
-        self.assertEqual(result.rollout_states[0][0].response_mask, [1])
+        self.assertEqual(result.rollout_states[0][0].labels, [-100, 2])
+
+    async def test_take_train_batch_bakes_full_sequence_agentic_per_token_staleness(self):
+        # agentic 全序列形态（response_ids 为 input_ids 去掉 prompt 前缀的后缀，含环境 token）
+        # 同样按 token 级 staleness 烙制：逐 token 按 response_model_steps 判定，仅过期 token
+        # 的监督位被清零。
+        state = RolloutState(
+            rollout_id=1,
+            group_id=1,
+            message=[{"role": "user", "content": "prompt"}],
+            prompt_ids=[1, 2],
+            response_ids=[30, -1, 40],
+            response_model_steps=[0, 4, 4],
+            input_ids=[1, 2, 30, -1, 40],
+            labels=[-100, -100, 30, -100, 40],
+            status=Status.COMPLETED,
+        )
+        strategy = _FakeProduceStrategy(token_stale_threshold=4)
+        manager = AgentLoopManager(
+            task_runners=[
+                _TaskRunner(
+                    task_name="task",
+                    agent_loop=_fake_agent_loop(),
+                    produce_strategy=strategy,
+                    sampler=_FakeSampler(),
+                    weight=1.0,
+                    order=0,
+                )
+            ],
+            replay_buffer=_FakeReplayBuffer(
+                rollout_states_by_task={"task": [[state]]},
+                leftover_counts={},
+            ),
+            rollout_controller=_fake_rollout_controller(),
+        )
+
+        result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4)
+
+        self.assertEqual(result.rollout_states[0][0].labels, [-100, -100, -100, -100, 40])
+
+    async def test_take_train_batch_sync_path_leaves_labels_untouched(self):
+        # 同步路径（无 token staleness）零操作：labels 在生成期已是最终态（语义洞已烙入）。
+        state = RolloutState(
+            rollout_id=1,
+            group_id=1,
+            message=[{"role": "user", "content": "prompt"}],
+            prompt_ids=[1, 2],
+            response_ids=[3, 4, 5],
+            input_ids=[1, 2, 3, 4, 5],
+            labels=[-100, -100, 4, 5, -100],
+            status=Status.COMPLETED,
+        )
+        manager = AgentLoopManager(
+            task_runners=[
+                _TaskRunner(
+                    task_name="task",
+                    agent_loop=_fake_agent_loop(),
+                    produce_strategy=_FakeProduceStrategy(),
+                    sampler=_FakeSampler(),
+                    weight=1.0,
+                    order=0,
+                )
+            ],
+            replay_buffer=_FakeReplayBuffer(
+                rollout_states_by_task={"task": [[state]]},
+                leftover_counts={},
+            ),
+            rollout_controller=_fake_rollout_controller(),
+        )
+
+        result = await manager.produce_batch(batch_size=1, train_step=5, model_step=4)
+
+        self.assertEqual(result.rollout_states[0][0].labels, [-100, -100, 4, 5, -100])
 
     async def test_produce_batch_allocates_by_weight_and_returns_task_sorted_results(self):
         # 共卡 produce_batch 按 task 权重分配 batch，并按 task 名稳定返回训练数据和 leftover 统计。