FlashREINFORCE: one stale rollout per prompt, corrected, gated, length-normalized

​

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 (different splits, and decontaminate still runs and counts).
2. Reward, both arms: the binary outcome, MathEqual against the GSM8K gold number. A program, not a judge. The paper changes the update, not the reward.
3. The regime, both arms: rollouts come from a stale sampler, a frozen copy of the policy’s LoRA adapter refreshed only every 4 optimizer steps, and the next batch is drawn before the current update lands, the way an asynchronous trainer overlaps generation with training [2]. A batch is therefore 1 to 4 updates old (the first is 0; mean 2.4 over 40 steps). One rollout per prompt, 64 prompts per step, 256 new tokens, temperature 1.0 and top-p 1.0 (no truncation, so the ratio has common support). Each token’s behavior_logprobs are the sampler’s, teacher-forced under the sampler’s weights right after it wrote them; logprobs are the policy’s at update time, one teacher-forced pass. LoRA rank 16, alpha 32, AdamW at 1e-5 on the adapter, gradient clip 1.0, one full-batch step per batch and the batch is discarded. 40 steps. Seed 17, set right before the adapter is built, so both arms draw one init and the same first five batches (the sampler is not refreshed until step 4, and batch 5 is drawn before that update).
4. Baseline arm: the paper’s uncorrected ablation. Batch-mean advantage A_i = R_i - mean R, token-mean loss: every token of trajectory i gets A_i / (total tokens in the batch). No importance ratio, no trust gate, no 1/T_i.
5. Recipe arm: wai.FlashReinforce().update(batch) at its defaults (trust 3e-3, FLASH_REINFORCE_TRUST): coefficient m_i * A_i * rho_it / (T_i * B). Both arms minimize -(coefficients * logprobs).sum() with the coefficients held constant; the selftest pins the recipe’s coefficients to the values tests/api/test_flash_reinforce.py pins. The FlashReinforce statistics (admitted share, sequence KL, ratio) are logged for both arms, so the baseline’s staleness is measured even though it ignores it.
6. Eval: pass@1 on the same 120 held-out tasks, 4 samples per task at temperature 0.7. The untrained base is evaluated three times first, and that spread is the noise floor a delta has to clear. Paired delta with a 95% interval (wai.pass_at, wai.compare), the training reward named as proxy (it is the target, so None), hack_scan on the last training batch.

​

Run

python recipe.py --selftest   # the reward, the lag bookkeeping, both coefficient rules, offline
python recipe.py              # both arms, 44 GPU minutes on one L40S, $1.48
python recipe.py --arm recipe --learning-rate 1e-4 --lag 8

​

Result

• The rollouts were barely stale. At 1e-5 on a rank-16 adapter, four updates move the policy very little: the mean sampled-action KL between sampler and learner per step was 2.9e-4 in the recipe arm and 3.7e-4 in the baseline (max over the run 9.1e-3 on one trajectory), and the mean ratio stayed within 0.9998 to 1.0010 of 1. The trust gate at 3e-3 masked one of 64 trajectories on 5 of 40 steps (admitted share 0.998 mean, never below 0.984); it would have masked one or two on 8 steps of the baseline. At lag 0 (step 1) the KL is exactly 0 and the ratio exactly 1, which is the check that the behavior log-probabilities are the sampler’s own.
• The uncorrected baseline did not degrade. It went from 0.40 to 0.47 on the holdout and its training reward climbed from 0.44 to the mid 0.50s (mean 0.49, last five steps 0.52 to 0.63). With the ratio this close to 1 and the gate this quiet, the correction had almost nothing to correct.
• What differed was the length weighting. With the ratio inert, the two arms differ in 1/T_i against a token mean. The token-mean baseline gives a long wrong trajectory more total push than a short one, and its completions shortened from 202 to 130 tokens over training (455 characters after, from 712 before). The recipe weights every trajectory the same, its completions stayed at 177 tokens (663 characters), and the hack scan’s strongest pooled feature on its last batch is truncated (r -0.58): the rollouts that ran into the 256-token cap and never gave a number. Its training reward averaged 0.45 (last five steps 0.42 to 0.55).

​

Checks

​

Climb

​

Learned

• The stale sampler costs nothing to build: a frozen copy of the adapter’s tensors, swapped in to generate and out to train, with the behavior log-probabilities teacher-forced under the swapped-in weights. Step 1 at lag 0 reports a KL of exactly 0 and a ratio of exactly 1, which is the one line that says the bookkeeping is right, and it is worth printing before anything else.
• wai.FlashReinforce().update(batch) is the whole update: the loop hands it 64 dicts of reward, logprobs and behavior_logprobs and multiplies its coefficients into a padded log-probability tensor. The baseline is the same loop with a different coefficient list, so the two arms cannot drift apart in anything but the rule.
• A correction can only be tested where there is something to correct. At 1e-5 on a rank-16 adapter the four-update lag left the ratio within a tenth of a percent of 1 and the gate touching one trajectory in five steps, so the paper’s stability claim was not exercised, and what the delta measured instead is the paper’s 1/T_i against the token mean, which at this scale kept long truncated failures alive that the baseline learned to cut. Check the KL the log prints on the first refreshed step before reading the delta.

​

Artifacts on Hugging Face

​

References

1. Lambert, N. Reinforcement Learning from Human Feedback. arXiv:2504.12501, 2025. Chapter Policy Gradient Algorithms.
2. Hu, Y., Zhang, Y., Zhang, J., Xu, Y., Zhang, Y., Peng, W., Yu, K., Molchanov, P., Kautz, J., Dong, Y. FlashREINFORCE: Critic-Free Single-Rollout Asynchronous RL for Agentic Language Models. NVIDIA, September 2026. https://yifanzhang-pro.github.io/FlashREINFORCE/FlashREINFORCE.pdf
3. Lambert, N. Reinforcement Learning from Human Feedback. arXiv:2504.12501, 2025. Chapter Evaluation.
4. Lambert, N. et al. Tülu 3: Pushing Frontiers in Open Language Model Post-Training. arXiv:2411.15124, 2024.
5. Gao, L., Schulman, J., Hilton, J. Scaling Laws for Reward Model Overoptimization. ICML 2023. arXiv:2210.10760.