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

Yifan Zhang; Yifeng Liu; Huizhuo Yuan; Quanquan Gu; Andrew C Yao

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 (k1/k2/k3) — 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 k3 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.

Regularized Policy Gradient

We derive policy gradients and corresponding surrogate losses for Forward/Reverse KL, in normalized (KL) and unnormalized (UKL) forms, under off-policy sampling with importance weights.
We give both fully differentiable surrogates and REINFORCE-style losses (with stop-gradient) and prove their gradient-equivalence to the intended regularized objective.
We introduce RPG-Style Clip, a clipped-importance-weighted REINFORCE estimator that substantially improves stability and variance control while preserving the RPG gradients.
We reveal the equality between the k3 estimator and unnormalized KL, and show that GRPO’s KL penalty omits an essential importance weight under off-policy sampling; we provide a corrected estimator and loss.
We present an iterative training framework that periodically updates the reference model to satisfy KL constraints while allowing the policy to depart meaningfully from the initial checkpoint.
On math reasoning, RPG-REINFORCE (with RPG-Style Clip) yields stable and scalable training and outperforms DAPO by up to +6 absolute points on AIME24/25.

Experimental Results

AIME24/AIME25 results, 4K context — Combined performance on AIME24 and AIME25, showing “Last” and “Best” scores for 4K context length. Best in each column is **bold**, second best underlined.

AIME24/AIME25 results, 2K context — Combined performance on AIME24 and AIME25, “Last” and “Best” scores for 2K context length.

Training dynamics, 4K context — Training dynamics and benchmark performance for RPG and REINFORCE-style RPG vs. baselines (GRPO, DAPO) at 4K context length.

Training dynamics, 2K context — Training dynamics and benchmark performance for RPG and REINFORCE-style RPG vs. baselines (GRPO, DAPO) at 2K context length.

Fully Differentiable Surrogate Loss Functions

Regularized policy gradients with fully differentiable surrogate losses

REINFORCE-Style Regularized Policy Gradients

Citation

If you find RPG useful, please cite:

@article{zhang2025design,
    title   = {On the Design of KL-Regularized Policy Gradient Algorithms for LLM Reasoning},
    author  = {Zhang, Yifan and Liu, Yifeng and Yuan, Huizhuo and Gu, Quanquan and Yao, Andrew C},
    journal = {arXiv preprint arXiv:2505.17508},
    year    = {2025}
}