On the Design of KL-Regularized Policy Gradient Algorithms for LLM Reasoning

TL;DR

This paper introduces a unified framework called RPG for designing KL-regularized policy gradient algorithms that improve reasoning in large language models, ensuring stable, off-policy training with better accuracy on reasoning benchmarks.

Contribution

It provides a unified derivation of KL variants, corrects importance-weighting mismatches, and introduces RPG-Style Clip for stable off-policy training of LLMs.

Findings

RPG-REINFORCE with RPG-Style Clip improves accuracy by up to 6 percentage points on reasoning benchmarks.

Achieves 52% accuracy on AIME25 with 8K context length, surpassing some existing models.

Demonstrates stability and scalability of the proposed RL algorithm for LLM reasoning.

Abstract

Policy gradient algorithms have been successfully applied to enhance the reasoning capabilities of large language models (LLMs). KL regularization is ubiquitous, yet the design surface, choice of KL direction (forward vs. reverse), normalization (normalized vs. unnormalized), and estimator ($k_1/k_2/k_3$), is scattered across the literature and often intertwined with off-policy estimation. We ask a focused question: under the off-policy setting, what weighting is required for each KL variant so that the surrogate we optimize yields the exact gradient of the intended KL-regularized objective? We answer this with a compact, unified derivation we call the Regularized Policy Gradient (RPG) view. RPG (i) unifies normalized and unnormalized KL variants and shows that the widely-used $k_3$ penalty is exactly the unnormalized KL; (ii) specifies conditions under which REINFORCE-style losses with stop-gradient are gradient-equivalent to fully differentiable surrogates; (iii) identifies and corrects an off-policy importance-weighting mismatch in GRPO's KL term; and (iv) introduces RPG-Style Clip, a clipped-importance-sampling step within RPG-REINFORCE that enables stable, off-policy policy-gradient training at scale. On mathematical reasoning benchmarks (AIME24, AIME25), RPG-REINFORCE with RPG-Style Clip improves accuracy by up to $+6$ absolute percentage points over DAPO. We extend our experiments to 8K context length, and RPG-REINFORCE with RPG-Style Clip achieves 52% accuracy on AIME25, surpassing the official Qwen3-4B-Instruct model (47%). Notably, RPG is a stable and scalable RL algorithm for LLM reasoning, realized via (a) a KL-correct objective, (b) clipped importance sampling, and (c) an iterative reference-policy update scheme. Project Page: https://github.com/complex-reasoning/RPG.