KL-Regularized Reinforcement Learning

• KL-Regularized RL is a class of reinforcement learning algorithms that integrate a KL divergence penalty to balance reward maximization with adherence to a trusted baseline policy.
• It stabilizes learning by constraining policy updates through a tunable regularization coefficient, improving safety and sample efficiency in applications like LLM fine-tuning and offline RL.
• Various algorithmic implementations, including policy iteration, actor-critic, and Q-learning, benefit from robust theoretical guarantees and enhanced performance in safe and transferable policy learning.

Kullback-Leibler (KL)-Regularized Reinforcement Learning (KL-RL) is a broad class of reinforcement learning algorithms that augment the standard RL objective with a penalty term involving the Kullback-Leibler divergence between the learned policy and a reference (or prior) policy. KL regularization is used to promote safe, stable, and sample-efficient policy optimization across deep RL, offline RL, policy transfer, and LLM fine-tuning. The core principle is to constrain policy updates so as to balance behavioral improvement (expected reward) with distributional proximity to a known or trusted baseline, thereby stabilizing learning and controlling policy drift, exploitation, or unsafe exploration.

1. Mathematical Foundations and Principal Objectives

KL-RL extends standard Markov Decision Process (MDP) objectives by introducing a KL penalty weighted by a regularization coefficient. In the episodic/horizon-$H$ setting, with agent policy $\pi$, reference policy $\pi_0$, reward $r(s, a)$, and discount $\gamma$, the general KL-regularized objective is

$$J_{\mathrm{KL}}(\pi) = \mathbb{E}_{\tau \sim \pi}\bigg[\sum_{t=1}^H \gamma^t r(s_t, a_t)\bigg] - \eta \cdot \mathbb{E}_{\tau \sim \pi}\left[\sum_{t=1}^H \mathrm{KL}\big(\pi(\cdot \mid s_t) \| \pi_0(\cdot \mid s_t)\big)\right],$$

where $\eta$ controls the strength of the regularization. The form generalizes to infinite-horizon and entropy-regularized settings, and encompasses both online and offline RL, as well as policy iteration and policy gradient methods (Cai et al., 2022, Tirumala et al., 2019, Zhao et al., 2024).

In the context of LLM fine-tuning, the KL penalty is crucial to preventing distributional collapse and aligning the output distribution with safe baselines, often implemented as reverse KL between the fine-tuned and base LLM (Shah et al., 26 Dec 2025, Korbak et al., 2022). The choice of KL direction—reverse ($\mathrm{KL}(\pi \| \pi_0)$, “mode-seeking”) versus forward ($\mathrm{KL}(\pi_0 \| \pi)$, “mass-covering”)—directly affects the learned policy's characteristics, coverage, and theoretical properties (GX-Chen et al., 23 Oct 2025).

2. Algorithms and Mechanistic Design

KL-RL can be instantiated via value-based, policy-gradient, and actor-critic schemes. Canonical update rules leverage the analytically tractable Gibbs (Boltzmann) policy form:

$$\pi^*(a \mid s) = \frac{\pi_0(a \mid s)\exp(Q^*(s, a) / \eta)}{Z(s)}$$

where $\pi$0 is the regularized action-value function and $\pi$1 is a normalization factor (Tirumala et al., 2019, Zhao et al., 11 Feb 2025, Hong et al., 4 Jun 2026). The reference policy can be fixed (pretrained, demonstrator, or behavior policy in offline RL) or learned jointly to induce inductive biases, e.g., hierarchical structure or information bottlenecks (Galashov et al., 2019, Tirumala et al., 2019).

KL regularization appears in multiple algorithmic motifs:

• KL-regularized Policy Iteration: Alternates regularized policy evaluation and greedy improvement to the current Gibbs-optimal policy (Vieillard et al., 2020, Zhu et al., 2021, Kitamura et al., 2021).
• KL-regularized Q-learning: Replaces max in the Bellman target with a KL-regularized operator, enabling “soft” target backups and regularization toward a behavior or safe-support policy (Lim et al., 28 Apr 2026, Cai et al., 2022).
• KL-regularized Actor-Critic: Adds the KL term in the actor (policy) update or in the loss for policy-gradient-based optimization, often realized via PPO- or SAC-style objectives (Cai et al., 2022, Wang et al., 14 Mar 2025).
• Reverse-KL for Mode-Seeking BC: In offline RL with multi-behavior datasets, reverse KL regularization enforces “mode-seeking” updates, avoiding out-of-distribution actions by concentrating the policy on high-density support regions of $\pi$2 (Cai et al., 2022).

Recent algorithmic advances focus on:

• Exponential Moving Average (EMA) reference policies for stability in non-stationary RL (Zhang et al., 4 Feb 2026).
• Dynamic KL coefficients tuned per iteration to adaptively balance exploration/stability (Kitamura et al., 2021).
• Mode-Anchored Reward Augmentation (MARA) to actively enforce uniform coverage over all optimal modes when KL regularization would naively collapse to the highest-prior mode (GX-Chen et al., 23 Oct 2025).
• Top-$\pi$3 unbiased KL estimators for scalable, memory-efficient training in high-dimensional action spaces (Zhang et al., 4 Feb 2026, Shah et al., 26 Dec 2025).

Empirical recipes for estimator choice, placement in the computational graph, and stability are critical for LLM-RL training (Shah et al., 26 Dec 2025).

3. Theoretical Properties and Guarantees

KL regularization imparts strong convexity to the optimization landscape in the policy distribution, yielding several theoretical benefits:

• Sample Complexity Improvements: For KL-regularized contextual bandits and RLHF, reverse KL regularization yields a sharp $\pi$4 sample complexity rate for $\pi$5-optimality, under reasonable coverage assumptions on the reference policy (Zhao et al., 2024). This is a substantial improvement over the $\pi$6 rate in standard unregularized settings.
• Logarithmic Regret and Benign Exploration: Optimism-driven KL-regularized bandit and RL algorithms achieve $\pi$7 regret, where $\pi$8 is the reward class size and $\pi$9 the function class complexity (Zhao et al., 11 Feb 2025).
• Error Averaging and Robustness: KL regularization causes filtered averaging of value estimation errors (“dual averaging”), resulting in performance bounds linear in the planning horizon and smoothing of error propagation (Vieillard et al., 2020, Kitamura et al., 2021).
• Function Approximation and Misspecification: With general function approximators, high-probability KL-regret bounds degrade gracefully with misspecification, with explicit additive terms entering the regret (Hong et al., 4 Jun 2026).

Table: Sample Complexity Dependence in KL-Regularized RL

Regime	Required Samples	Coverage Dependence
Data coverage (additive D^2)	$\pi_0$0	Additive in $\pi_0$1 (coverage constant)
Local KL-ball (mult. C_ρ)	$\pi_0$2	Multiplicative in $\pi_0$3
No coverage	–	No guarantee

(Zhao et al., 2024, Zhao et al., 11 Feb 2025, Hong et al., 4 Jun 2026)

4. Structural Insights and Inductive Bias

KL-RL provides a route to hierarchical, modular, and reusable policies:

• Default/Reference Policy Learning: Learning the “prior” policy together with the agent policy—under capacity or information constraints—yields inductive biases (temporal abstraction, motor primitives, goal-agnostic skills) leading to faster transfer and more efficient reuse (Tirumala et al., 2019, Galashov et al., 2019).
• Latent Variable Hierarchies: Augmenting both the agent and default policies with latent factors (high-level and low-level) decouples task and control knowledge; KL penalties at each level allow modular recombination across tasks and morphologies (Tirumala et al., 2019).
• Information Asymmetry: Restricting observations available to the default policy forces learning of “core” behaviors, while the main policy specializes via the regularization gradient; this parallels information bottleneck or variational EM methods (Galashov et al., 2019).

Hierarchical KL-regularized methods show dramatic gains in multi-task learning and body/skill transfer for complex control tasks (Tirumala et al., 2019).

5. Practical Applications and Empirical Performance

KL-RL permeates contemporary deep RL and RLHF practice:

• Offline RL from Mixed Datasets: TD3+RKL employs per-state weighted reverse KL-based behavior cloning in MuJoCo tasks, outperforming TD3+BC and forward KL regularized approaches on both standard and mixed-expert datasets by up to 25.5% normalized score (Cai et al., 2022).
• Safe RL and Exploration Constraints: KL regularization toward a behavior policy with support restricted to safe actions enables “Safe-Support Q-Learning,” guaranteeing no unsafe state visitation and stable, calibrated Q-function learning (Lim et al., 28 Apr 2026). This is distinct from entropy-based methods such as SAC, as the support constraint aligns with safety enforcement.
• Imitation and Behavioral Cloning: KL-regularized actor-critic and Q-learning algorithms benefit from non-parametric, uncertainty-calibrated (e.g., GP) behavioral policies to avoid gradient explosions and unreliable policy evaluation seen with overconfident parametric networks (Rudner et al., 2022).
• LLMs and Language Agents: Regularizing RL fine-tuning of LLMs toward the pretrained base model is essential for preventing distribution collapse. The exact estimator and placement of the KL term are critical: “K1 in reward” is unbiased and empirically superior, both for stability and downstream performance on in-domain and out-of-domain evaluations (Shah et al., 26 Dec 2025, Korbak et al., 2022).
• Fine-Tuning and Policy Customization: Residual Policy Gradient (RPG) demonstrates that KL-RL is equivalent to maximum-entropy RL on an augmented reward, supporting customization while preserving desirable priors (Wang et al., 14 Mar 2025).
• Diversity, Mode Collapse, and MARA: Contrary to naive “mode-seeking/mass-covering” dichotomies from variational inference, actual mode coverage in KL-RL depends on regularization strength and relative reference/prior scoring. For standard hyperparameterizations, RL with reverse or forward KL is prone to mode collapse. Augmented reward (MARA) ensures uniform coverage of high-quality modes and superior out-of-distribution diversity (GX-Chen et al., 23 Oct 2025).

6. Limitations, Pathologies, and Safety Considerations

KL-RL is not without pitfalls, especially when the base policy is misspecified or the regularization or support is improperly set:

• Support Mismatch and KL Singularity: Classical KL regularization is infinite if the policy visits actions with zero probability under the reference, leading to degenerate or ill-posed control, particularly in low-noise and deterministic dynamics. State-space-aware divergences (Wasserstein-KL) resolve singularities and yield well-posed LQR controllers (Stein et al., 2 Feb 2026).
• Pathological Instabilities: Behavioral reference policies derived from poorly-calibrated or low-variance expert fits can result in exploding KL gradients and training collapse. Non-parametric policies or careful variance regularization are essential remedies (Rudner et al., 2022).
• Bayesian Predictive Base Policy Vulnerability: KL-regularizing to a “Bayesian predictor” of a trusted demonstrator is not sufficient to block unsafe RL behaviors in novel situations—a simple, high-return novel action may incur minimal KL cost once the agent enters a new regime. The “don’t do anything I mightn’t do” principle, realized via pessimistic Bayes imitation, is suggested as an alternative safety anchor (Cohen et al., 2024).
• Estimator Bias in LLM Fine-Tuning: Many open-source RL for LLMs incorrectly implement the KL regularizer, resulting in biased gradients and unstable or sub-optimal training. Empirically, unbiased estimator placements (naïve in-reward) are mandatory for both stability and performance (Shah et al., 26 Dec 2025).

7. Future Directions and Open Questions

KL-regularized RL continues to be a focal point for advancing statistical and algorithmic efficiency, robust transfer, safety, and scalability:

• Integrating alternative f-divergences or transport-based regularizers (Wasserstein-KL, Kalman-Wasserstein-KL) to circumvent singularities and maintain well-posedness in control (Stein et al., 2 Feb 2026).
• Hierarchical, modular, and structured regularization in multi-agent, multi-task, and lifelong-learning domains (Tirumala et al., 2019, Galashov et al., 2019).
• Automatic and dynamic tuning of KL coefficients for robustness to non-stationarity, reward misspecification, or on-the-fly adaptation (Kitamura et al., 2021).
• Deeper theoretical understanding of mode coverage and collapse under practical constraints (reward scale, prior support, batch estimation) (GX-Chen et al., 23 Oct 2025).
• Safe RL with rigorous support constraints and explicit divergence-based certification, beyond “soft” policy regularization (Lim et al., 28 Apr 2026, Cohen et al., 2024).

KL-regularized RL, rooted in both control theory and probabilistic inference, continues to unify and advance the design of scalable, reliable, and safe autonomous agents across a spectrum of real-world domains.