Skip to content
Docs
Open Console
Guides

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.

3 of 5 · Previous: Asynchronous training · Next: Tools and images

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 settingPinned trajectoryPer-segment policy
sampling_policy"trajectory" (default)"segment"
Sampling weightsThe policy selected for the first segment.A policy selected for each sampling request, subject to the KV bound.
Retained prefixSame-policy KV reuse.Bounded-age KV reuse, or an explicitly allowed re-prefill.
Training ageEvery 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_policyBehaviorTradeoff
"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 choiceTwo finished rolloutsTwo 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:

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.

BoundCompared policies
Trainer max_staleness=2Each generated segment versus the policy used to train it.
KV max_staleness=1The 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 choiceWhen policy 11 is availableWhen 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=TrueCan 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:

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.

Next: Add tools and images →