Command Palette
Search for a command to run...
REUSING ROLLOUTS UNDER POLICY LAG: PREFIX-NORMALIZED POLICY OPTIMIZATION FOR LLM REIN-FORCEMENT LEARNING
REUSING ROLLOUTS UNDER POLICY LAG: PREFIX-NORMALIZED POLICY OPTIMIZATION FOR LLM REIN-FORCEMENT LEARNING
Abstract
Autoregressive rollout generation is a major computational cost in reinforcement learning for large language models. Reusing each rollout batch for additional learner updates amortizes this cost, but later updates become increasingly off policy as the learner departs from the behavior policy. At a token position, exact off-policy correction must account for both the current action and the probability of reaching its prefix. The cumulative importance ratio provides this correction, but its product form can produce an unwieldy dynamic range. We study Prefix-Normalized Policy Optimization (PNPO), which replaces the cumulative ratio with the geometric mean of likelihood ratios along each causal prefix, preserving causal-prefix dependence at each position while compressing the log-weight scale. In controlled long-context mathematical reasoning experiments, we induce two off-policy regimes by using one or four policy-update epochs per rollout batch. PNPO does not consistently outperform GSPO with one epoch. With four epochs, it attains the highest observed Avg@32 on each benchmark; the unweighted mean of the three independently selected benchmark peaks is 50.24, 3.00 percentage points above GSPO. Under a matched 2,400-update budget, four-epoch PNPO reaches a final macro Avg@32 of 49.66 after 150 rollout batches, comparable to the 49.56 reached after 600 batches with one epoch. These results provide preliminary evidence that PNPO can be advantageous as training moves further off-policy.
One-sentence Summary
Tencent et al. propose Prefix-Normalized Policy Optimization (PNPO), which replaces the cumulative importance ratio with the geometric mean of likelihood ratios along each causal prefix to compress log-weight scale while preserving causal-prefix dependence, and in long-context mathematical reasoning experiments, PNPO attains a 3 percentage point higher Avg@32 than GSPO under four-epoch updates, demonstrating efficient off-policy LLM reinforcement learning.
Key Contributions
- Prefix-Normalized Policy Optimization (PNPO) replaces the exact cumulative importance ratio with the geometric mean of token likelihood ratios over each causal prefix, preserving prefix-conditioned dependence while compressing the log-weight dynamic range.
- In controlled mathematical reasoning experiments, PNPO attains a 3.00 percentage point higher unweighted mean peak Avg@32 than GSPO in a four-epoch off-policy regime, with the advantage becoming more pronounced as learner–behavior mismatch increases.
- Under a matched 2,400-update budget, four-epoch PNPO reaches a final macro Avg@32 of 49.66 after 150 rollout batches, comparable to the 49.56 reached after 600 batches with one epoch, providing preliminary evidence of more effective rollout reuse under greater off-policy mismatch.
Introduction
Policy-gradient methods are a standard approach for post-training language models, but collecting rollout batches through autoregressive generation is expensive. To amortize this cost, methods like PPO reuse collected rollouts for multiple learner updates, which introduces off-policy drift as the learner changes while the behavior data remains fixed. Proximal policy methods address this mismatch with local-ratio surrogates, but these omit a full state–action correction and weaken as reuse increases. The authors propose Prefix-Normalized Policy Optimization (PNPO), which uses the geometric mean of likelihood ratios along the causal prefix to correct for prefix visitation while controlling the scale of cumulative log-ratios. Their main contribution is formulating this prefix-normalized weight and evaluating it in controlled off-policy regimes, where PNPO shows a pronounced advantage over sequence-shared weighting when repeated updates induce greater learner–behavior mismatch.
Method
The authors propose Prefix-Normalized Policy Optimization (PNPO), a method for reinforcement learning fine-tuning of large language models that addresses the state-distribution shift inherent in standard proximal policy optimization (PPO) and its variants such as GRPO. These methods typically rely on a local surrogate that freezes the state occupancy at the behavior policy, using only the token-level action ratio ρt=πθ(at∣st)/πβ(at∣st) and a clipped advantage. While the resulting policy gradient is unbiased at the behavior policy, it omits the prefix-state correction required for exact off-policy estimation when the learner drifts away from the behavior policy.
In the autoregressive setting, the prefix-state ratio factorizes as a product of past action ratios, leading to the exact joint state–action change-of-measure ratio
Ct=k=1∏tρk.This cumulative ratio, when paired with the appropriate advantage, can yield an unbiased gradient estimate. However, using Ct directly introduces a scale problem: its log-variance grows with prefix length and policy drift, making gradient estimates unstable. Furthermore, retaining the behavior-advantage Atβ while using the current-token score zt=∇θlogπθ(at∣st) would be biased; to pair Atβ with an unbiased score, one must use the cumulative prefix score ∑k=1tzk, not the current-token score alone. The authors therefore develop a practical surrogate that interpolates between exact correction and scale control.
The core of PNPO is the prefix-normalized policy weight, which compresses the cumulative log-ratio by taking the t-th root of Ct:
wi,tPN=Ci,t1/t=exp(t1k=1∑tlogρi,k).This transform preserves the sign of logCt and the order across responses, while reducing the dynamic range. At t=1, it recovers the local action ratio; at t=Li, it equals the geometric mean of the full-response ratio. For intermediate positions it depends only on the prefix, excluding future-suffix likelihood shifts, and thus is not a true density ratio but a monotone approximation.
To further control the influence of tokens with extreme weight values, PNPO employs an acceptance gate that hard-rejects score terms at positions where the weight falls outside a position-dependent tolerance interval. The gate uses a heuristic scaling factor h(t,Li)=Li/t, which widens the bounds at early positions and shrinks them to the base interval [1−ϵ−,1+ϵ+] at the final token. The mask Mi,t is a binary indicator; if the weight is outside the bounds, the corresponding token’s score term is simply omitted from the objective, while later positions remain eligible.
The complete PNPO objective is a detached score-function surrogate that treats the gate, prefix-normalized weight, and advantage as constant coefficients via stop-gradient. For a batch of prompts, the authors sample a group of G responses from the behavior policy, compute a group-relative advantage Ai,t (using the same normalization as GRPO), and then maximize
JPNPO(θ)=E[G1i=1∑GLi1t=1∑Lisg[Mi,twi,tPNAi,t]logπθ(yi,t∣x,yi,<t)].The per-response averaging first normalizes over valid tokens (rejected tokens contribute zero), then averages the per-response means across responses. The objective thus retains the current-token-score structure of PPO/GRPO while incorporating a prefix-aware weight that approximates the state-distribution correction.
During training, PNPO reuses a fixed rollout batch for multiple optimizer epochs. The behavior-policy log probabilities, group-derived advantages, and response lengths are held constant. In each minibatch update, the learner evaluates the current numerator log probabilities, recomputes the cumulative ratios, the prefix-normalized weight, and the acceptance mask, and applies the objective. This allows the weight to adapt as the policy evolves, while the advantage and behavior-policy denominator remain fixed. The overall design trades exact off-policy correction for stable scale control, aiming to mitigate the state-distribution drift that can accumulate when a surrogate is repeatedly optimized over the same stale trajectories.
Experiment
The experiments evaluate PNPO on long-context mathematical reasoning using DeepSeek-R1-Distill-Qwen-1.5B fine-tuned on DAPO-Math-17k, comparing against GSPO and GRPO on AMC 2023, AIME 2024, and AIME 2025. Under four PPO epochs per rollout batch, PNPO achieves the best performance on all benchmarks, with a macro average 3.00 points above GSPO, and the advantage persists across the entire training trajectory rather than only at peak scores. Four-epoch PNPO matches the final performance of the one-epoch setting while using only a quarter of the newly generated responses, demonstrating more effective rollout reuse and reaching a given reward threshold substantially earlier than GSPO. GRPO fails to maintain early gains over long horizons, and while GSPO remains stable, a gap in evaluation metrics remains, consistent with the benefit of PNPO's token-level weighting granularity.
Under four PPO epochs, PNPO yields the highest average evaluation scores across all benchmarks, with a 3 percentage point advantage over GSPO. This lead persists across the evaluation trajectory, and PNPO achieves comparable final performance to its one-epoch counterpart using only a quarter of the newly generated responses, indicating more efficient rollout reuse. PNPO also reaches a fixed reward threshold hours earlier than GSPO in the four-epoch setting. With four PPO epochs, PNPO attains the best observed Avg@32 on all three benchmarks, averaging 50.24 compared to GSPO's 47.24. In the one-epoch setting, PNPO and GSPO remain close and exchange the lead, while in the four-epoch setting PNPO leads at 14 of 15 evaluations and finishes 2.66 percentage points ahead. At the same total optimizer updates, four-epoch PNPO reaches a final macro Avg@32 of 49.66 after 150 rollout batches, comparable to the 49.56 reached after 600 one-epoch batches, showing more effective rollout reuse. Four-epoch PNPO reaches a centered reward threshold of 0.25 in 16.0 hours, 6.4 hours earlier than GSPO and much earlier than the one-epoch runs. GRPO's four-epoch evaluation declines after peaking, and GSPO's stable reward does not close the evaluation gap, consistent with PNPO's position-dependent prefix statistic providing an advantage.
In experiments comparing PNPO, GSPO, and GRPO across benchmarks under one and four PPO epochs, PNPO consistently achieves the highest evaluation scores, particularly in the four-epoch setting where it maintains a clear lead and reaches reward thresholds faster. PNPO demonstrates more efficient reuse of rollouts, matching the final performance of one-epoch training with far fewer new responses, while GRPO's performance degrades after peaking and GSPO fails to close the gap. The advantage is attributed to PNPO's position-dependent prefix statistic.