Optimal Transport Meets Reinforcement Learning: A Survey
stat.ML, cs.AI, cs.LG
Submitted: 2026-10-01
Updated: 2026-10-01
Code: https://github.com/PythonOT/POT
License: http://creativecommons.org/licenses/by/4.0/
The gist: Reinforcement learning (RL) algorithms frequently compare probability distributions, such as state visitation distributions induced by policies and experts, action distributions from learned policies
Terminology
Abstract
Reinforcement learning (RL) algorithms frequently compare probability distributions, such as state visitation distributions induced by policies and experts, action distributions from learned policies and offline datasets, or transition distributions from learned models and environments. However, commonly used divergences may become ineffective when these distributions overlap weakly, which is frequently encountered in imitation learning, offline RL, and deployment under distribution shift. Optimal transport (OT) offers an alternative by measuring the cost of moving probability mass from one distribution to another under a ground cost that encodes task geometry. This survey covers how OT is used inside RL objectives and algorithms. For each method, we identify: the role OT plays, the distributions compared, the OT formulation used, and the treatment of temporal structure. Beyond categorising existing methods, we discuss the motivations behind different OT choices, practical considerations such as cost design and computational challenges, and highlight open problems including scalable trajectory-level transport, principled handling of mass mismatch, and theoretical analysis for OT-regularised RL.
Sources
- Wasserstein Robust Reinforcement Learning
- Equivalence Between Wasserstein and Value-Aware Loss for Model-based Reinforcement Learning
- Risk-Aware Reinforcement Learning through Optimal Transport Theory
- Understanding Reward Ambiguity Through Optimal Transport Theory in Inverse Reinforcement Learning
- The Synergy Between Optimal Transport Theory and Multi-Agent Reinforcement Learning
- The Cramer Distance as a Solution to Biased Wasserstein Gradients
- Align Your Intents: Offline Imitation Learning via Optimal Transport
- Wasserstein Distance Maximizing Intrinsic Control
- Assessing the Impact of Distribution Shift on Reinforcement Learning Performance
- Distributionally-Constrained Policy Optimization via Unbalanced Optimal Transport
- Combining Reinforcement Learning and Optimal Transport for the Traveling Salesman Problem
- Robust Reinforcement Learning with Wasserstein Constraint
- A Note on Loss Functions and Error Compounding in Model-based Reinforcement Learning
- Safe Wasserstein Constrained Deep Q-Learning
- Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems
- Improved Complexity Analysis of the Sinkhorn and Greenkhorn Algorithms for Optimal Transport
- Merging without Forgetting: Continual Fusion of Task-Specific Models via Optimal Transport
- Heterogeneous Federated Reinforcement Learning Using Wasserstein Barycenters
- On Wasserstein Reinforcement Learning and the Fokker-Planck equation
- Online Distributionally Robust LLM Alignment via Regression to Relative Reward
Related papers
- Behavior of prediction performance metrics with rare events
- Optimal Estimation of Generic Dynamics by Path-Dependent Neural Jump ODEs
- A Posterior-Dynamics Framework for Imaging Inverse Problems with Pretrained Diffusion Priors
- Efficient quadratic entropy with distance sketches
- Embedding-Bias in Conditional Independence Testing
- From Log-Odds to Shapley Values: An Explanatory Geometry for the Weighted Naive Bayes Classifier