Optimal transport improves distribution comparison in reinforcement learning
Optimal Transport Meets Reinforcement Learning: A Survey
Artificial IntelligenceMachine Learning
Summary
Reinforcement learning methods often need to compare different sets of outcomes or behaviors, which can be tricky when these sets differ a lot. The authors review how a mathematical tool called optimal transport helps measure differences more effectively by considering the cost to change one behavior into another. This can improve learning from examples, using offline data, or handling changes in the environment. They also discuss practical challenges and future questions in using optimal transport for reinforcement learning.
What this means in practice
- •For reinforcement learning engineers: Design reinforcement learning algorithms that better handle weakly overlapping distributions using optimal transport based objectives.
- •For machine learning tool developers: Create software libraries or modules that implement optimal transport techniques to improve imitation and offline learning workflows.
A survey. It maps existing work.
Authors
Yujie Zhu, Charles A. Hepburn, Matthew Thorpe, Giovanni Montana
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 \emph{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.