KL-Regularized Policy Gradient

Updated 24 November 2025

KL-Regularized Policy Gradient is a reinforcement learning method that integrates a KL divergence penalty to a reference policy to stabilize policy updates.
It leverages both forward and reverse KL formulations to control multimodal output, highlighting the role of regularization strength and reward-reference scaling.
Innovations like MARA enhance outcome diversity and mitigate mode collapse, benefiting applications from language modeling to molecule design.

KL-regularized policy gradient methods are a central class of algorithms that augment the standard reinforcement learning (RL) policy objective with a Kullback–Leibler (KL) divergence penalty to a reference policy. This regularization is used pervasively in deep RL, fine-tuning of LLMs, preference optimization for diffusion policies, and RL from human feedback. KL-regularization stabilizes policy updates, can provide trust-region effects, shapes the coverage of the policy’s distribution over output trajectories, and allows for policy customization. The characterization of how the choice between forward KL (FKL) and reverse KL (RKL), regularization strength, and reward-reference scaling influences learning dynamics is an area of active research, with recent work rigorously refuting many traditional intuitions about the “mode-seeking” and “mass-covering” properties of FKL and RKL in a policy gradient setting (GX-Chen et al., 23 Oct 2025).

1. Mathematical Formulation: KL-Regularized Objectives

KL-regularized reinforcement learning typically seeks to optimize the expected reward of a policy $\pi(y)$ over outputs $y$ while penalizing divergence to a fixed reference policy $\pi_{\text{ref}}(y)$ :

Reverse KL (RKL):

$J_{\text{rev}}(\pi) = \mathbb{E}_{\pi(y)}[R(y)] - \beta D_{\text{KL}}(\pi\,\|\; \pi_{\text{ref}})$

$= \sum_y \pi(y) R(y) - \beta \sum_y \pi(y) [\log \pi(y) - \log \pi_{\text{ref}}(y)]$

Forward KL (FKL):

$J_{\text{fwd}}(\pi) = \mathbb{E}_{\pi(y)}[R(y)] - \beta D_{\text{KL}}(\pi_{\text{ref}}\,\|\; \pi)$

$= \sum_y \pi(y) R(y) - \beta \sum_y \pi_{\text{ref}}(y) [\log \pi_{\text{ref}}(y) - \log \pi(y)]$

Here $\beta>0$ determines the regularization strength. State-dependent (per-step) KL penalties are also implemented in sequential and RLHF settings (Pan et al., 2023, Zhang et al., 23 May 2025). The policy gradient of such objectives introduces an additional term to the standard advantage-based update, effectively incorporating $\log \pi_{\text{ref}}$ (or $\log \pi$ ) terms depending on the chosen divergence.

2. Analytical Optima and Determinants of Mode Coverage

Analytical Solution

Assuming unconstrained distributions, the KL-regularized objective yields closed-form targets:

RKL optimum:

$\pi^*(y) \propto \pi_{\text{ref}}(y) \cdot \exp(R(y)/\beta)$

FKL optimum:

$\pi^*(y) = \frac{\beta \pi_{\text{ref}}(y)}{\Lambda - R(y)}$

for normalizing constant $\Lambda>\max_y R(y)$ .

Determinants of Mode Coverage

Contrary to the classical belief that RKL is “mode-seeking” and FKL “mass-covering,” mode diversity is not an intrinsic property of the KL direction. Instead, mode inclusion and balance in the optimal policy are governed by:

The regularization coefficient $\beta$ relative to reward gaps $\Delta R$ .
Relative log-mass differences between modes under $\pi_{\text{ref}}$ . Small $\beta$ yields exponential preference for higher reward—producing “mode collapse.” If multiple outputs achieve equal reward, modes with higher $\pi_{\text{ref}}$ mass dominate, regardless of KL direction. Thus, the choice of FKL or RKL defines the mathematical form of the target, but not its fundamental capacity for multimodal support (GX-Chen et al., 23 Oct 2025).

3. Algorithmic Instantiations and Extensions

Standard Policy Gradient Update

Given the RL objective $J(\pi_\theta)$ , the policy-gradient estimator under KL regularization is:

$\nabla_\theta J(\pi_\theta) = \mathbb{E}_{\pi_\theta} \left[ \nabla_\theta \log \pi_\theta(a|s) \big( A^\pi(s,a) - \beta (\log \pi_\theta(a|s) - \log \pi_{\text{ref}}(a|s)) \big) \right]$

where $A^\pi(s,a)$ is an advantage function. This directly implements the KL-penalty as a form of reward augmentation (Wang et al., 14 Mar 2025).

Mode-Anchored Reward Augmentation (MARA)

MARA addresses mode collapse by explicitly flattening the reward among all $\tau$ -good modes (where $R(y)\geq\tau$ ). For any $y$ with $R(y)\geq\tau$ :

$\bar{R}(y) = R(z) + \beta [\log \pi_{\text{ref}}(z) - \log \pi_{\text{ref}}(y)]$

with $z$ being the $\max_{\pi_{\text{ref}}}$ anchor in the good set. This produces a uniform distribution over high-reward modes under the regularized optimum—provably correcting the exponential mass imbalance otherwise induced by small $\beta$ (GX-Chen et al., 23 Oct 2025).

Pseudocode for Policy Gradient with MARA

Sample {y_i} ~ π_θ
Find anchor z = argmax_{i: R(y_i) ≥ τ} π_ref(y_i)
For each y_i:
    if R(y_i) ≥ τ:
        bar_r_i = R(z) + β * (log π_ref(z) - log π_ref(y_i))
    else:
        bar_r_i = R(y_i)
Policy gradient update using {bar_r_i} as rewards

Implementation is a trivial substitution into RL code, yielding strictly improved mode coverage and solution diversity (GX-Chen et al., 23 Oct 2025).

4. Empirical Findings and Benchmarks

Experiments on language modeling, creative QA, and chemical design reveal:

Toy LLM tasks: Without MARA, both FKL and RKL converge to the highest- $\pi_{\text{ref}}$ answer even when ground-truth is ambiguous. MARA restores full uniform valid-answer coverage.
Creative QA ("NoveltyBench"): MARA (both FKL/RKL) outperforms GRPO and RLOO both on reward and diversity metrics, including n-gram entropy and mean distinct functional classes.
Chemical LM molecule design: On SYNTH and ALL-AMIDE, MARA boosts both the yield of high-scoring molecules and diversity metrics over standard RL and REINVENT baselines, with substantial improvements in sample efficiency.

The table below summarizes selected creative QA results from the main text of (GX-Chen et al., 23 Oct 2025):

Algorithm	Out-dist Reward ↑	N-grams EAD ↑	MeanDistinct ↑
Base model	1.166 ± .076	0.413 ± .015	4.01 ± .25
MARA (rev-KL)	1.451 ± .103	0.543 ± .014	4.14 ± .23
MARA (fwd-KL)	1.604 ± .113	0.568 ± .012	4.62 ± .26

5. Theoretical and Practical Implications

Refutation of Classic FKL/RKL Heuristics

The dichotomy of “mode-seeking” vs. “mass-covering” for RKL/FKL is not generically valid in KL-regularized policy gradient RL. Both forms define target distributions whose effective mode support is a deterministic function of the reward gaps, reference log-probabilities, and $\beta$ (GX-Chen et al., 23 Oct 2025).
For small $\beta$ , both FKL and RKL induce severe mode collapse regardless of the initial multimodality of rewards or $\pi_{\text{ref}}$ .

Algorithmic Guidance

The regularization strength $\beta$ must be tuned relative to both the scale of reward differences and log-masses under $\pi_{\text{ref}}$ ; naive reduction leads to loss of diversity.
Reward-anchoring (MARA) provides a simple, theoretically justified, and empirically validated mechanism to induce robust multimodality without external diversity signals or complex reward engineering.
MARA requires only minor code modifications and yields strict Pareto improvements in mode entropy vs. average reward.

Broader Context

The findings challenge widespread heuristics motivating popular methods in LP-RLHF and LLM post-training, necessitating a reevaluation of the role and tuning of KL-regularized objectives.
The formalism and algorithms are relevant for any domain where mode preservation and solution diversity under partial reward supervision are critical.

6. Limitations and Future Directions

The analysis in (GX-Chen et al., 23 Oct 2025) assumes access to a full evaluation of $R(y)$ and $\pi_{\text{ref}}(y)$ for minibatch sampled outputs, which may be computationally costly or impractical in ultra-large output spaces.
Extensions to hierarchical, latent-variable, or adaptive reference policies, or settings with highly structured output spaces, require further investigation.
The development of mechanistically grounded diversity metrics to inform $\tau$ selection and further augment reward-anchoring strategies remains open for future work.

In summary, KL-regularized policy gradient is a principled approach to controlling both the exploration/exploitation tradeoff and adherence to prior knowledge in RL and structured sequence modeling. Recent advances show that neither forward nor reverse KL uniquely ensures multimodality or prevents mode collapse—only reward-reference scaling and direct modification of the reward or reference mass can guarantee broad solution coverage. Reward-anchored regularization such as MARA represents a robust, general mechanism with minimal implementation burden, yielding improved outcome quality and diversity across reinforcement learning domains (GX-Chen et al., 23 Oct 2025).