KL-Convex Policy Constraint in RL

• KL-convex policy constraint is a class of methods that use KL divergence to regulate policy updates by enforcing trust regions or regularizing against a reference policy.
• These methods offer distinct formulations—hard constraints versus regularizers—that influence update dynamics, optimization landscapes, and theoretical regret guarantees in reinforcement learning.
• They are embedded within broader f‐divergence frameworks and find practical use in multi-agent settings, safe RL, and offline learning to enhance stability and performance.

Across the cited literature, the term KL-convex policy constraint is most naturally understood as a family of policy optimization mechanisms in which the Kullback-Leibler divergence is used to restrict or regularize policy change through a trust-region bound, a convex penalty in the objective, or a reference-policy regularizer. In generalized policy iteration, the policy improvement step is augmented with a trust region constraint bounding information loss, and the size of that trust region is commonly determined by the KL divergence because it “captures the notion of distance well” and “yields closed-form solutions” (Belousov et al., 2017). In approximate policy iteration and related actor-critic methods, the same idea appears either as a hard KL constraint between consecutive policies or as a KL regularizer that continuously trades off policy improvement against closeness to the previous or reference policy (Lazić et al., 2021). The resulting design choices are not interchangeable: the literature associates them with different update rules, critic forms, optimization landscapes, regret guarantees, and statistical rates.

1. Canonical formulations in policy improvement

A canonical constrained formulation uses the expected advantage objective

$$L_{\pi_k}(\pi) = E_{s \sim \mu_k,\, a \sim \pi(\cdot|s)}[\widehat A_{\pi_k}(s,a)]$$

and imposes a trust-region condition of the form

$$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$

This is the implementation choice associated in the literature with practical algorithms such as TRPO, MPO, and VMPO (Lazić et al., 2021).

The corresponding regularized formulation replaces the hard constraint by a KL term in the objective: $$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$ For the per-state mirror-descent-style update, the policy is

$$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$

with analytic form

$$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$

The distinction is explicit: the constrained variant is “stay within a trust region” around $\pi_k$, whereas the regularized variant “trade[s] off improvement and closeness” continuously in the objective (Lazić et al., 2021).

This distinction extends beyond notation. The cited work argues that it is “not merely cosmetic,” because the constraint and the regularizer induce different optimization dynamics and different regret behavior (Lazić et al., 2021).

2. KL within the broader $f$-divergence framework

KL-constrained policy improvement is a special case of a broader $f$-divergence construction. The general framework derives policy update rules for a class of $f$-divergences, with the generic solution expressed through the derivative of the convex conjugate function to $f$; the KL solution appears as one special case (Belousov et al., 2017).

Within this framework, the one-parameter family of $$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$0-divergences is used to study how the choice of divergence changes policy improvement. Previously known as well as new policy updates emerge for different values of $$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$1, and every type of policy update is paired with a compatible policy evaluation rule. Two identifications are especially prominent. First, the KL divergence yields the soft-max policy update and a log-sum-exp critic. Second, mean-squared Bellman error minimization is closely related to policy evaluation with the Pearson $$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$2-divergence penalty (Belousov et al., 2017).

The paper also reports asymptotic analysis for different $$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$3 values and demonstrations on a multi-armed bandit problem and on common standard reinforcement learning problems. Taken together, these results place the KL construction inside a larger convex-analytic design space rather than treating it as the only possible trust-region geometry (Belousov et al., 2017).

3. Optimization issues, instability, and landscape effects

A central controversy concerns whether a hard KL constraint is preferable to a KL regularizer. In a two-armed stochastic bandit with rewards

$$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$4

and noisy samples $$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$5, $$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$6, the constrained TRPO-style update can be solved exactly: $$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$7 The reported pathology is that when the empirical estimate has the wrong sign, $$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$8, the update moves in the wrong direction until it hits the KL boundary. The paper proves that for $$\max_\theta \; L_{\pi_k}(\pi_\theta) \quad \text{s.t.} \quad E_{x \sim \mu_k}\!\left[D_{KL}\!\left(\pi_k(\cdot|x)\parallel \pi_\theta(\cdot|x)\right)\right] \le \eta.$$9, there exists a bandit instance $$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$0 such that

$$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$1

so the expected probability of choosing the optimal arm may not increase even near optimality. It also proves a linear lower bound on expected regret,

$$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$2

which is linear in $$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$3 (Lazić et al., 2021).

By contrast, the regularized mirror-descent update is reported to average noise over iterations and to achieve sublinear regret: the paper cites $$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$4 from the known analysis of Politex and notes that this can be improved to $$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$5 if all data are reused to estimate $$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$6 (Lazić et al., 2021).

The same paper links these differences to optimization geometry under function approximation. The expected-advantage objective optimized by constrained methods can be highly non-convex with suboptimal plateaus and even exponentially many local minima in the worst case for softmax policies. Adding KL regularization can make the landscape “more connected,” mitigates softmax saturation, and leads to surrogate objectives such as

$$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$7

and

$$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$8

that are convex in the policy activations $$\max_\theta \; L_{\pi_k}(\pi_\theta) - \eta^{-1} E_{x \sim \mu_k} \left[D_{KL}\!\left(\pi_\theta(\cdot|x)\parallel \pi_k(\cdot|x)\right)\right].$$9. Empirically, on noisy two-armed bandits, MNIST contextual bandits, and discretized-action DeepMind Control Suite tasks, the regularized methods are reported to be smoother and often faster, while CPO oscillates more and is more sensitive to $$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$0 (Lazić et al., 2021).

4. Penalization, maximal KL, and exact trust-region enforcement

A separate line of work studies how to obtain an exact trust-region guarantee without solving the full nonlinear constrained optimization problem. The starting point is the TRPO-style objective

$$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$1

where $$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$2 denotes the maximum KL divergence over states, together with the PPO-style KL-regularized update

$$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$3

The proposed method, FixPO, uses a primary phase with combined loss

$$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$4

and a separate multiplier update

$$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$5

If $$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$6, the update drives $$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$7; if $$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$8, then the optimum is moved to

$$\pi_{k+1}(\cdot|x) = \arg\max_{\pi \in \Delta_A} \widehat A_{\pi_k}(x,\pi) - \eta^{-1}D_{KL}(\pi\|\pi_k(\cdot|x)),$$9

The paper characterizes KL penalization alone as “nearly sufficient”: it usually drives the policy inside the trust region, but minibatching, optimization error, and approximate $$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$0-tuning can still leave violations (Zentner et al., 2023).

The fixup phase then explicitly checks the per-state maximal-KL constraint and keeps applying KL-only updates until

$$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$1

Because the loop does not terminate until this condition holds, the paper states that the trust region is enforced between every policy update. The reported practical overhead is fewer than $$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$2 additional gradient steps, and the method is described as applicable to a variety of policy architectures and action spaces while retaining minibatching (Zentner et al., 2023).

5. Multi-agent allocation, prior-policy views, and statistical formulations

In multi-agent reinforcement learning, KL constraints need not be assigned uniformly across agents. HATRPO uses sequential updates with the same fixed KL threshold for each agent, but the cited work argues that this is conservative in heterogeneous settings. It replaces uniform per-agent thresholds by a global KL budget,

$$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$3

and proposes two allocation schemes. HATRPO-W uses a KKT-based rule with allocations of the form

$$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$4

while HATRPO-G ranks agents by the improvement-to-divergence ratio

$$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$5

Reported gains include faster convergence, improved escape from local optima, and final-performance improvements of about $$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$6 for HATRPO-G and about $$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$7 for HATRPO-W, with HATRPO-W exhibiting lower variance (Shek et al., 14 Aug 2025).

A different extension reinterprets KL regularization at the reward level. The standard KL-regularized return

$$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$8

is decomposed as

$$\pi_{k+1}(\cdot|x) \propto \exp\!\Big(\eta \sum_{i=1}^k \widehat A_{\pi_i}(x,\cdot)\Big).$$9

Under the paper’s stated assumptions, this means the KL-regularized objective is exactly optimizing a maximum-entropy policy on the MDP

$\pi_k$0

The paper stresses that this equivalence depends on the prior policy being an optimal maximum-entropy policy, on matching entropy coefficients $\pi_k$1, and on $\pi_k$2; otherwise the neat interpretation need not hold exactly (Wang et al., 14 Mar 2025).

The statistical role of KL regularization is also explicit in offline learning. In KL-regularized two-player zero-sum games with a fixed reference policy, the regularized objective induces smooth best responses and a unique Nash equilibrium, and the paper proves the first $\pi_k$3 sample-complexity bound for offline learning without pessimism (Zhang et al., 8 Apr 2026). In offline contextual bandits, the forward-KL objective

$\pi_k$4

has the closed-form optimizer

$\pi_k$5

and the paper gives the first $\pi_k$6 upper bounds in both tabular and general function approximation settings under single-policy concentrability (Zhao et al., 9 May 2026).

6. Relation to adjacent convex policy-constraint methods

A recurrent misconception is that any convex policy constraint in reinforcement learning is KL-based. The adjacent literature shows otherwise. Several works construct convex policy constraints without using KL divergence at all.

In safe reinforcement learning with nonlinear function approximation, one approach replaces the original nonconvex objective and safety constraint by convex quadratic surrogates

$\pi_k$7

$\pi_k$8

leading to a convex QCQP; the paper explicitly states that it has no KL regularizer and no KL trust-region constraint (Yu et al., 2019). A related off-policy CMDP method, SCAOPO, also uses strongly convex quadratic surrogates in parameter space and a relaxed update

$\pi_k$9

again without any explicit KL trust region (Tian et al., 2021).

Other nearby formulations constrain occupancy measures, mixed-policy measurements, or expected costs rather than policy divergence. Examples include optimization over the convex hull of base policies or occupancy measures (Banijamali et al., 2018), Euclidean distance to a convex set of long-term measurement vectors in constrained deep reinforcement learning (Cai et al., 2021), expected linear cost constraints in safe linear bandits (Afsharrad et al., 2023), and convex functionals of visitation-measure embeddings handled by Lagrangian and Fenchel duality in constrained convex MDPs (Li et al., 2024). A further source of confusion is the “Kalman constraint” in linear control policy fitting, which is an inverse-LQR optimality constraint and not a Kullback-Leibler constraint (Palan et al., 2020).

The boundary that emerges from these comparisons is precise. KL-convex policy constraints are distinguished not by convexity alone, but by the use of KL geometry to control policy movement, encode proximity to a previous or reference policy, or induce softmax-type regularized best responses. Adjacent convex-constraint methods may share trust-region or proximal motivations, yet they rely on different objects—Euclidean parameter penalties, occupancy geometry, measurement-space projections, or control-theoretic feasibility conditions—rather than KL divergence itself.