Well-Posed KL-Regularized Control via Wasserstein and Kalman-Wasserstein KL Divergences
Not on the map. Why?
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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 21 of its 79 references that are in the map. First posted on arXiv 2 Feb 2026.
Builds on
21 of 79 references.
References · 21
- 2015Trust Region Policy Optimizationcited by 2,503
- 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,277
- 2016Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithmcited by 498
- 2017Wasserstein Auto-Encoderscited by 371
- 2020Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Losscited by 267
- 2017Stein Variational Gradient Descent as Gradient Flowcited by 149
- 2018A Tour of Reinforcement Learning: The View from Continuous Controlcited by 129
- 2019On the geometry of Stein variational gradient descentcited by 82
- 2018Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusionscited by 77
- 2019Accelerated Information Gradient flowcited by 38
- 2021KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Supportcited by 37
- 2022On Pathologies in KL-Regularized Reinforcement Learning from Expert Demonstrationscited by 18
9 more
- 2021On a linear Gromov-Wasserstein distancecited by 17
- 2021Stein Variational Gradient Descent: many-particle and long-time asymptoticscited by 15
- 2024Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spacescited by 15
- 2023Towards Understanding the Dynamics of Gaussian-Stein Variational Gradient Descentcited by 13
- 2024Wasserstein Gradient Flows of MMD Functionals with Distance Kernel and Cauchy Problems on Quantile Functionscited by 7
- 2025Accelerated Stein Variational Gradient Flowcited by 5
- 2025Gradient Flows of Potential Energies in the Geometry of Sinkhorn Divergencescited by 3
- 2025Towards understanding Accelerated Stein Variational Gradient Flow -- Analysis of Generalized Bilinear Kernels for Gaussian target distributionscited by 2
- 2025Information Geometry for Wasserstein KL Divergence of Gaussian Measures on ℝncited by 1
Built on by
Not cited yet.