feat: max_off_policy_steps masks stale tokens instead of cancelling episodes

docs/training.md

@@ -64,7 +64,7 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli
 | `orchestrator.batch_size` | Tasks per trainer step. |
 | `orchestrator.constant_trainer_batch_size` | Keep trainer batches constant when samples have no training signal, such as zero advantage on all tokens. Enabled by default. Disable it for faster collection with variable trainer batch sizes. |
 | `orchestrator.train.group_size` | Rollouts generated per task. |
-| `orchestrator.max_off_policy_steps` | Maximum staleness of a trained rollout (default 8): the version a batch trains on minus the oldest version that generated the rollout, queue time included. Episodes past the bound are dropped; a group shares one dispatch version, so its episodes age out together. The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `off_policy/*` and `mismatch_kl/all/mean` when tuning. |
+| `orchestrator.max_off_policy_steps` | Maximum staleness of a trained token (default 8): the version a batch trains on minus the policy version that sampled the token, queue time included. Older tokens are masked out of the loss; the rest of the sample still trains. Tokens without a policy label (Dynamo engines) count as old as their episode's dispatch version. A sample is dropped only when nothing is left to train, and rollouts are never cancelled for staleness. The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `off_policy/token/masked`, `off_policy/*` and `mismatch_kl/all/mean` when tuning. |
 | `[orchestrator.train.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `rae`, `hierarchical_grpo`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). |
 | `trainer.model.freeze_moe_router` | Freeze MoE router gates (default `true` for RL, `false` for SFT). Router updates in RL shift the trainer's expert choices away from the ones inference sampled, amplifying the train-inference mismatch, so RL keeps the pretrained router. Set `false` to train it. |
 | `[[orchestrator.train.source]]` | Training sources. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Training sources](configuration.md#training-sources-orchestratortrainsource). |

packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py

@@ -549,7 +549,7 @@ class OrchestratorConfig(BaseConfig):
     """Maximum training steps. If None, runs indefinitely."""
 
     max_off_policy_steps: int = Field(8, ge=0)
-    """Maximum staleness of a trained rollout: the version a batch trains on (v{step-1}) minus the oldest version that generated the rollout (a rollout can span several weight updates), queue time included. Episodes past the bound are dropped, in-flight and queued; a group shares one dispatch version, so its episodes age out together. Higher values yield better throughput at the cost of off-policy noise."""
+    """Maximum staleness of a trained token: the version a batch trains on (v{step-1}) minus the policy version that sampled the token, queue time included. Engines label each sampled token with the policy step (vLLM ``weight_versions``); unlabelled tokens (Dynamo engines) count as old as their episode's dispatch version. Older tokens are masked out of every loss component; the rest of the sample still trains with its own advantage and inference logprobs. A sample is dropped only when nothing is left to train. Rollouts are never cancelled for staleness. Higher values yield better throughput at the cost of off-policy noise."""
 
     heartbeat: HeartbeatConfig | None = None
     """BetterStack heartbeat configuration for monitoring training progress."""

src/prime_rl/eval/runner.py

@@ -123,7 +123,6 @@ async def setup(self, *, skip_first_step: bool = False, is_resumed: bool = False
             initial_max_inflight=self.concurrency.max_inflight,
             max_inflight_ceiling=config.concurrency.max_inflight,
             dispatch_per_minute=config.dispatch_per_minute,
-            max_off_policy_steps=0,
             run_id=self.run_id,
             run_name=self.run_name,
             on_episode_complete=self.concurrency.record_episode,

src/prime_rl/orchestrator/dispatcher.py

@@ -14,13 +14,10 @@
   schedule next. Transitions are level-triggered (driven by the eval
   source's emptiness), so in-flight episodes of the opposite kind drain
   naturally on either side of an eval boundary.
-- ``on_version_pending`` (called by the watcher before the engines pause for
-  the weight update) drops train groups already past ``max_off_policy_steps`` — a
-  compute-saving early cancel; the sink's queue sweep is what guarantees the
-  bound. Eval episodes are measurements for the policy version they started
-  with. Online evals may explicitly cancel them when a newer checkpoint is ready. Train
-  episodes sampled from a frozen model never go stale — their generation
-  source doesn't change with policy updates.
+- Staleness never cancels a rollout: tokens past ``max_off_policy_steps`` are
+  masked by the sink's queue sweep. Eval episodes are measurements for the
+  policy version they started with. Online evals may explicitly cancel them
+  when a newer checkpoint is ready.
 """
 
 from __future__ import annotations
@@ -55,7 +52,6 @@
     Progress,
     WorkKind,
 )
-from prime_rl.orchestrator.utils import min_fresh_version
 from prime_rl.utils.async_utils import safe_cancel, safe_cancel_all
 from prime_rl.utils.logger import get_logger
 
@@ -124,8 +120,8 @@ class Dispatcher:
     """``await dispatcher.start()`` runs the dispatch loop until ``stop()``.
     Pulls examples from ``TrainSource`` / ``EvalSource``, schedules episodes
     under shared capacity, and emits native verifier episodes to ``out_q``.
-    The watcher drives ``on_version_pending`` for staleness
-    cancellation; the orchestrator triggers eval epochs."""
+    The watcher drives ``on_version_pending`` around weight updates; the
+    orchestrator triggers eval epochs."""
 
     def __init__(
         self,
@@ -140,7 +136,6 @@ def __init__(
         initial_max_inflight: int,
         max_inflight_ceiling: int | None,
         dispatch_per_minute: float | None,
-        max_off_policy_steps: int,
         run_id: str,
         run_name: str | None,
         policy_weights_change: bool = True,
@@ -156,7 +151,6 @@ def __init__(
         self.policy_clients = policy_clients
         self.train_source = train_source
         self.eval_source = eval_source
-        self.max_off_policy_steps = max_off_policy_steps
         self.run_id = run_id
         self.run_name = run_name
         # Called with ``total_tokens`` per completed episode
@@ -402,45 +396,15 @@ async def publish_live(self) -> None:
             await self.flush_live()
 
     async def on_version_pending(self, step: int) -> None:
-        """Drop train groups past ``max_off_policy_steps``: a group dispatched at
-        v{k} ships at earliest in the batch currently collecting, at staleness
-        ``(progress.step - 1) - k`` — beyond the bound it can never train, so
-        cut it before more inference sinks in. This is a compute saver; the
-        sink's queue sweep is what guarantees no stale episode ships. Eval
-        groups never go stale (they are tied to their start-time policy
-        version), nor do frozen-sourced train groups (their generation source
-        doesn't change with policy updates).
-
-        Runs *before* the inference engines are paused for the weight update so
-        the resulting aborts are processed while the engine is still stepping —
-        otherwise the orphaned KV transfers crash the decode engine on resume
-        (see ``WeightWatcher.apply_policy_update``)."""
+        """Block train scheduling across the weight update. Staleness never
+        cancels a rollout: it keeps sampling fresh tokens, and the sink masks
+        the ones that age past ``max_off_policy_steps``."""
         self.policy_update_pending = True
         # Wait for a scheduling call that started before the pending update.
         # No rollout can cross the inference weight swap after this barrier.
         async with self.scheduling_lock:
             pass
 
-        if self.train_envs is None or self.progress is None:
-            return
-        min_version = min_fresh_version(self.progress.step, self.max_off_policy_steps)
-        stale_groups = [
-            gid
-            for gid, group in self.groups.items()
-            if group.kind == "train"
-            and self.train_envs.get(group.env_name).uses_live_policy
-            and group.policy_version_at_start < min_version
-        ]
-        cancelled = 0
-        for gid in stale_groups:
-            cancelled += await self.drop_group(gid, reason="stale")
-
-        if cancelled:
-            get_logger().warning(
-                f"Cancelled {cancelled} train episodes past max_off_policy_steps={self.max_off_policy_steps}. "
-                "Consider increasing it to avoid this."
-            )
-
     async def on_new_version(self, step: int) -> None:
         """Resume rollout scheduling after inference applies the new policy."""
         self.policy_update_pending = False

src/prime_rl/orchestrator/metrics.py

@@ -31,7 +31,7 @@ def dispatch_failure_metrics(failures: list[DispatchFailure], *, prefix: str, to
 @dataclass(frozen=True)
 class TraceRecord:
     """A trace joined with its episode and training selection state.
-    ``cancelled`` marks episodes the pipeline voided (stale drop) — a
+    ``cancelled`` marks episodes the pipeline voided (stale or shutdown) — a
     selection verdict disjoint from trace errors."""
 
     episode: vf.Episode
@@ -635,7 +635,7 @@ def observe(
         self.attempts += attempts
         discarded = [episode for episode in episodes if not any(t.id in queued_trace_ids for t in episode.traces)]
         self.discarded += len(discarded) + len(failures) + cancelled
-        if cancellation is not None and cancellation.reason != "stale":
+        if cancellation is not None:
             self.cancelled += sum(episode.id in episodes.cancelled for episode in discarded) + cancelled
         else:
             self.stale += sum(episode.id in episodes.cancelled for episode in discarded) + cancelled

src/prime_rl/orchestrator/orchestrator.py

@@ -336,7 +336,6 @@ async def setup(self) -> None:
             initial_max_inflight=self.concurrency.max_inflight,
             max_inflight_ceiling=config.concurrency.max_inflight,
             dispatch_per_minute=config.dispatch_per_minute,
-            max_off_policy_steps=config.max_off_policy_steps,
             run_id=self.run_id,
             run_name=self.run_name,
             on_episode_complete=self.concurrency.record_episode,