Async RL: enforce staleness by pacing sampling instead of discarding rollouts

[joschu]

Async RL: enforce staleness by pacing sampling instead of discarding rollouts

Problem

do_async_training enforced max_steps_off_policy by discarding any completed trajectory group whose sampler version was too old and requeuing its problem. When rollout durations vary (multi-turn/agentic tasks), this systematically starves slow rollouts:

• slow (usually harder) problems are regenerated over and over and rarely make it into a batch — the trained data distribution is silently biased toward fast-completing problems;
• every discard throws away completed sampler work;
• a problem whose rollout reliably takes more than max_steps_off_policy steps of wall-clock is never trained on.

This matches a failure mode reported by an external user running multi-turn tasks with 30–60 min rollouts ("longer rollouts... are constantly regenerated and rarely make it to a batch"). No published async RL system uses discard-and-requeue; the standard design (e.g. AReaL, arXiv:2505.24298) bounds staleness by admission control and trains oldest-first. Magistral (arXiv:2506.10910) states the invariant directly: the completion-length distribution of trained data must match the data distribution "even though shorter completions finish more quickly."

New semantics

AsyncConfig(max_steps_off_policy=K, groups_per_batch=B, pipeline_depth=D) — no data loss, and staleness is controlled by two separate knobs (max_steps_off_policy keeps its meaning; pipeline_depth is new, default 1):

1. Admission control: at most B·(D+1) group rollouts may be outstanding (started but not yet trained on), so sampling runs at most D iterations ahead — D is the typical staleness whenever sampling outpaces training (D=0 is synchronous; D ≤ K enforced). When the trainer stalls, rollout starts pause instead of racing ahead. Slots are released only after the optimizer step and sampling-client swap, so newly admitted rollouts sample from the new weights.
2. Deadline waiting: before training step t, the loop waits for every in-flight rollout with sampler version ≤ t − K — its last chance to be trained within the bound — instead of dropping it. K is the worst-case staleness that only slow-rollout stragglers approach; it can be generous without dragging bulk staleness up.
3. Stalest-first batches from a buffer of completed groups.

Together with one rollout worker per batch slot, these guarantee: every collected trajectory group is trained on exactly once, with staleness ≤ K (measured in training iterations = published sampler versions). K=0 reduces to exactly synchronous training. If the loop must wait a long time for a straggler, it logs a diagnostic warning every 2 minutes rather than stalling silently.

Removed: the discard/requeue path (per review, nobody should use it). async_config + stream_minibatch_config now raises ConfigurationError (the streaming path had no staleness enforcement of its own; previously the combination silently used discard semantics).

New metrics: async/staleness_{max,min,mean}, async/in_flight_groups, async/outstanding_groups, async/completed_buffer_size, time/waiting_for_stragglers.

Tests

tinker_cookbook/rl/async_training_test.py (~11 s, no network) runs the real do_async_training against in-process fakes with adversarial latency patterns. The fake sampling client stamps its weight version into the sampled tokens and the fake training client records when each datum is trained, so the staleness bound and exactly-once are verified from the training data itself, not the loop's bookkeeping. Covered: slow rollouts spanning many steps (30×/150× latency outliers), K=0 ⇒ fully synchronous, async beats sync wall-clock on straggler workloads, constant-reward (None) groups, datasets holding more groups than training consumes (no deadlock; leftovers skipped), num_substeps>1, resume from start_batch>0, config validation, logged metrics cross-checked against data.

These tests caught two real bugs during development (a slot-release race that let rollouts start on about-to-be-stale weights, and an off-by-one in the original capacity analysis), and were themselves hardened by mutation testing — 14 hand-written mutations of the staleness machinery; the suite now catches all of them.

Experiments

A/B on GSM8K (Qwen3.5-9B-Base, B=32, G=8, K=3, 40 steps) with difficulty-correlated env latency (problems with ≥5 reference-solution steps sleep 60–300 s per rollout — ~25% of data, disproportionately the hard problems), old semantics (main @ fad40e0) vs this branch, plus a sharper faster-step regime (B=16, K=2):

	blocking (new)	discard (old)
groups trained / sampled	1280 / 1280	1248 / 1360
slow-problem coverage	100%	92% (sharp regime: 92%)
slow share of trained batches (offered 0.25)	0.25	0.23 (sharp: 0.21)
max staleness	3 = K	≤ K by discarding
sampled tokens discarded	0	~143k (8%) (sharp: ~56k, 13%)

The experiments also motivated the pipeline_depth split: with a single knob, the pipeline filled its whole B·(K+1) budget whenever sampling outpaced training, so typical staleness sat at K rather than merely bounded by it (the old code masked this by discarding). Under the default unclipped importance_sampling loss, training uniformly at K=3 off-policy destabilized (sampler-trainer KL grew ~30×, entropy collapse); the same A/B with loss_fn="ppo" is stable — 94.5% (new) vs 94.0% (old) held-out accuracy after 40 steps, with the new arm training on every offered group and discarding nothing while the old arm still burned 105 group rollouts. A depth-1 rerun with the unclipped loss stayed healthy substantially longer (92.7% at step 20, measured mean staleness 0.97 vs 2.49 fully pipelined) but still destabilized late at this learning rate, so the recommendation is unconditional: prefer a trust-region loss for async training; pipeline_depth (default 1) controls how much off-policyness and trainer-waiting you trade for throughput, not stability. Follow-up controls and 3-seed replications strengthen this: a fully on-policy control with the unclipped loss is stable at the same LR (the instability is staleness-caused); with ppo, blocking (92.4 ± 1.8%) and discard (93.7 ± 0.5%) are statistically indistinguishable on held-out accuracy while blocking removes all waste and bias in every seed; one-sided cispo (caps 4.0 and 2.0) also destabilizes at staleness ≈ 2.5; and with a generous bound (K=8) the staleness deadline rarely binds, giving 1.8× faster wall-clock at lower realized staleness. Guidance in the docstring: prefer a trust-region loss for async training, and set max_steps_off_policy to comfortably cover your slowest rollouts.

Review

Adversarially reviewed by two independent reviewers (one with mutation testing — findings included a leaked-slot deadlock class the tests initially missed, an atomicity hazard now made structural, and an unsound proof sketch in a docstring, all fixed and re-verified; plus a GPT-5.6 review that caught the num_substeps staleness-unit ambiguity and the leftover-rollout waste at shutdown).

Follow-ups (out of scope here)

• Refresh the sampling client between turns of multi-turn rollouts (each turn is an independent sample call), so later turns of long rollouts come from newer weights — the turn-level analog of in-flight weight updates. Per-token sampling logprobs already make mixed-version trajectories valid for the loss, and the start version remains the staleness minimum, so bookkeeping is unchanged (TODO in trajectory_group_worker_loop).
• Wall-clock rollout budgets: on timeout, grade the partial trajectory and train on it with reward clamped ≤ 0 (needs Env-protocol support for partial-trajectory grading).
• Streaming-minibatch + async support (reserving per-minibatch slots for deadline stragglers).
• Async checkpoint/resume does not persist in-flight rollouts (pre-existing TODO).

🤖 Generated with Claude Code

https://claude.ai/code/session_0133j6tAVqf1Zd7nVeBVtKZQ

[YujiaBao]

A single permanently-hung rollout now wedges the whole run here, with no escape hatch. A rollout that never completes (stuck sandbox, hung sampling request — the exact failure mode RetryOnFailure.per_rollout_timeout's docstring says exists) never leaves _in_flight_by_version; once i_batch ≥ hung_version + K this wait can never satisfy num_deadline_stragglers == 0, and the shutdown sentinel needs the hung worker to exit. Since slots are only released by training, admission control then freezes all sampling too. The default config (rollout_error_tolerance=False → FailFast; per_rollout_timeout=0 disabled) has no mitigation, and multi-turn/agentic envs — this PR's target workload — are where hangs live.

Verified by execution with this PR's own fakes (48 problems, B=4, K=2, 12 batches, problem 0 hung):

	this branch	main (discard)
steps trained before stalling	2 / 12 (wedged at hung_version + K)	11 / 12 (only the final batch, one group short)
problems ever sampled	16 / 48 (admission control frozen)	48 / 48

So on main a hang costs one worker and the last step; here it halts all progress within K steps. It's loud (the warning below fired correctly, showing 7/4 groups collected, 1 in-flight rollouts at the staleness deadline) and arguably the intended no-data-loss tradeoff — but the warning doesn't tell the user what to do about it. Cheap fix: point the straggler warning text and the AsyncConfig docstring at RetryOnFailure(per_rollout_timeout=...) as the escape hatch. (The wall-clock-budget follow-up in the description would be the real fix.)

Review assisted by Claude.

[YujiaBao]

time/waiting_for_stragglers over-attributes: it charges the full duration of every queue.get() to stragglers whenever ≥1 deadline straggler is in flight — even when len(completed_groups_buffer) < groups_per_batch, i.e. the trainer would have blocked on batch-fill regardless (a wait that exists at any K, even in synchronous training). With K=0 it degenerates: the deadline is the current step, so every in-flight rollout counts as a straggler and the metric reports the entire collection wait. Someone tuning from this metric would raise K/pipeline_depth to fix a bottleneck that isn't staleness.

Since the loop can only reach the wait with a full buffer when a deadline straggler exists, restricting the attribution to that case makes the metric the marginal cost its name promises:

finally:
    # Only count waits where the deadline straggler was the sole blocker;
    # otherwise this was ordinary batch-fill waiting.
    if len(completed_groups_buffer) >= groups_per_batch:
        t_straggler_wait += time.time() - t_start

Review assisted by Claude.

[YujiaBao]

Nit: the three async gauges in one metrics row are snapshots from different instants — async/in_flight_groups and async/completed_buffer_size are captured in collect_batch before the train step, async/outstanding_groups here after release_slots. Fine once you know it, but the row can look self-inconsistent when debugging staleness (e.g. outstanding < in_flight + buffer). Maybe worth a one-line comment or capturing them at one point.

Review assisted by Claude.

[YujiaBao]

The torture tests are excellent for the happy-path scheduling space. Four gaps in the new machinery still worth covering, since the staleness guarantee depends on bookkeeping that only these tests protect:

1. Failing/raising rollouts: no test injects a rollout failure into do_async_training. An exception between record_started and record_completed_and_enqueue leaks a slot and an in-flight record; today that's masked by fail-fast crash semantics, but a future catch-and-log hardening would deadlock the deadline wait in every flaky-env run — with this suite still green (rollout_error_resilience_test.py never touches the async path).
2. save_every > 0 (the Config default is 20): all tests use 0, and FakeTrainingClient.create_sampling_client's NOTE says the fake can't model the periodic-checkpoint path. That path (save_periodic_async + create_sampling_client) is a different client-minting route than the one every staleness assertion validates.
3. evaluation_loop: never exercised (all tests set eval_every=0), including its interaction with the reworked shutdown cascade.
4. groups_per_batch=1: single worker, capacity = depth+1 — the tightest admission geometry.

Related nit: of the new metrics only async/staleness_max is cross-checked against ground truth; staleness_min/mean and the three gauges are unasserted.

Review assisted by Claude.

[YujiaBao]

Nit: spacing deviates from the file's convention — the other entries have a blank line between the --- separator and the ### heading, and this entry adds a doubled blank line after its closing --- (#802/#803 were exactly this class of cleanup).

Review assisted by Claude.

[YujiaBao]

Recipes can't reach this knob, but their existing flag silently changed meaning: five recipes (math_rl, code_rl, harbor_rl, rubric ×2) construct AsyncConfig with only the two old fields. A user re-running e.g. math_rl with their existing max_steps_off_policy=4 drops from an effectively 4-deep sampling pipeline to depth 1 with no CLI way back — and "set pipeline_depth=max_steps_off_policy to reproduce the previous behavior" is advice no recipe user can follow. Suggest plumbing pipeline_depth through the recipe CLIs (or at least the ones people use for long-rollout async runs).

Review assisted by Claude.

[YujiaBao]

Since this guidance is the PR's headline recommendation: the shipped research skill still teaches the opposite — skills/research/references/rl.md:304-308 recommends AsyncConfig(max_steps_off_policy=4, groups_per_batch=8) with the default importance_sampling loss and no mention of pipeline_depth, i.e. the configuration these experiments found destabilizes. Worth updating alongside this PR (out of this diff, so flagging it here).

Review assisted by Claude.