Reusing Rollouts under Policy Lag: Prefix-Normalized Policy Optimization for LLM Reinforcement Learning

2026-08-02

AI Summary

This work addresses the high computational cost of autoregressive rollouts and the issues of policy lag and distribution shift arising from multi-epoch reuse in large language model reinforcement learning. To overcome these challenges, the paper proposes Prefix-Normalized Policy Optimization (PNPO), which innovatively replaces conventional cumulative importance weights with the geometric mean of likelihood ratios computed over causal prefixes. This approach preserves prefix dependencies while compressing the dynamic range of importance weights, enabling stable and efficient off-policy updates. Experimental results demonstrate that under a four-epoch update setting, PNPO achieves an Avg@32 of 50.24% across multiple mathematical reasoning benchmarks, outperforming GSPO by 3 percentage points. Moreover, within the same computational budget, PNPO attains performance equivalent to that of single-epoch training with 600 rollout batches using only 150 batches.

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.

Problem

• policy lag
• off-policy learning
• rollout reuse
• large language models
• reinforcement learning

Innovation

• Prefix-Normalized Policy Optimization
• off-policy correction
• rollout reuse
• importance weighting
• large language model reinforcement learning