Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Well-Posed KL-Regularized Control via Wasserstein and Kalman-Wasserstein KL Divergences

Viktor Stein, Adwait Datar, Nihat Ay

TL;DR

The paper addresses the instability of KL regularization under support mismatch and low-noise limits by formulating state-space-aware KL divergences through Wasserstein-based and Kalman–Wasserstein geometries on the density manifold. It develops explicit, closed-form divergences between elliptic distributions, establishes gradient-flow dynamics and dual geodesics under these geometries, and shows how KWKL regularization yields well-posed, interpretable penalties in LQR problems, even as process noise vanishes. The authors verify the approach with two canonical, fully tractable examples—the double integrator and cart-pole—demonstrating that WKL and KWKL regularizers outperform classical KL in low-noise regimes, while smoothly interpolating between KL and transport behavior via a tunable parameter. The work lays groundwork for integrating these divergences into RL and model-based control, offering robust alternatives to KL and insights into regulator design under uncertainty and near-deterministic dynamics.

Abstract

Kullback-Leibler divergence (KL) regularization is widely used in reinforcement learning, but it becomes infinite under support mismatch and can degenerate in low-noise limits. Utilizing a unified information-geometric framework, we introduce (Kalman)-Wasserstein-based KL analogues by replacing the Fisher-Rao geometry in the dynamical formulation of the KL with transport-based geometries, and we derive closed-form values for common distribution families. These divergences remain finite under support mismatch and yield a geometric interpretation of regularization heuristics used in Kalman ensemble methods. We demonstrate the utility of these divergences in KL-regularized optimal control. In the fully tractable setting of linear time-invariant systems with Gaussian process noise, the classical KL reduces to a quadratic control penalty that becomes singular as process noise vanishes. Our variants remove this singularity, yielding well-posed problems. On a double integrator and a cart-pole example, the resulting controls outperform KL-based regularization.

Well-Posed KL-Regularized Control via Wasserstein and Kalman-Wasserstein KL Divergences

TL;DR

The paper addresses the instability of KL regularization under support mismatch and low-noise limits by formulating state-space-aware KL divergences through Wasserstein-based and Kalman–Wasserstein geometries on the density manifold. It develops explicit, closed-form divergences between elliptic distributions, establishes gradient-flow dynamics and dual geodesics under these geometries, and shows how KWKL regularization yields well-posed, interpretable penalties in LQR problems, even as process noise vanishes. The authors verify the approach with two canonical, fully tractable examples—the double integrator and cart-pole—demonstrating that WKL and KWKL regularizers outperform classical KL in low-noise regimes, while smoothly interpolating between KL and transport behavior via a tunable parameter. The work lays groundwork for integrating these divergences into RL and model-based control, offering robust alternatives to KL and insights into regulator design under uncertainty and near-deterministic dynamics.

Abstract

Kullback-Leibler divergence (KL) regularization is widely used in reinforcement learning, but it becomes infinite under support mismatch and can degenerate in low-noise limits. Utilizing a unified information-geometric framework, we introduce (Kalman)-Wasserstein-based KL analogues by replacing the Fisher-Rao geometry in the dynamical formulation of the KL with transport-based geometries, and we derive closed-form values for common distribution families. These divergences remain finite under support mismatch and yield a geometric interpretation of regularization heuristics used in Kalman ensemble methods. We demonstrate the utility of these divergences in KL-regularized optimal control. In the fully tractable setting of linear time-invariant systems with Gaussian process noise, the classical KL reduces to a quadratic control penalty that becomes singular as process noise vanishes. Our variants remove this singularity, yielding well-posed problems. On a double integrator and a cart-pole example, the resulting controls outperform KL-based regularization.
Paper Structure (33 sections, 15 theorems, 109 equations, 11 figures)

This paper contains 33 sections, 15 theorems, 109 equations, 11 figures.

Table of Contents

  1. Introduction
  2. Prior work and contributions
  3. Notation
  4. The Riemannian structure of the density manifold
  5. Metric gradient flows on the density manifold
  6. Dual geodesics on the density manifold
  7. state-space-aware KL-divergences
  8. Explicit expressions for the state-space-aware KL divergence
  9. The state-space-aware KL divergence between elliptic distributions
  10. Elliptic distributions
  11. Quadratic potentials
  12. Optimal Control Problems with Divergence-Based Regularization
  13. Example 1: Double integrator
  14. Example 2: Cart-pole system
  15. Conclusions and future work
  16. ...and 18 more sections

Key Result

proposition 2.7

For $f \in \mathop{\mathrm{\mathcal{C}}}\nolimits^{\infty}(M)$, we have This property uniquely characterizes these metrics.

Figures (11)

  • Figure 1: For $\mu \sim \mathop{\mathrm{\mathcal{N}}}\nolimits(m_0, \sigma^2)$ and $\nu \sim \mathop{\mathrm{\mathcal{N}}}\nolimits(m_1, \sigma^2)$, we compare $D^{\mathop{\mathrm{KW}}\nolimits}(\mu \mid \nu) = \frac{1}{2(\sigma^2 + \lambda)}(\Delta m)^2$ with $\mathop{\mathrm{KL}}\nolimits(\mu \mid \nu) = \frac{1}{2 \sigma^2}(\Delta m)^2$ and $\mathop{\mathrm{WKL}}\nolimits(\mu \mid \nu) = \frac{1}{2}(\Delta m)^2$, where $\lambda = 1$ and $(\Delta m)^2 = \frac{1}{4}$.
  • Figure 2: Optimal feedback gains for $\mathop{\mathrm{KL}}\nolimits-$, $\mathop{\mathrm{WKL}}\nolimits-$, and $\mathop{\mathrm{KW}}\nolimits-$ regularized LQR as function $\rho$. Each feedback gain $F\in \mathbb{R}^{1\times 2}$ shown component-wise; both entries are shown. KL gains shrink to zero as noise disappears, WKL gains are constant because they do not depend on $\rho$, and KW gains smoothly interpolate between the two.
  • Figure 3: Closed-loop trajectories of the double integrator under KL-, WKL-, and KW-regularized controllers for $\rho \in \{10^{-1},10^{-2},10^{-3},10^{-4}\}$. For small $\rho$, KL yields weak feedback and large oscillations, whereas WKL- and KW-regularized controllers lead to reduced oscillations. See Figure \ref{['fig:traj_all_vary_lambda']} for other values of $\lambda$.
  • Figure 4: Closed-loop spectral radius as a function of noise $\rho$ for different values of parameter $\lambda$. KL-regularized controllers approach the unit circle as $\rho\rightarrow 0$, indicating near-instability, while WKL- and KW-regularized controllers preserve maintain spectral radii well below one across noise regimes. Varying $\lambda$ interpolates between KL-like and WKL-like behavior.
  • Figure 5: Cart-pole system
  • ...and 6 more figures

Theorems & Definitions (56)

  • Definition 2.1: Inertia operator
  • Definition 2.2: Regular Riemannian metric on $P_+^{\infty}(M)$
  • Example 2.3: Fisher-Rao metric
  • Example 2.4: Otto's Wasserstein metric with non-linear mobility
  • Definition 2.5: Linear functional derivative
  • Definition 2.6: Metric gradient flow on $P_+^{\infty}(M)$
  • proposition 2.7: Metric gradients of the expectation functional
  • proof
  • Definition 2.8: $G$-parallel transport
  • Theorem 2.9: Characterizing $G$-dual geodesics
  • ...and 46 more