Enabling Token-wise Importance Sampling in Off-Policy RLVR with Proximal KL Regularization
AuthorsHao Bai, Jiarui Yao, Tong Zhang
AffiliationsUniversity of Illinois Urbana-Champaign
A new proximal-KL approach makes token-wise off-policy learning more stable for training language models with verifiable rewards.
Key results
Largest absolute improvement over GRPO across matched off-policy comparisons.
Proximal KL has higher micro-averaged accuracy in every comparison.
Evaluation suite includes MATH500 and AIME24.
Approximate optimizer updates to reach 50%.
Approximate optimizer updates to reach 50%.
What the paper found
The paper introduces Proximal KL, an off-policy reinforcement-learning method for verifiable rewards that separates the frozen proximal reference, which defines the improvement target, from the behavior policy that generated each response. It corrects each token with a token-wise importance ratio and assigns it bounded KL costs from future tokens, replacing ordinary PPO clipping without fitting a critic. Under coverage and local-expressiveness assumptions, the ideal update with the full, unclipped future-KL tail has the proximal optimizer as its unique stationary policy, even for stale behavior-policy mixtures; separate experiments show why practical cost clipping and finite horizons can help. On Qwen3-1.7B and Qwen3-4B, Proximal KL beats GRPO in all 7 matched off-policy comparisons across 8 math benchmarks, including MATH500 and AIME24, with up to 2.58% higher micro-averaged accuracy; on-policy aggregate accuracy is comparable. In a separate 8B periodic-refresh experiment, it reaches 50% training accuracy after about 200 updates, compared with about 640 for GRPO, and avoids collapses observed with PG and CISPO. Ablations find that token-wise correction outperforms sequence-level or unweighted alternatives, while finite future-KL windows can improve AIME24 results and reduce correction variance.
Original 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.
Read the original paperMore in Reinforcement Learning
Browse all 63 papers →LLM Value Functions for Credit Assignment in Agentic RL
Kaijie Zhu, Weixiang Yan, Omar Attia, Abbas Kazerouni, Alexander Toshev, Andrew Szot
LVRL uses an LLM’s value estimates to identify which actions help or hurt an agent, improving learning from sparse rewards.
MagicBeanStalk: Escaping the RLVR Zero Advantage Trap with Recursive Data Synthesis
Ximing Lu, Aviral Kumar, Jan Kautz, Yejin Choi, David Acuna
MagicBeanStalk recursively creates calibrated reasoning problems that help RL-trained models learn more effectively than simply scaling or filtering existing data.
On KL-Regularized Policy Optimization
Yifan Zhang
KLPO offers a cheaper, critic-free way to train asynchronous language-model agents from stale rollouts by using the sampler itself as the KL reference.