cs.LGOct 13, 2026

Enabling Token-wise Importance Sampling in Off-Policy RLVR with Proximal KL Regularization

Authors: Hao Bai, Jiarui Yao, Tong Zhang

Organizations: University of Illinois Urbana-Champaign

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Off-policy data in reinforcement learning with verifiable rewards (RLVR) can arise from asynchronous generation, rollout-training mismatch, and response reuse. Token-wise importance weighting avoids multiplying likelihood ratios across a response but generally biases policy gradients. We prove that, under coverage and local expressiveness conditions, full future KL credit makes the exact proximal optimizer the update's unique stationary policy, even for mixtures of older policies, without full-response weights or a fitted critic. We further establish convergence of continuous-time tabular population updates under persistent coverage and time-varying behavior policies, and a finite-sample reward guarantee for linear logistic models under additional assumptions. These results motivate Proximal KL, a practical off-policy RLVR method that separates the proximal reference from the sampling policy and replaces PPO clipping with bounded future KL credit. Across all seven matched off-policy comparisons on Qwen3-1.7B and Qwen3-4B, spanning two learning rates, response reuse, and rollout staleness, Proximal KL outperforms GRPO in micro-averaged accuracy over eight math benchmarks by up to 2.58%. A separate on-policy comparison yields comparable aggregate accuracy. In periodic-refresh experiments, Proximal KL avoids the collapses observed in PG and CISPO, exceeds GRPO's training accuracy, and reaches higher accuracy in fewer policy updates. Ablations show that token-wise correction outperforms sequence-level and unweighted controls. With token-wise correction retained, future KL feedback reduces sensitivity to the tested baseline perturbation relative to PPO clipping. A finite future horizon improves AIME24 mean@16 and pass@16 over the full tail, while fixed-policy measurements show that truncation reduces correction variance over the tested horizons.

Figures & tables

Figure 1: Finite future-KL window. The policy conditions on the full prefix, while KL credit includes the current token and up to H−1H-1 subsequent tokens.
Figure 2: At fixed ρi>1+ϵclip\rho_{i}>1+\epsilon_{\mathrm{clip}} , the unperturbed advantage is zero and A^m=±σ\widehat{A}_{m}=\pm\sigma with equal probability. (a) PPO clipping removes positive update shifts. (b) The mean shift is −σ/2-\sigma/2 with clipping and zero without it; clipping also reduces conditional coefficient variance. The common importance-weighted policy score is factored out.
Figure 3: Periodic-refresh training with FP8 rollouts and BF16 optimization at 1×10−61\times 10^{-6} , with a rollout refresh interval of 64 optimizer updates. Panels (a) through (d): 1.7B, 4B, 8B, and 32B. All panels use BF16 KV. EMA α=0.10\alpha=0.10 . Configurations: Appendix B.2 .
Figure 4: Stability across six settings. Top: AIME24 pass@16; bottom: policy entropy on logarithmic axes. The first four columns match Figure 5 ; the last two use 1.7B, 1×10−51\times 10^{-5} , at staleness levels 4 and 16. EMA α=0.20\alpha=0.20 . Bottom groups distinguish the smaller-model/lower-rate settings from larger-model or higher-rate settings; each column specifies its rollout staleness. Display limits and configurations: Appendix C.2 .
SettingStale-LRMethodMATH500OlympiadMinervaGaokaoAIME24AIME25AMC23AMC24Average
 ness  pass@1pass@1pass@1pass@1mean@16mean@16mean@16mean@16 
Model: Qwen3-1.7B
---Base56.1023.1522.4351.174.272.9227.6613.4724.15
On-policy05e-6GRPO68.93 ±\pm 1.0332.84 ±\pm 1.2828.92 ±\pm 0.7759.53 ±\pm 0.458.89 ±\pm 0.965.97 ±\pm 0.9446.30 ±\pm 2.1726.53 ±\pm 0.2434.50
On-policy05e-6Proximal KL67.80 ±\pm 1.2033.14 ±\pm 0.6729.90 ±\pm 1.0658.40 ±\pm 0.9210.07 ±\pm 1.205.62 ±\pm 0.7545.10 ±\pm 0.3925.42 ±\pm 0.2834.09
Off-policy05e-6GRPO65.40 ±\pm 2.4628.24 ±\pm 0.6228.68 ±\pm 2.0558.05 ±\pm 1.068.96 ±\pm 1.576.74 ±\pm 0.4342.24 ±\pm 2.2622.78 ±\pm 0.9731.99
Table 1: Raw checkpoint accuracy (%). 1.7B: mean ±\pm sample SD over three seeds. Average pools 4,149 responses per seed across eight benchmarks and reports the mean across seeds; Bold: higher mean in each column within a matched comparison; both bold if the unrounded gap is below 0.50%0.50\% . Protocols are in Appendix C.1 .
Figure 5: AIME24 at 5×10−65\times 10^{-6} . Top: mean@16; bottom: pass@16. Columns: 1.7B staleness 4, 1.7B staleness 16, 4B staleness 4, and 4B staleness 16. Configurations are given in Appendix C.2 . EMA α=0.20\alpha=0.20 ; shared Base within each model at step 0.
StalenessMethodMATH500OlympiadMinervaGaokaoAIME24AIME25AMC23AMC24Average
  pass@1pass@1pass@1pass@1mean@16mean@16mean@16mean@16 
4→164\rightarrow 16GRPO↓ 1.60\downarrow\,1.60↓ 0.59\downarrow\,0.59↓ 0.37\downarrow\,0.37↓ 2.09\downarrow\,2.09↓ 0.42\downarrow\,0.42↑ 0.83\uparrow\,\mathbf{0.83}↓ 8.59\downarrow\,8.59↑ 4.31\uparrow\,\mathbf{4.31}↓ 1.04\downarrow\,1.04
4→164\rightarrow 16Proximal KL↓ 1.40\downarrow\,\mathbf{1.40}↑ 1.34\uparrow\,\mathbf{1.34}→ 0.00\rightarrow\,\mathbf{0.00}↑ 1.04\uparrow\,\mathbf{1.04}↑ 2.71\uparrow\,\mathbf{2.71}↓ 0.42\downarrow\,0.42↓ 4.53\downarrow\,\mathbf{4.53}↑ 2.92\uparrow\,2.92↑ 0.22\uparrow\,\mathbf{0.22}
Table 2: Qwen3-4B accuracy changes (%) from staleness 4 to 16 at 5×10−65\times 10^{-6} . Differences use unrounded Table 1 scores. ↑\uparrow / ↓\downarrow : increase/decrease; →\rightarrow : unchanged. Bold: larger signed change.
Figure 6: Importance-weighting comparison: 1.7B, 5×10−65\times 10^{-6} , staleness 16, full tail, β=0.30\beta=0.30 . Token-specific future KL credit is retained in all three updates. EMA α=0.20\alpha=0.20 ; shared evaluation step 0.
Figure 7: Baseline perturbation at 1.7B, 5×10−65\times 10^{-6} , staleness 4. Solid/dashed: unperturbed (clean)/perturbed. Shading shows the score difference between the two curves. EMA α=0.20\alpha=0.20 .
Figure 8: Future-cost clipping: 1.7B, 5×10−65\times 10^{-6} , staleness 4, full tail, β=0.30\beta=0.30 ; EMA 0.200.20 , shared evaluation step 0. Gradient norms precede optimizer clipping.
Figure 9: Future KL horizon ablation. (a,b) 1.7B, 5×10−65\times 10^{-6} , staleness 4; EMA 0.200.20 . (H,β)=(1,4.80)(H,\beta)=(1,4.80) and (64,0.075)(64,0.075) match βH=4.8\beta H=4.8 . (c) Fixed-policy correction variance.
Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Ideal conditionImplemented choiceEffect on the sampled update
Full future KL tailHorizon HH , including H=fullH=\mathrm{full}Selects how far credit extends; a finite window omits later costs.
Unclipped future log ratioOptional log-ratio clippingBounds each future cost and changes its expectation when active; the choice is independent of the sampling policy.
History-only baselineGroup-normalized reward advantageSets the reward scale from sampled responses; the resulting field can differ from the ideal one.
Population response averageMinibatch token averageAdds sampling noise; token normalization can change the relative weights of responses.
Appendix
Table 3: From the stationary-policy result to the implemented update.
MethodSettingQwen3
 Model size1.7B4B8B32B
SharedLearning rate1×10−61\times 10^{-6}
 Rollout refresh interval64
 Prompts ×\times responses per batch32×832\times 8
 Optimizer updates per batch1
 Prompt / response limits (tokens)1,024 / 16,384
Appendix
Table 4: Key hyperparameters for training with periodic policy refresh. Merged cells denote shared settings; refresh intervals count optimizer updates.
MethodSettingOff-policy
 ModelQwen3-1.7BQwen3-4B
 Learning rate5×10−65\times 10^{-6}1×10−51\times 10^{-5}5×10−65\times 10^{-6}
GRPOPPO clipYes
 Prompts per batch128
 Responses per prompt16
 Response limit (tokens)8,192
Appendix
Table 5: Hyperparameters for training with controlled staleness. Columns specify controlled rollout staleness 4; the comparison at staleness 16 uses the corresponding setting. At 1.7B, the 1×10−51\times 10^{-5} Proximal KL run with staleness 16 uses H=64H=64 , β=0.60\beta=0.60 . At 4B, staleness 4 uses a 32-token unit-weight prefix and a linearly decayed remainder; staleness 16 uses H=8H=8 , with β=0.25\beta=0.25 in both cases. PG shares GRPO’s batch, sampling, optimizer, reward-normalization, and staleness settings. Merged cells denote shared values.
Figure 10: Qwen3-1.7B GRPO learning-rate sweep, staleness 4. Dotted line: 1,000 updates; EMA α=0.20\alpha=0.20 .
Figure 11: GRPO clipping-width comparison: Qwen3-1.7B, 1×10−51\times 10^{-5} , staleness setting 4. Accuracy uses EMA α=0.20\alpha=0.20 ; gradient norms are raw and shown on a logarithmic axis.
More in this topic

Off-Policy RL

Off-policy reinforcement learning studies how agents learn from data generated by other or earlier policies, improving sample efficiency while managing distribution shift.

See all 124 papers

Explore similar work

CardsList
Sep 30, 2026

Advancing Entropy-Level Credit Assignment in RLVR via Proximal Entropy Policy Optimization

Value-model-free RLVR methods such as GRPO assign uniform advantages to all tokens in a rollout, ignoring that tokens contribute unequally. Recent methods use token entropy as an importance proxy but compute it globally across the batch, conflating importance with prompt difficulty and positional trends. We argue that importance should instead be measured relative to the local context of each token. We introduce proximal entropy, a local measure of token importance relative to neighboring tokens, and prove it is invariant to both confounders. Proximal Entropy Policy Optimization (PEPO) uses it to weight per-token advantages and outperforms GRPO and entropy-based baselines on mathematical reasoning across Qwen3-1.7B, Qwen3-4B, and Llama-3.2-3B-Instruct. We also show the formulation generalizes to other algorithms where substituting proximal entropy into existing methods improves, and applying it to single-stream RL succeeds where global entropy fails.

Yun Kim, Nojun Kwak

Seoul National University

RL for Language Model ReasoningCredit Assignment in RLStep-Level Credit Assignment in RLReinforcement Learning with Verifiable Rewards
AskPDF
May 8, 2026

Rethinking Importance Sampling in LLM Policy Optimization: A Cumulative Token Perspective

Reinforcement learning, including reinforcement learning with verifiable rewards (RLVR), has emerged as a powerful approach for LLM post-training. Central to these approaches is the design of the importance sampling (IS) ratio used in off-policy policy-gradient estimation. Existing methods face a fundamental bias-variance dilemma: token-level IS ratios, as adopted by PPO (Schulman et al., 2017) and GRPO (Shao et al., 2024), introduce bias by ignoring prefix state distribution mismatch; full sequence ratios provide exact trajectory-level correction but suffer from high variance due to the multiplicative accumulation of per-token ratios, while GSPO (Zheng et al., 2025) improves numerical stability via length normalization at the cost of deviating from the exact full-sequence IS correction. In this work, we identify the cumulative token IS ratio, the product of per-token ratios up to position tt, as a theoretically principled solution to this dilemma. We prove that, under the token-level policy-gradient formulation, this ratio provides an unbiased prefix correction for each token-level gradient term and has strictly lower variance than the full sequence ratio. Building on this insight, we propose CTPO (Cumulative Token Policy Optimization), which combines the cumulative token IS ratio with position-adaptive clipping that scales log-space clip bounds according to the natural t\sqrt{t} growth of the cumulative log-ratio. This yields more consistent regularization across token positions. We implement and evaluate CTPO in the tool-integrated reasoning setting on several challenging mathematical reasoning benchmarks, achieving the best average performance across both model scales compared with strong GRPO and GSPO baselines. Code will be available at https://github.com/horizon-llm/CTPO.

Yuheng Zhang, Chenlu Ye, Shuowei Jin +4

UIUC · University of Michigan · Amazon

RL for Language Model ReasoningImportance SamplingPolicy Gradient MethodsRL Post-TrainingOff-Policy RL
AskPDF
Oct 6, 2026

On KL-Regularized Policy Optimization

Asynchronous reinforcement learning (RL) for large language model (LLM) agents trains one policy on trajectories generated by another: rollouts come from stale checkpoints, and the inference engine's probabilities differ from the trainer's even at identical parameters. Standard remedies either clip importance ratios, which biases the update, or, as in GRPO, sample a group of responses per prompt, which is costly when episodes are long. We propose KL-Regularized Policy Optimization (KLPO), a framework that anchors the KL regularizer at the sampler. The regularized improvement step then has a closed-form Gibbs solution, and KLPO fits its log-ratio optimality condition by least squares on the sampler's own trajectories, so the sampler probability enters through a log-ratio and no importance weights are needed. Profiling out the regression intercept replaces the intractable log-partition function with the signal's sampler mean plus a sampler-to-trainer KL divergence. For token-level policy mirror descent targets, we show that the resulting gradient can be computed from terminal returns without a critic, via sampler-centered scores or a single trajectory residual, even under stochastic tool outputs. We further prove that independent Monte Carlo estimates of the KL term keep these gradients unbiased, derive the exact KL gap of cheaper top-KK and binary approximations, and show that SPPO, GPO, REBEL, and BPO arise as special cases of KLPO. The result is a critic-free update that uses one rollout per prompt and requires neither a learned normalizer nor a group of responses.

Yifan Zhang

Princeton University

RL for LLM AgentsKL-Regularized RLOff-Policy RLAsynchronous RLCritic-Free RL
AskPDF