Skip to main content
The scripts are in the repository, not in the installed package. Clone it, then cd recipes/papers/sao-single-rollout before running the commands below. Browse this recipe on GitHub.
Paper: Single-Rollout Asynchronous Optimization for Agentic Reinforcement Learning, Hou, Li, Tang, Dong (Tsinghua), arXiv:2607.07508, July 2026. https://arxiv.org/abs/2607.07508 Book: the policy gradient with a learned value baseline and generalized advantage estimation, and importance sampling as the correction when the rollouts were written by an older policy than the one being trained [1]. Claim: one rollout per prompt is enough when a value critic supplies the baseline (length-adaptive GAE, lambda = 1 - 1/(1.5 L) [5]) and the staleness is handled by direct double-sided importance sampling, which masks any token whose current/rollout probability ratio leaves (0.7, 6.0) instead of clipping it; the policy then trains stably on rollouts one to several updates old, the shape a production trace arrives in. The change: the recipe arm masks a token outside the ratio band to zero; the baseline arm applies the same ratio to the same critic advantages and never masks.

​
Recipe

  1. Base: Qwen/Qwen2.5-1.5B-Instruct. Data: GSM8K, 512 train prompts from the train split, 120 held out from the test split, decontaminated against the holdout before training.
  2. The production-trace regime, both arms: one rollout per prompt, 64 prompts per optimizer step, 256 new tokens at most, temperature 1.0 with no top-p, top-k or repetition penalty (the ratio needs common support). The rollouts come from a stale sampler: a frozen copy of the policy’s LoRA adapter, refreshed from the policy only every LAG = 4 optimizer steps, so a batch was written by a policy 1 to 4 updates old (0 on the very first policy step). Each generated token carries the sampler’s own log-probability, read off the generation scores as the token is drawn, and the current policy’s log-probability from one teacher-forced pass at update time.
  3. The critic, both arms: one linear layer on the policy trunk’s last hidden state at each generated token, on the same forward pass, with the hidden states detached so the critic never trains the trunk. Its own AdamW at 1e-4 (the paper’s 5e-6 is a full separate 30B value model, so this is an override: a one-layer head at 5e-6 does not leave its zero init in 50 steps), critic_steps = 2 squared-error updates per policy update toward update.value_targets, and critic_warmup = 10 critic-only steps first. Both are read off wai.SAO().
  4. The policy, both arms: LoRA rank 16, alpha 32 (TRAINING_LORA_RANK, TRAINING_LORA_ALPHA), AdamW at 1e-5 on the adapter (the paper’s 1e-6 is full weights of a 30B MoE; 1e-5 is the adapter rate the FlashReinforce and BPCO recipes share, so the three compare), gradient clip 1.0, one step per batch, 40 steps after the warmup, seed 17 set right before the adapter is built. The reward is the binary outcome, MathEqual against the GSM8K gold: a program, not a judge.
  5. Baseline arm: wai.SAO().update(batch) for the length-adaptive GAE advantages over the critic, then update.advantages times the unmasked ratio, over the same 1/N token mean: plain importance-weighted policy gradient, the paper’s “no DIS” case. Recipe arm: update.coefficients as they come, the ratio inside (0.7, 6.0) and zero outside. The loss on both sides is -(coefficients * logprobs).sum() with the coefficients held constant. Same update call, same advantages, same critic: the arms differ in the mask and nothing else (--selftest checks both on a hand-made batch).
  6. Eval: pass@1 on the same 120 held-out tasks, 4 samples per task at temperature 0.7, the untrained base three times first for the noise floor, each arm once after training; paired delta with a 95% interval (wai.compare with run_std, run_std_runs; proxy= left unset because the training reward is the target).

​
Run

train_log.json holds every step of both arms (reward, masked share, ratios, advantage, explained variance, critic loss, length, staleness); the curves block of results.json folds each to first, last, mean, max and min.

​
Result

Run 2026-09-22, both arms, at the protocol above (round 1 of the Climb table). Recipe vs baseline: -0.004 [-0.040, +0.031] over 120 paired tasks. Verdict: unresolved. One training seed per arm; a second seed on each arm, passed as train_runs=, would resolve it, and at these settings it can only resolve to flat: the two arms ran the same experiment. The band masked 0.03% of tokens on the recipe arm (0.04% would have been masked on the baseline), because four adapter updates at 1e-5 do not move a token’s probability out of (0.7, 6.0). The median per-step maximum ratio was 1.5 under a lag of 4 and 1.7 at lag 0, where it is pure generate-versus-teacher-forced numerics; the band starts at 6.0 and 0.7. Both arms went 0.36 to 0.46, which is the critic-baselined single-rollout policy gradient working on stale rollouts either way, not the mask. So the protocol run cannot test the paper’s claim, and the Climb table’s round 2 is the run that can: at 1e-4 the band binds (1 to 4% of tokens at lag 4), the unmasked baseline sees ratios of 20 to 38 on its worst token and applies them, and both arms still reach 0.65 to 0.66. That round is also unresolved, and its interval also covers zero.

​
Checks

Nothing in this table is ticked by hand: every cell is written by recipe.py into results.json. These are round 1’s numbers. The critic: explained variance of the value head against the per-token return went from 0.10 at the end of the warmup to 0.12 to 0.13 over the last ten steps on both arms (critic loss 0.22 -> 0.21). A one-layer head on the frozen trunk reads the batch mean and a little of the prompt’s difficulty; it is a baseline, not a predictor. Gradient norm 0.05 on average, so the clip at 1.0 never bound. No coefficient was non-finite and no step was skipped on either arm in either round.

​
Climb

Round 1 could not test the paper, and the log says so. The recipe’s own per-step masked share is the check: 0.0000 at lag 0, 0.0001 at lag 1, 0.0004 at lag 4, 0.0024 at its worst step. A mask that fires on three tokens in ten thousand cannot change a gradient by anything the eval could see. The same line was in the smoke run before the real one, so the second round was launched alongside the first rather than after it. Round 2 makes the band live and the arms still agree. At 1e-4 the masked share on the recipe arm rises with the lag as it should, 0.002 at lag 1, 0.008 at lag 2, 0.013 at lag 3, 0.017 at lag 4, 0.033 at its worst step, and the baseline’s would-be-masked share is about twice that at every lag (0.007, 0.018, 0.029, 0.040, peak 0.100). The median per-step maximum ratio at lag 4 is 6.5 on the recipe arm and 21 on the baseline (its worst 38), so the masked policy drifts about a third as far from its sampler as the unmasked one: DIS holds the policy near the rollout distribution by refusing the gradient on the tokens it has already left. The baseline applied those ratios of 20 to 38 as weights and did not blow up: gradient norm peaked at 0.30 against the recipe’s 0.13, the clip at 1.0 never bound, training reward went 0.45 -> 0.76 on the baseline and 0.46 -> 0.78 on the recipe, and pass@1 0.66 against 0.65. The eval noise floor in that round is 0.132 (run_std 0.0217), so the delta of -0.010 is inside it. The recipe arm’s rollouts got shorter faster (133 tokens against 174 over the last ten steps, truncation 0.04 against 0.14, eval reply 384 chars against 483), which is the only place the arms visibly differ. What would move this off “unresolved” is a second training seed per arm on round 2. What would let it test the paper’s stability claim is a lag the unmasked ratio cannot survive: the paper’s regime is a 30B model, full weights at 1e-6, on rollouts many updates old from a separate rollout engine, and that is where it argues the unmasked ratio is the problem; four adapter updates at 1e-4 on a 1.5B model is not, and the baseline here is the evidence. --lag 16 at 1e-4 is the next knob.

​
Learned

  • The band is dead at the protocol’s step size, and the recipe now says so in its own log before a GPU hour is spent. masked_share per step is the reachability check: 3 tokens in 10,000 at lr 1e-5, lag 4. Check that the change is reachable before you read a delta over it.
  • Masking outside the band keeps the policy near its sampler. With the band live, the recipe arm’s worst ratio at lag 4 was 6.5 against the baseline’s 21, and its masked share half the baseline’s would-be share, at every lag. That is the mechanism working. Whether it matters for the eval needs a regime where the unmasked ratio actually hurts, which four adapter updates on a 1.5B model is not.
  • One rollout per prompt trains, with a one-layer critic. Both arms went 0.36 -> 0.46 at 1e-5 and 0.38 -> 0.66 at 1e-4 with no group baseline, on rollouts 1 to 4 updates stale, with a value head whose explained variance never passed 0.21. A batch-mean-quality baseline was enough here; a second seed per arm would say whether the delta between masked and unmasked is anything at all.
Verified 2026-09-22, whileai 0.114 (the checkout’s whileai/ package, mounted into the container: wai.SAO is newer than the wheel), transformers 4.54.0 + PEFT 0.16.0 on torch 2.7.1, a plain PyTorch loop, no TRL. Round 1: 45.8 GPU minutes, 1.53ontwoL40Sinparallel(baseline25.7withthethreebaseevals,recipe20.1);round2:42.0minutes,1.53 on two L40S in parallel (baseline 25.7 with the three base evals, recipe 20.1); round 2: 42.0 minutes, 1.40; smoke: 14.1 minutes, $0.47. Run page: none (no WHILEAI_API_KEY in the run).

​
Artifacts on Hugging Face

Part of the Papers, replicated collection in the while-ai org.

​
References

  1. Lambert, N. Reinforcement Learning from Human Feedback. arXiv:2504.12501, 2025. Chapter Policy Gradient Algorithms.
  2. Lambert, N. Reinforcement Learning from Human Feedback. arXiv:2504.12501, 2025. Chapter Evaluation.
  3. Lambert, N. et al. Tülu 3: Pushing Frontiers in Open Language Model Post-Training. arXiv:2411.15124, 2024.
  4. Gao, L., Schulman, J., Hilton, J. Scaling Laws for Reward Model Overoptimization. ICML 2023. arXiv:2210.10760.
  5. Yue, Y. et al. VAPO: Efficient and Reliable Reinforcement Learning for Advanced Reasoning Tasks. arXiv:2504.05118, 2025.
Last modified on September 22, 2026