# Policy versions and trajectory control A long rollout can outlive several optimizer updates. River lets you choose which weights generate its next tokens, how old those tokens may be when trained, and whether a group should wait for unfinished rollouts. These are separate decisions. The starting recipe is trajectory-pinned sampling, a bounded training age with `stale_policy="wait"`, and `GroupCompletion(mode="wait")`. ## One trajectory can contain several policies A **policy version** identifies model weights after an optimizer update. A **turn** is an assistant response, possibly ending in tool calls. A **segment** is the portion generated by one sampling request. Reaching `segment_tokens` continues the same turn in another request, subject to the rollout budget. `sampling_policy="trajectory"` pins the first segment's policy for the whole conversation. With `sampling_policy="segment"`, each request can select newer weights. This includes continuations within a turn, so a single assistant response can contain tokens generated by different policies. Every generated span retains its actual policy version and sampling log probabilities.
Pinned versus per-segment sampling Both trajectories contain two sampling segments in the first assistant turn and a third segment after a tool result. The pinned trajectory uses policy 10 throughout. Per-segment sampling uses policies 10, 11, and 12 as the trainer advances. At training policy 12 their ages are respectively 2, 2, 2 and 2, 1, 0. Pinned trajectory Trainer advances: 10 → 11 → 12 Segment 1 policy 10 Segment 2 policy 10 Segment 3 policy 10 Turn 1 Turn 2 ↑ tool result Policy ages at training version 12: 2, 2, 2 One policy for the entire conversation. Per-segment policy Trainer advances: 10 → 11 → 12 Segment 1 policy 10 Segment 2 policy 11 Segment 3 policy 12 Turn 1 Turn 2 ↑ tool result Policy ages at training version 12: 2, 1, 0 One conversation, multiple policies.
Each segment is one sampling request. A request keeps its selected weights until it completes; newer weights can be selected for the next request.
| RolloutEngine setting | Pinned trajectory | Per-segment policy | | --- | --- | --- | | `sampling_policy` | `"trajectory"` (default) | `"segment"` | | Sampling weights | The policy selected for the first segment. | A policy selected for each sampling request, subject to the KV bound. | | Retained prefix | Same-policy KV reuse. | Bounded-age KV reuse, or an explicitly allowed re-prefill. | | Training age | Every segment has the same policy age. | Each segment has its own age; newer tokens do not make earlier tokens younger. | The changing-policy diagram assumes compatible weights and retained-KV support are available. Per-segment sampling does not guarantee the latest weights: `KVCache(on_limit="hold")` can keep it on an older compatible policy. ## Wait for old trajectories or mask old spans `AsyncTrainer.max_staleness` bounds the distance from **each generated span's sampling policy** to the policy used for forward/backward. With a bound of 2, policy-10 tokens may train against policies 10, 11, or 12.
Wait versus mask with a two-update staleness bound A trajectory has segments from policies 10, 11, and 12. Wait preserves its opportunity to train against policy 12, where all segments are within the age bound. If the trainer advances to policy 13 under mask_span, policy 10 tokens are masked out of the training loss while policy 11 and 12 tokens remain trainable. Wait Training policy 12 · max_staleness=2 policy 10 age 2 train policy 11 age 1 train policy 12 age 0 train All segments contribute. Hold updates before this group expires. Sampling and tools can keep working. Mask old spans Training policy 13 · max_staleness=2 policy 10 age 3 masked policy 11 age 2 train policy 12 age 1 train Only in-bound spans contribute. Old tokens remain in the context. A pinned trajectory can be fully masked.
The sampling-policy choice determines which versions a trajectory contains. The stale-policy choice determines which generated tokens may contribute to training.
| `stale_policy` | Behavior | Tradeoff | | --- | --- | --- | | `"wait"` (default) | Holds optimizer progress so outstanding whole groups can train within the bound. Sampling and tools continue. | Preserves the groups' opportunity to train; slow trajectories can delay updates. | | `"mask_span"` | Masks generated spans that exceed the bound at training time. Their tokens remain in the conditioning context. | Updates need not wait for old spans, but some or all of a trajectory's training signal can be lost. Requires `allow_trajectory_loss=True`. | | `"keep"` | Trains spans regardless of age. | Does not enforce `max_staleness`; requires `allow_unbounded_staleness=True`. | The wait policy accounts for outstanding groups before they expire. It can hold an update earlier than the illustrated limit when several groups still need to fit into the remaining policy versions. A trajectory's own token and turn limits still apply. Masking changes which tokens contribute to the loss; it does not regenerate them. With a pinned trajectory, every segment ages together, so the entire trajectory can become masked. A batch with no usable gradient produces no optimizer update. ## Finish a group or bound its stragglers `GroupCompletion` answers a different question: when are enough rollouts from one question ready to form a training group? `mode="wait"` waits for every member. `mode="deadline"` starts a generated-token allowance once `min_members` have completed. For example, suppose two of four rollouts have finished: | Completion choice | Two finished rollouts | Two unfinished rollouts | | --- | --- | --- | | `mode="wait"` | Wait for the whole group. | Continue within their rollout budgets. | | `mode="deadline", on_stragglers="truncate"` | Join the group with the truncated members. | Stop at the deadline; the truncation policy decides their reward and training treatment. | | `mode="deadline", on_stragglers="discard"` | Form a partial group. | Stop and contribute no training data. | | `mode="deadline", on_stragglers="carry_over"` | Become available before the tails finish. | Continue for later training; requires `Batchwise` advantages. | A deadline is reached when any unfinished member uses its `max_straggler_tokens` allowance after that quorum. It is checked between sampling segments and cannot interrupt an in-flight request. It is not a wall-clock timeout. **Discard is a group-completion option, not a stale-policy option.** Partial-group completion (`discard` or `carry_over`) cannot use `stale_policy="wait"`; it requires `"mask_span"` or `"keep"`. Deadline truncation and discard also require `allow_trajectory_loss=True`. Carry-over requires asynchronous admission. ## Keep older KV or rebuild the prefix Trajectory-pinned sampling can reuse the unchanged prefix under the same policy in both synchronous and asynchronous training. With `sampling_policy="segment"`, the engine can extend an older prefix using newer weights. Configure its KV policy separately from the trainer's age limit: ```python kv_cache = rl.KVCache(max_staleness=1, on_limit="hold") ``` Pass this as `kv_cache=kv_cache` to `RolloutEngine` with `sampling_policy="segment"`. The trainer can still use `max_staleness=2`. | Bound | Compared policies | | --- | --- | | Trainer `max_staleness=2` | Each generated segment versus the policy used to train it. | | KV `max_staleness=1` | The retained cache's anchor policy versus the policy sampling the next segment. | Consider a prefix anchored at policy 10 with a KV age limit of 1: | `KVCache` choice | When policy 11 is available | When policy 12 is available | | --- | --- | --- | | `max_staleness=1, on_limit="hold"` | Can extend the retained cache with policy 11. | Keeps sampling at a compatible policy instead of advancing beyond the bound. | | `max_staleness=1, on_limit="refill", allow_reprefill=True` | Can extend the retained cache with policy 11. | Can switch to policy 12 by rebuilding the prefix under those weights. | If using current weights matters more than preserving the prefix, choose: ```python kv_cache = rl.KVCache( max_staleness=1, on_limit="refill", allow_reprefill=True, ) ``` Crossing the KV age limit starts a new prefill under current weights. The explicit flag acknowledges that cost. Availability depends on the model's capabilities; the engine checks support before sampling. `sampling_policy="segment"` requires either a positive KV age limit or explicit re-prefill permission. A zero KV age limit with `on_limit="hold"` is rejected. Trajectory pinning uses same-policy KV and does not need these cross-policy controls. Cache reuse still depends on residency and model support; a worker change or eviction can require a prefill in either mode. ### Check actual cache reuse `sampling/cached_prompt_fraction` in `step.metrics` is the fraction of reported prompt tokens served from cache. It is not a percentage of requests or a promise that a retained prefix stayed resident. Cache eviction, worker changes, or a new context can require a prefill. Image handles save repeated uploads but do not themselves guarantee a KV-cache hit.