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

arXiv: https://arxiv.org/abs/2602.02250 · HTML: https://mlsanity.com/paper/2602.02250/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 21 of its 79 references that are in the map.
- First posted on arXiv 2 Feb 2026.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 21 of 21)

- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Wasserstein Auto-Encoders](https://mlsanity.com/paper/1711.01558/) (2017)
- [Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss](https://mlsanity.com/paper/2002.04486/) (2020)
- [Stein Variational Gradient Descent as Gradient Flow](https://mlsanity.com/paper/1704.07520/) (2017)
- [A Tour of Reinforcement Learning: The View from Continuous Control](https://mlsanity.com/paper/1806.09460/) (2018)
- [On the geometry of Stein variational gradient descent](https://mlsanity.com/paper/1912.00894/) (2019)
- [Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions](https://mlsanity.com/paper/1806.08141/) (2018)
- [Accelerated Information Gradient flow](https://mlsanity.com/paper/1909.02102/) (2019)
- [KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Support](https://mlsanity.com/paper/2106.08929/) (2021)
- [On Pathologies in KL-Regularized Reinforcement Learning from Expert Demonstrations](https://mlsanity.com/paper/2212.13936/) (2022)
- [On a linear Gromov-Wasserstein distance](https://mlsanity.com/paper/2112.11964/) (2021)
- [Stein Variational Gradient Descent: many-particle and long-time asymptotics](https://mlsanity.com/paper/2102.12956/) (2021)
- [Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spaces](https://mlsanity.com/paper/2402.04613/) (2024)
- [Towards Understanding the Dynamics of Gaussian-Stein Variational Gradient Descent](https://mlsanity.com/paper/2305.14076/) (2023)
- [Wasserstein Gradient Flows of MMD Functionals with Distance Kernel and Cauchy Problems on Quantile Functions](https://mlsanity.com/paper/2408.07498/) (2024)
- [Accelerated Stein Variational Gradient Flow](https://mlsanity.com/paper/2503.23462/) (2025)
- [Gradient Flows of Potential Energies in the Geometry of Sinkhorn Divergences](https://mlsanity.com/paper/2511.14278/) (2025)
- [Towards understanding Accelerated Stein Variational Gradient Flow -- Analysis of Generalized Bilinear Kernels for Gaussian target distributions](https://mlsanity.com/paper/2509.04008/) (2025)
- [Information Geometry for Wasserstein KL Divergence of Gaussian Measures on ℝn](https://mlsanity.com/paper/2503.24022/) (2025)

## Abstract (arXiv)

Kullback-Leibler (KL) divergence regularization is widely used in reinforcement learning, but it becomes infinite under support mismatch and can degenerate in low-noise regimes. Using a unified information-geometric framework, we introduce KL analogs by replacing the Fisher-Rao geometry in the dynamical formulation of the KL with transport-based geometries, and derive closed-form expressions for common distribution families. Between elliptic distributions, these divergences remain finite for degenerating equal covariances 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 and yield well-posed problems. In both the double integrator and cart-pole examples, the resulting controls preserve nontrivial feedback and achieve better closed-loop performance.