[Feat] Add off-policy masking for partial rollouts

docs/en/rl/advanced_tutorial/rl_trainer.md

@@ -140,7 +140,7 @@ Parameter meanings:
 | `over_sample_threshold` | The ratio of extra samples that may be generated. A larger value makes the rollout side easier to keep fully loaded, but may produce more samples that are not from the current step. |
 | `enable_partial_rollout` | Whether rollouts paused before weight synchronization may continue after synchronization. Before using this for tool calling or multi-turn tasks, confirm that the AgentLoop supports continuation. |
 | `max_staleness` | The number of synchronization cycles by which samples may lag behind the current training progress. A larger value gives more throughput flexibility but weakens the on-policy property. |
-| `tail_batch_trigger_size` | When expired samples accumulate to this number, tail batch mode is entered and these samples are retried first. |
+| `tail_batch_trigger_size` | Expired-sample retry policy: `-1` disables rerollout, `0` retries immediately without entering tail batch mode, and a positive value enters tail batch mode after that many expired groups accumulate. |
 
 `max_staleness` is counted in "weight synchronization cycles". The actual expiration threshold used in code is:
 
@@ -161,12 +161,11 @@ Both oversampling and partial rollout are affected by `max_staleness`:
   by the earliest model version in the response, so continuation across synchronization cycles also needs room from
   `max_staleness`.
 
-Tail batch is used to handle samples that have expired during asynchronous production. When the number of `expired`
-samples reaches `tail_batch_trigger_size`, `AsyncProduceStrategy` enters tail batch mode: this round no longer
-oversamples according to `over_sample_threshold`, only fills the required target, and retries samples from the
-expired sample pool first. You can understand it as a non-oversampling synchronous fill-up production. Its goal is
-not to improve throughput, but to collect long-tail expired samples again and avoid leaving them in the buffer for
-too long.
+Expired samples can be rerolled out according to `tail_batch_trigger_size`. `-1` disables rerollout. `0` retries
+expired groups as soon as they appear while retaining the normal asynchronous production and oversampling policy.
+For a positive value, `AsyncProduceStrategy` waits until the expired pool reaches the configured size, then enters
+tail batch mode: this round no longer oversamples according to `over_sample_threshold`, only fills the required
+target, and retries samples from the expired sample pool first.
 
 Note: it is not recommended to set `max_staleness>0` and `enable_partial_rollout=False` at the same time. With this
 combination, long-tail oversampled samples may be reset after weight synchronization because partial rollout is not

docs/zh_cn/rl/advanced_tutorial/rl_trainer.md

@@ -131,7 +131,7 @@ produce_strategy_config = AsyncProduceStrategyConfig(
 | `over_sample_threshold` | 允许额外生成的比例。值越大，rollout 侧越容易保持满载，但也可能产生更多非当前 step 的样本。 |
 | `enable_partial_rollout` | 权重同步前被暂停的 rollout 是否允许在同步后续跑。工具调用或多轮任务使用前需要确认 AgentLoop 支持续跑。 |
 | `max_staleness` | 允许样本相对当前训练进度滞后的同步周期数。值越大，吞吐更宽松，on-policy 程度更弱。 |
-| `tail_batch_trigger_size` | 过期样本累计到一定数量后，进入 tail batch 模式，优先重试这些样本。 |
+| `tail_batch_trigger_size` | 过期样本重试策略：`-1` 关闭 rerollout，`0` 立即重试但不进入 tail batch 模式，正数表示累计到指定 group 数量后进入 tail batch 模式。 |
 
 `max_staleness` 按“权重同步周期”计数。代码中实际使用的过期阈值是：
 
@@ -146,7 +146,7 @@ stale_threshold = (max_staleness + 1) * sync_weights_interval
 - `over_sample_threshold>0` 会为未来 step 提前生成样本。如果这些样本跨过下一次权重同步点，只有 `max_staleness` 允许时才会继续保留为可训练样本。
 - `enable_partial_rollout=True` 会让被暂停的 response 在同步后续跑。样本的 staleness 按 response 中最早的模型版本计算，因此跨同步周期续跑时也需要 `max_staleness` 留出空间。
 
-tail batch 用于处理异步生产中已经过期的样本。当 `expired` 样本数量达到 `tail_batch_trigger_size` 时，`AsyncProduceStrategy` 会进入 tail batch 模式：本轮不再按 `over_sample_threshold` 超发，只补齐必要目标，并优先从过期样本池中取样重试。可以把它理解为一次非超发的同步补齐生产；它的目的不是提高吞吐，而是把长尾过期样本重新收集起来，避免它们长期留在 buffer 中。
+`tail_batch_trigger_size` 控制过期样本的 rerollout。设置为 `-1` 时关闭 rerollout；设置为 `0` 时，过期 group 一出现就立即优先重试，但仍保持普通异步生产和 oversampling 策略；设置为正数时，`AsyncProduceStrategy` 等待 expired pool 累积到指定数量后进入 tail batch 模式，本轮不再按 `over_sample_threshold` 超发，只补齐必要目标，并优先从过期样本池中取样重试。
 
 注意：不建议同时设置 `max_staleness>0` 且 `enable_partial_rollout=False`。这种组合下，长尾超发样本在权重同步后可能因为不支持 partial rollout 被重置（当前在 `RolloutWorker` 中重置样本只保留 prompt 字段）；但由于每次重置过期信息归0，它们不会过期，tail batch 不会及时接管，下一轮同步窗口内仍然可能生成不完并反复重试。当前还没有支持 `tail_batch_max_tries` 机制来按重试次数触发 tail batch。因此 `max_staleness>0` 时，优先开启 `enable_partial_rollout=True`。
 

examples/v1/config/rl_disagg_multi.py

@@ -70,7 +70,7 @@
 sync_weights_interval = int(os.environ.get("SYNC_WEIGHTS_INTERVAL", "1"))
 over_sample_threshold = float(os.environ.get("OVER_SAMPLE_THRESHOLD", "0.0"))
 partial_rollout = os.environ.get("PARTIAL_ROLLOUT", "0") == "1"
-tail_batch_trigger_size = int(os.environ.get("TAIL_BATCH_TRIGGER_SIZE", "0"))
+tail_batch_trigger_size = int(os.environ.get("TAIL_BATCH_TRIGGER_SIZE", "-1"))
 max_staleness = int(os.environ.get("MAX_STALENESS", "0"))
 enable_evaluate = os.environ.get("ENABLE_EVALUATE", "0") == "1"
 gsm8k_task_weight = float(os.environ.get("GSM8K_TASK_WEIGHT", "3.0"))

examples/v1/config/rl_disagg_single.py

@@ -80,7 +80,7 @@
 sync_weights_interval = int(os.environ.get("SYNC_WEIGHTS_INTERVAL", "1"))
 over_sample_threshold = float(os.environ.get("OVER_SAMPLE_THRESHOLD", "0.0"))
 partial_rollout = os.environ.get("PARTIAL_ROLLOUT", "0") == "1"
-tail_batch_trigger_size = int(os.environ.get("TAIL_BATCH_TRIGGER_SIZE", "0"))
+tail_batch_trigger_size = int(os.environ.get("TAIL_BATCH_TRIGGER_SIZE", "-1"))
 max_staleness = int(os.environ.get("MAX_STALENESS", "0"))
 prompt_repeat_k = int(os.environ.get("PROMPT_REPEAT_K", "4"))
 rollout_tp_size = int(os.environ.get("ROLLOUT_TP_SIZE", "1"))

tests/rl/test_multi_task_agent_loop_manager.py

@@ -12,15 +12,19 @@
 import asyncio
 import unittest
 from typing import Any
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
-from xtuner.v1.data_proto.rl_data import Status
+from xtuner.v1.data_proto.rl_data import RolloutState, Status
 from xtuner.v1.rl.agent_loop_manager.agent_loop_manager import (
     AgentLoopManager,
     AgentLoopManagerConfig,
     TaskSpecConfig,
 )
-from xtuner.v1.rl.agent_loop_manager.disagg_agent_loop_manager import DisaggAgentLoopManager
+from xtuner.v1.rl.agent_loop_manager.disagg_agent_loop_manager import (
+    DisaggAgentLoopManager,
+    DisaggAgentLoopManagerConfig,
+    DisaggTaskSpecConfig,
+)
 from xtuner.v1.rl.agent_loop_manager.produce_utils import (
     GROUP_GENERATE_TIME_KEY,
     ProduceBatchStatus,
@@ -41,10 +45,12 @@ def __init__(
         self,
         cleanup_pause_time_s: float = 0.0,
         stale_threshold: int = 1,
-        tail_batch_trigger_size: int = 0,
+        tail_batch_trigger_size: int = -1,
+        token_stale_threshold: int | None = None,
     ):
         self.cleanup_pause_time_s = cleanup_pause_time_s
         self.stale_threshold = stale_threshold
+        self.token_stale_threshold = token_stale_threshold
         self.tail_batch_trigger_size = tail_batch_trigger_size
         self.called_batch_sizes: list[int] = []
         self.called_train_steps: list[int] = []
@@ -122,6 +128,7 @@ def __init__(self, rollout_states_by_task: dict[str, list[list[Any]]], leftover_
         self._leftover_counts = leftover_counts
         self.refresh_staleness_calls: list[tuple[str, int, int, tuple[Status, ...]]] = []
         self.expired_groups_retryable_calls: list[dict[str, bool]] = []
+        self.task_token_stale_threshold_calls: list[dict[str, int]] = []
 
     async def get(self, batch_size: int, task_name: str, group_status: Status):
         assert group_status == Status.COMPLETED
@@ -137,11 +144,13 @@ async def refresh_staleness(
         self,
         *,
         task_stale_thresholds: dict[str, int],
+        task_token_stale_thresholds: dict[str, int] | None = None,
         expired_groups_retryable_by_task: dict[str, bool] | None = None,
         current_train_step: int,
         statuses: list[Status] | None = None,
     ):
         self.expired_groups_retryable_calls.append(dict(expired_groups_retryable_by_task or {}))
+        self.task_token_stale_threshold_calls.append(dict(task_token_stale_thresholds or {}))
         expired_counts = {}
         for task_name, stale_threshold in task_stale_thresholds.items():
             self.refresh_staleness_calls.append(
@@ -191,6 +200,8 @@ def _fake_rollout_controller():
 
 
 class TestMultiTaskAgentLoopManager(unittest.IsolatedAsyncioTestCase):
+    """共卡与多 task manager 的 batch 生产、消费和统计行为。"""
+
     def test_manager_config_accepts_single_task_spec(self):
         # 单 task 配置可以直接传入，兼容最小 AgentLoopManager 配置。
         task = TaskSpecConfig.model_construct(
@@ -205,11 +216,80 @@ def test_manager_config_accepts_single_task_spec(self):
 
         self.assertEqual(manager_config.tasks.task_name, "single_task")
 
+    async def test_take_train_batch_applies_token_staleness_mask(self):
+        # 启用 token staleness 时，公开 produce_batch 路径应返回最终 effective mask。
+        state = RolloutState(
+            rollout_id=1,
+            group_id=1,
+            message=[{"role": "user", "content": "prompt"}],
+            prompt_ids=[1, 2],
+            response_ids=[3, 4],
+            response_model_steps=[0, 4],
+            status=Status.COMPLETED,
+        )
+        strategy = _FakeProduceStrategy(token_stale_threshold=4)
+        replay_buffer = _FakeReplayBuffer(
+            rollout_states_by_task={"task": [[state]]},
+            leftover_counts={},
+        )
+        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=replay_buffer,
+            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].response_mask, [0, 1])
+        self.assertEqual(replay_buffer.task_token_stale_threshold_calls, [{"task": 4}])
+
+    async def test_take_train_batch_skips_agentic_token_staleness_mask(self):
+        state = RolloutState(
+            rollout_id=1,
+            group_id=1,
+            message=[{"role": "user", "content": "prompt"}],
+            input_ids=[1, 2],
+            labels=[-100, 2],
+            response_mask=[1],
+            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].response_mask, [1])
+
     async def test_produce_batch_allocates_by_weight_and_returns_task_sorted_results(self):
         # 共卡 produce_batch 按 task 权重分配 batch，并按 task 名稳定返回训练数据和 leftover 统计。
         strategy_a = _FakeProduceStrategy(tail_batch_trigger_size=2)
         strategy_b = _FakeProduceStrategy()
-        strategy_c = _FakeProduceStrategy()
+        strategy_c = _FakeProduceStrategy(tail_batch_trigger_size=0)
         replay_buffer = _FakeReplayBuffer(
             rollout_states_by_task={
                 "task_a": [["a-0"], ["a-1"]],
@@ -268,7 +348,7 @@ async def test_produce_batch_allocates_by_weight_and_returns_task_sorted_results
         self.assertIn("task_c", result.task_results)
         self.assertEqual(
             replay_buffer.expired_groups_retryable_calls,
-            [{"task_b": False, "task_a": True, "task_c": False}],
+            [{"task_b": False, "task_a": True, "task_c": True}],
         )
         self.assertEqual(strategy_a.called_expired_groups_retryable, [True])
         self.assertEqual(strategy_a.cleanup_expired_groups_retryable, [True])

tests/rl/test_prepare_train_data.py

@@ -18,7 +18,7 @@
 import numpy as np
 import torch
 
-from xtuner.v1.data_proto.rl_data import RolloutState, Status
+from xtuner.v1.data_proto.rl_data import RolloutState, Status, reset_rollout_response
 from xtuner.v1.train.rl_trainer import BaseRLTrainer
 
 
@@ -111,6 +111,22 @@ def test_text_path_builds_shifted_training_tensors(self):
         self.assertEqual(info["response_len/mean"], 3.0)
         self.assertEqual(info["prompt_len/mean"], 3.0)
 
+    def test_rerolled_state_without_semantic_mask_uses_all_response_tokens(self):
+        trainer = self._build_trainer([1.0])
+        state = reset_rollout_response(self._state(response_mask=[0, 1, 0]))
+        state.response = "rerolled response"
+        state.response_ids = [30, 31]
+        state.logprobs = [0.1, 0.2]
+        state.reward = {"score": 1.0}
+        state.status = Status.COMPLETED
+        state.finish_reason = "stop"
+
+        data_batches, _ = self._prepare(trainer, [[state]])
+
+        self.assertIsNone(state.response_mask)
+        self.assertEqual(data_batches[0]["shifted_labels"].tolist(), [[-100, -100, 30, 31]])
+        self.assertEqual(data_batches[0]["advantage"], [1.0, 1.0, 1.0, 1.0, 1.0])
+
     def test_multi_sample_group_uses_each_sample_reward_and_advantage(self):
         # 同一个 prompt 下的多个 response 要分别使用自己的 reward 和 advantage。
         trainer = self._build_trainer([1.5, -2.0])