feat(orchestrator): expose target_lag config

docs/training.md

@@ -58,7 +58,8 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli
 |---|---|
 | `orchestrator.batch_size` | Tasks per trainer step. |
 | `orchestrator.group_size` | Rollouts generated per task. |
-| `orchestrator.max_off_policy_steps` | How many distinct policies may have contributed to one rollout before it's discarded (default 8). The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `errored_rollouts` and `mismatch_kl/all/mean` when tuning. |
+| `orchestrator.max_off_policy_steps` | Per-rollout staleness cap (default 8). Lag = `(trainer_step - 1) - policy_version_at_generation`, including rollouts waiting in the train sink. Bump for throughput on long agentic rollouts; lower for tighter on-policyness. Watch `errored_rollouts` and `mismatch_kl/all/mean` when tuning. |
+| `orchestrator.target_lag` | Dispatch gate: max batches the orchestrator may ship ahead of `policy.version` (default 1). Pauses rollout scheduling until the trainer publishes the next checkpoint — **not** the same knob as `max_off_policy_steps`. Watch `time/wait_for_policy` when tuning. |
 | `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). |
 | `[[orchestrator.train.env]]` | Training environments. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Environments](configuration.md#environments-orchestratortrainenv). |
 | `[[orchestrator.eval.env]]` + `orchestrator.eval.interval` | Eval environments and cadence (default every 100 steps). |

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

@@ -529,7 +529,16 @@ class OrchestratorConfig(BaseConfig):
     """Maximum training steps. If None, runs indefinitely."""
 
     max_off_policy_steps: int = Field(8, ge=0)
-    """Maximum policies allowed to generate a single rollout. Rollouts generated more than ``max_off_policy_steps`` ahead of training are discarded. Higher values yield better throughput at the cost of off-policy noise."""
+    """Maximum trainer–policy lag allowed per rollout before it is discarded.
+
+    Lag is measured as ``(trainer_step - 1) - policy_version_at_generation``,
+    covering both in-flight rollouts and rollouts waiting in the train sink.
+    This is **not** the orchestrator dispatch gate (see ``target_lag``)."""
+
+    target_lag: int = Field(1, ge=1)
+    """Maximum batches the orchestrator may ship ahead of ``policy.version``.
+    Default ``1`` preserves the standard async trainer–inference pipeline.
+    Unrelated to ``max_off_policy_steps`` (per-rollout staleness)."""
 
     bench: bool = False
     """Benchmark mode. Sets ``max_steps`` to 5 and disables W&B."""

src/prime_rl/orchestrator/orchestrator.py

@@ -97,8 +97,9 @@
 
 # Maximum batches the orchestrator may run ahead of the trainer. The
 # dispatcher is paused via ``update_dispatch_gate`` once this is exceeded;
-# resumed when the watcher advances ``policy.version``.
-TARGET_LAG = 1
+# resumed when the watcher advances ``policy.version``. Configurable via
+# ``orchestrator.target_lag`` (default 1). Unrelated to ``max_off_policy_steps``.
+DEFAULT_TARGET_LAG = 1
 
 
 class Orchestrator:
@@ -825,7 +826,7 @@ def update_dispatch_gate(self) -> None:
         )
         gate = self.dispatcher.dispatch_allowed
         was_set = gate.is_set()
-        if lead > TARGET_LAG and not building_final_batch_nccl:
+        if lead > self.config.target_lag and not building_final_batch_nccl:
             if was_set:
                 get_logger().info(
                     "Pausing dispatcher to prevent orchestrator from racing from trainer. Waiting for new policy..."

tests/unit/orchestrator/test_off_policy_lag.py

@@ -0,0 +1,16 @@
+"""Tests for target_lag orchestrator config (#2812)."""
+
+import pytest
+from pydantic import ValidationError
+
+from prime_rl.configs.orchestrator import OrchestratorConfig
+
+
+def test_target_lag_defaults_to_one():
+    cfg = OrchestratorConfig()
+    assert cfg.target_lag == 1
+
+
+def test_target_lag_rejects_zero():
+    with pytest.raises(ValidationError):
+        OrchestratorConfig(target_lag=0)