🔢 mathematics

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

This paper introduces well-posed KL-regularized control methods using Wasserstein and Kalman-Wasserstein divergences to resolve the singularities and support mismatch issues inherent in classical KL divergence, thereby enabling stable and high-performing feedback control in low-noise regimes.

Original authors: Viktor Stein, Adwait Datar, Nihat Ay

Published 2026-06-02

📖 4 min read🧠 Deep dive

CC BY 4.0

Original authors: Viktor Stein, Adwait Datar, Nihat Ay

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to walk. To do this, you compare two scenarios: how the robot actually moves when you give it a command, and how it would move if you just let it drift on its own. The difference between these two movements is a "penalty" you add to your training to keep the robot from going crazy.

For decades, scientists have used a mathematical tool called KL Divergence to measure this difference. Think of KL Divergence as a very strict, old-fashioned ruler. It works great when the robot is moving in a noisy, chaotic environment (like walking on a bumpy, windy day). But, as the paper explains, this ruler has a fatal flaw: it breaks completely when the environment becomes too quiet.

If the "noise" (the wind, the bumps) disappears and the robot is in a perfectly smooth, predictable world, the KL ruler tries to measure a distance that becomes infinitely large. It's like trying to measure the distance between two points on a map that has suddenly shrunk to a single dot; the math explodes, and the robot stops learning entirely.

The New Solution: A Flexible, "State-Aware" Ruler

The authors of this paper, Viktor Stein, Adwait Datar, and Nihat Ay, propose a new way to measure this difference. Instead of using the old, rigid KL ruler, they invented three new types of "rulers" based on different mathematical geometries:

Wasserstein Ruler: This ruler understands the physical space. It knows that moving a robot one inch to the left is a small change, even if the math says the "probability" of that move is different. It treats the robot's movement like pouring water from one cup to another, measuring the actual effort required to move the "mass."
Kalman-Wasserstein Ruler: This is a hybrid. It's like the Wasserstein ruler but with a built-in safety net. It adds a tiny bit of "imaginary noise" to keep the math from breaking, even when the real world is perfectly quiet.
Linearized Ruler: A simplified version that works well for specific, symmetrical shapes.

The Magic Analogy:
Imagine you are pushing a heavy box across the floor.

The Old Way (KL): If the floor is perfectly smooth (no friction/noise), the old ruler says, "You can't push this box! The effort required is infinite!" and you stop trying.
The New Way (Wasserstein/Kalman): These new rulers say, "Okay, the floor is smooth, but you still need to apply a little force to get it moving." They keep the problem solvable, no matter how smooth the floor gets.

What They Actually Did

The paper doesn't just talk about theory; they proved these new rulers work in three specific ways:

Fixed the Math: They showed that for common shapes of data (like the bell curves used in robotics), they can write down exact formulas for these new rulers. Crucially, these formulas never blow up when the noise level drops to zero.
Solved a Control Problem: They tested these new rulers on a classic control problem called the Linear Quadratic Regulator (LQR). This is a standard test for controlling systems like a double pendulum or a cart with a pole on top.
- When they used the old KL ruler with very little noise, the controller (the "brain" of the robot) gave up and stopped reacting. The robot became unstable and wobbled wildly.
- When they used the new Kalman-Wasserstein ruler, the controller kept working. It maintained a steady, stable grip on the robot, even in the "perfectly quiet" limit.
Real-World Examples: They tested this on a "double integrator" (a simple model of a moving object) and a "cart-pole" (a classic balancing act). In every case where the noise was low, the new method produced smoother, more stable movements than the old method.

The Bottom Line

The paper claims that by swapping out the old, fragile mathematical tool (KL Divergence) for these new, "state-space-aware" tools, we can build control systems that don't crash when the world gets too predictable.

They aren't claiming to have built a new robot or a new AI algorithm for self-driving cars yet. Instead, they have provided a better mathematical foundation (a new way to measure "difference" and "effort") that prevents control systems from breaking down in low-noise situations. They demonstrated that this new foundation leads to better, more stable performance in standard control tests.

Technical Summary: Well-Posed KL-Regularized Control via Wasserstein and Kalman–Wasserstein KL Divergences

1. Problem Statement

The Kullback-Leibler (KL) divergence is a standard regularizer in reinforcement learning (RL) and optimal control, appearing in linearly-solvable Markov decision processes, path-integral control, and trust-region methods. However, the classical KL divergence exhibits a critical pathology in low-noise regimes. When comparing controlled dynamics to passive dynamics with small Gaussian process noise, the KL penalty reduces to a quadratic control cost whose coefficient scales inversely with the noise covariance. As the noise covariance vanishes (the deterministic limit), this penalty diverges to infinity, rendering the regularized control problem singular. This "low-variance instability" causes optimal feedback gains to vanish, leading to poor closed-loop performance and instability in deterministic or near-deterministic settings.

Furthermore, the classical KL divergence is defined via Fisher–Rao geometry, which is intrinsic to the statistical manifold but ignores the geometry of the underlying state space. Consequently, KL reacts discontinuously to support mismatches (e.g., $KL(\delta_x | \delta_y) = \infty$ for any $x \neq y$ ), failing to encode the distance between states.

2. Methodology

The authors propose a unified information-geometric framework to construct "state-space-aware" KL analogs. The core methodology involves replacing the Fisher–Rao geometry in the dynamical formulation of the KL divergence with transport-based geometries.

2.1 Geometric Framework

The paper defines a general class of divergences, $D_G$ , based on a Riemannian geometry $G$ on the manifold of probability densities.

Dynamical Formulation: The divergence is defined as a time-weighted action integral along $G$ -dual geodesics:
$D_G(\mu | \nu) = \inf \left\{ \int_0^1 t \|\dot{\gamma}(t)\|_{\gamma(t)}^2 dt : \gamma(0)=\mu, \gamma(1)=\nu \right\}$
where $\gamma$ is a $G$ -dual geodesic.
Geometric Choices: The authors investigate replacing the Fisher–Rao metric ($G=FR$) with:
- Wasserstein-2 ( $G=O$ ): Otto's metric, encoding state-space geometry.
- Linearized Wasserstein ( $G=\dot{H}^{-1}$ ): A flat structure.
- Kalman–Wasserstein ($G=KW$): A metric incorporating a regularization parameter $\lambda$ and covariance information.
- Stein Geometry: Using a bilinear kernel.

2.2 Theoretical Derivation

Dual Geodesics: The authors characterize $G$ -dual geodesics as metric gradient flows of a potential energy functional.
Closed-Form Expressions: They derive explicit formulas for these divergences between elliptic distributions (a class including Gaussians and Student's $t$ $t$ -distributions).
- For the Wasserstein case ($WKL$), the divergence remains finite even when covariances degenerate.
- For the Kalman–Wasserstein case ($KWKL$), the divergence incorporates a "noise floor" via a regularization parameter $\lambda$ , effectively interpolating between the classical KL and the Wasserstein KL.

2.3 Application to Control

The proposed divergences are applied as regularizers in Linear Quadratic Regulator (LQR) problems with Gaussian process noise.

Regularization Term: The divergence between the controlled transition distribution $p_t$ and the uncontrolled distribution $p_0$ yields an additive quadratic control penalty $u_t^\top R_\circ u_t$ .
Control Penalties:
- Classical KL ($FR $):$ R_{FR} = B^\top \Sigma_w^{-1} B$ (singular as $\Sigma_w \to 0$ ).
- Wasserstein ( $O$ ): $R_O = B^\top B$ (independent of noise).
- Kalman–Wasserstein ($KW $):$ R_{KW} = B^\top (\Sigma_w + \lambda I)^{-1} B$ (finite as $\Sigma_w \to 0$ ).

3. Key Contributions

General Construction: Introduction of $G$ -dual geodesics and the associated divergence $D_G$ for arbitrary geometries, recovering classical KL for Fisher–Rao and state-space-aware variants for transport geometries.
Gradient Flow Characterization: Proof that $G$ -dual geodesics correspond to metric gradient flows of potential energy, linking the divergences to geometric optimization and mean-field neural network dynamics.
Explicit Formulas: Derivation of closed-form expressions for $D_G$ between elliptic distributions in Wasserstein, linearized Wasserstein, and Kalman–Wasserstein geometries. These formulas explicitly demonstrate how transport-based geometries incorporate state-space information absent in classical KL.
Well-Posed Control: Demonstration that in LQR settings, Wasserstein and Kalman–Wasserstein regularizers remain finite as process noise vanishes, yielding well-posed optimal control problems where classical KL fails.
Empirical Validation: Analysis of double integrator and cart-pole systems showing that the proposed methods preserve nontrivial feedback and achieve better closed-loop performance in low-noise regimes compared to KL-regularized controls.

4. Results

Theoretical Limits: In the limit of vanishing noise ( $\Sigma_w \to 0$ ), the optimal feedback gain for classical KL regularization ( $F_{FR}$ ) converges to zero, effectively removing control. In contrast, the Kalman–Wasserstein gain ( $F_{KW}$ ) converges to a non-zero limit (specifically $F_O$ ), preserving stabilizing feedback.
Interpolation: The Kalman–Wasserstein divergence ( $D_{KW}$ ) acts as a smooth interpolation between classical KL (high noise) and Wasserstein KL (low noise) controlled by the parameter $\lambda$ .
Stability: Experiments on a double integrator show that as noise decreases, KL-regularized controllers approach the stability boundary (spectral radius $\to 1$ ), leading to large oscillations. WKL and KWKL controllers maintain spectral radii well below 1, ensuring stability.
Anisotropic Noise: In scenarios with anisotropic noise, KWKL adapts directionally, treating low-noise dimensions like WKL (ignoring noise scale) and high-noise dimensions like KL (scaling with noise), whereas KL suppresses control entirely in low-noise dimensions.

5. Significance and Claims

The paper claims that the proposed state-space-aware divergences provide a principled solution to the low-noise degeneracy of KL regularization without requiring a complete overhaul of RL algorithms. By replacing the Fisher–Rao geometry with transport-based geometries, the authors achieve:

Robustness: Control problems remain well-posed in the deterministic limit.
Geometric Consistency: The regularizers respect the underlying state-space geometry, avoiding the discontinuous behavior of KL regarding support mismatches.
Practical Utility: The derived closed-form expressions allow for immediate integration into standard LQR and Gaussian policy frameworks used in continuous control.

The authors position this work as a geometric framework for constructing KL-type divergences rather than a complete deep RL algorithm. They suggest that while validated in the tractable LQR setting, the next logical step is empirical evaluation in model-free and model-based reinforcement learning. The work is funded by the European Research Council (ERC) under the NEITALG project.