Optimal Transport Meets Reinforcement Learning: A Survey

Yujie Zhu; Charles A. Hepburn; Matthew Thorpe; Giovanni Montana

stat.MLOct 1, 2026

Optimal Transport Meets Reinforcement Learning: A Survey

Authors: Yujie Zhu, Charles A. Hepburn, Matthew Thorpe, Giovanni Montana

Organizations: University of Warwick

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.

Figures & tables

Panel of Figure 1 — Figure 1 : Where optimal transport enters the RL pipeline. Each spoke represents a distinct role that OT plays, with the representative objective and section reference. Here $\mathsf{W}$ denotes a generic OT-based discrepancy, which may differ across applications, e.g., $\mathsf{W}_{1},\mathsf{W}_{2}$ or an entropy-regularised formulation. Additional roles not shown include safety constraints (Sections 7.2 , 7.3 ), preference alignment (Section 7.4 ), transfer and curricula (Sections 7.6 , 8.3 ), and multi-policy/multi-agent diversity (Section 8 ).

Panel of Figure 2 — Figure 2 : The ground cost determines what OT counts as a good match. In both panels, the source and target marginals are the same, but the ground cost changes. Because OT minimises transport cost under that matrix, the optimal plan changes as well. This is why the ground cost is a modelling choice, not a numerical detail.

Panel of Figure 3 — Figure 3 : Offline reward labelling via optimal transport. An optimal coupling $\Gamma^{*}$ is computed between empirical state marginals of learner and expert trajectories, treating each as a uniform distribution over its states. This figure shows an example of a coupling, where each row sums to $1/T=0.20$ ; each column sums to $1/T_{E}=0.25$ . The reward at each learner state $s_{t}$ is the negative transport cost to matched expert states under the optimal plan, $r(s_{t})=-\sum_{t^{\prime}=1}^{T_{E}}c(s_{t},s^{E}_{t^{\prime}})\,\Gamma^{*}_{t,t^{\prime}}$ . A near-diagonal structure emerges when trajectories are temporally aligned. The coupling in the figure is an example and is not necessarily the optimal coupling.

Panel of Figure 4 — Figure 4 : Example of optimal transport between two continuous distributions. The heatmap in the left figure shows the standard OT plan, and the right figure shows the entropy-regularised OT plan ( $\varepsilon=0.01$ ), which is smoother. Brighter colour indicates more mass is transported.

Panel of Figure 5 — Figure 5 : Three OT formulations. (a) Balanced OT enforces exact marginal constraints: all source mass reaches the target. (b) Partial OT transports only a fraction $m$ of the total mass, leaving the rest unmatched ( $\times$ )—useful when data contain outliers or suboptimal demonstrations (Section 6 ). (c) Unbalanced OT allows mass creation and destruction at a cost (red arrows), softly penalising marginal deviations rather than enforcing them exactly. Dot sizes indicate relative mass.

Panel of Figure 6 — Figure 6 : OT in imitation learning and reward inference (Section 3 ). OT compares learner and expert behaviour distributions via dual critics or primal couplings, yielding shaping rewards or pseudo-reward labels (Sections 3.3 – 3.4 ). Temporal/intent-aware extensions modify the coupling or ground cost (Section 3.5 ); cross-domain settings use GW (Section 3.6 ).

Panel of Figure 7 — Figure 7 : Standard vs. temporally-aware OT matching between learner and expert trajectories. (a) Without temporal constraints, the coupling can align temporally distant states, producing crossing matches that yield noisy pseudo-rewards. (b) With a band constraint $|t-t^{\prime}|\leq 3$ (or context-window costs), matches are near-diagonal, reducing spurious alignments while tolerating modest speed variation.

Panel of Figure 8 — Figure 8 : OT in policy optimisation and regularisation (Section 4 ). Wasserstein trust regions constrain per-state action distributions or trajectory embeddings to stay close to a reference policy (Sections 4.2 – 4.4 ). Wasserstein-trained critics use OT losses on return or Q-value distributions (Section 4.5 ). The blue dashed bidirectional arrow connects two distributions compared by OT, while grey solid arrows represent connections in the RL algorithm.

Panel of Figure 9 — Figure 9 : OT in model-based RL (Section 5 ). A learned world model $P_{\phi}(\cdot\mid s,a)$ is trained to match the true transitions under a Wasserstein loss that upper-bounds worst-case value error. Bisimulation metrics embed this comparison into a fixed-point state metric (Section 5.3 ).

Panel of Figure 10 — Figure 10 : OT in offline RL (Section 6 ). The agent learns entirely from a fixed dataset $\mathcal{D}$ collected from the behaviour policy $\pi_{\beta}$ . OT enters as a behaviour constraint $\mathsf{W}(\pi_{\theta}(\cdot\mid s),\pi_{\beta}(\cdot\mid s))$ that keeps the learned policy geometrically close to the data-generating policy with both induced distributions approximated by collected trajectories (Section 6.2 ). Partial and unbalanced OT relax mass conservation to handle suboptimal demonstrations.

Panel of Figure 11 — Figure 11 : OT in robustness, safety, and transfer (Section 7 ). A Wasserstein ambiguity set $\mathcal{U}_{\varepsilon}$ around the nominal transitions defines worst-case Bellman updates (Section 7.1 ). OT-based visitation regularisation and preference alignment shape where the policy visits or how it aligns with human feedback (Sections 7.2 – 7.4 ).

Panel of Figure 12 — Figure 12 : The Wasserstein ambiguity set for distributionally robust RL. The nominal transition $P_{0}(\cdot\mid s,a)$ (dark-blue dot) sits at the centre of a Wasserstein ball $\mathcal{U}_{\delta}=\{P:\mathsf{W}_{p}(P,P_{0})\leq\delta\}$ inside the space of transition distributions $\mathcal{P}(\mathcal{S})$ . Grey dots denote other candidate distributions inside the ball. The worst-case is represented by $P^{\ast}$ (red dot). Two modelling choices shape this game: the ground metric $d$ determines which perturbations are considered admissible (through the $\mathsf{W}_{p}$ geometry), while the radius $\delta$ controls how far models may move away from the nominal one, determining the conservatism of the resulting policy.

Application	Start with	Distributions	Cost default	Key pitfall
Online imitation	$\mathsf{S}_{\varepsilon}$ (Sinkhorn div.)	occupancies $\rho_{\pi}$ vs $\rho_{\pi_{E}}$	learned features	$\varepsilon$ sensitive to cost scale
Offline reward labelling	entropic OT + mask	trajectory segments	cosine on encoder	spurious matches without temporal mask
Policy trust region	per-state Sinkhorn	$\pi(\cdot\|s)$ vs $\pi_{\mathrm{old}}(\cdot\|s)$	Euclidean (actions)	radius too tight $\to$ no improvement
Offline constraint	$\mathsf{W}_{1}$ penalty	$\pi(\cdot\|s)$ vs $\pi_{\beta}(\cdot\|s)$	action-space metric	overly strong $\to$ clones $\pi_{\beta}$
World model	$\mathsf{W}_{1}$ or $\mathsf{S}_{\varepsilon}$	$P_{\phi}(\cdot\|s,a)$ vs $P(\cdot\|s,a)$	latent features	optimising prediction, not control
Robustness	$\mathsf{W}_{p}$ ball	transitions $P(\cdot\|s,a)$ around $P_{0}$	env. geometry	uncalibrated $\delta$ $\to$ over-conservatism

Table 1 : Decision guide: the authors’ suggested starting points by application area . These are reasonable first experiments supported by the evidence discussed in Sections 3 – 8 , not empirically established best defaults across problem classes .

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

MDP and RL
$\mathcal{S},\mathcal{A}$	State and action spaces.
$\mathcal{M}=(\mathcal{S},\mathcal{A},P,R,\gamma,p_{0})$	Discounted Markov decision process (MDP).
$\gamma$	Discount factor.
$p_{0}\ (\text{or }\nu)$	Initial-state distribution (default $p_{0}$ ; $\nu$ when varying the start distribution).
$P(\cdot\mid s,a)$	Environment transition kernel.
$P_{\pi}$	Trajectory law generated by $(p_{0},\pi,P)$ .

Appendix

Table 2: Notation used throughout the survey.

Method (area)	OT role	Compared objects	OT form	Implementation / notes
Adversarial occupancy matching and discriminator-based IL
WAIL ( Xiao et al., 2019 ) , WDAIL ( Zhang et al., 2020 ) ( generative adversarial IL )	Discriminator / reward shaping	State-action occupancy measures	$\mathsf{W}_{1}$ (Euclidean cost)	Kantorovich duality with 1-Lipschitz critic (WGAN-style training).
W-IRL ( Wu et al., 2023 ) ( IRL )	Discriminator / reward shaping	State marginals	$\mathsf{W}_{1}$ (Euclidean cost)	Kantorovich duality with a Lipschitz-constrained critic.
LWAIL ( Yang et al., 2026 ) ( adversarial IL; twin delayed DDPG (TD3) )	Latent discriminator + reward	Latent state-pair (transition) occupancies (ICVF-pretrained features)	$\mathsf{W}_{1}$ (Euclidean cost)	Duality-based critic in latent space; pretraining via ICVF as proposed.
SIL ( Papagiannis and Li, 2022 ) ( adversarial IL )	Discriminator / reward shaping via primal OT	State-action occupancy measures	Entropy-regularised $\mathsf{W}_{\varepsilon}$ (OT-GAN-style cost ( Salimans et al., 2018 ) )	Sinkhorn iterations in minibatches; entropic OT objective.
PWIL ( Dadashi et al., 2021a ) ( IL )	Reward-from-transport (non-adversarial)	State-action occupancy measures	Upper bound on $\mathsf{W}_{1}$	Uses approximate primal coupling (e.g. POT ( Flamary et al., 2024 ) ); transport cost converted to rewards.

Appendix

Table 3 : Representative OT uses for imitation learning and reward inference.

Method (setting)	OT role	Compared objects	OT form	Implementation / notes
Policy update constraints and regularisers
WPO, SPO ( Song et al., 2023 ) ( on-policy; discrete )	Trust-region constraint	Per-state action distributions $\pi^{\prime}(\cdot\mid s)$ vs. $\pi(\cdot\mid s)$	$\mathsf{W}_{1}$ / Sinkhorn	Per-state Lagrangian dual; closed-form distribution update under task-dependent action costs (e.g. $0$ – $1$ , $\mathrm{L}_{1}$ ).
OT-TRPO ( Terpin et al., 2022 ) ( TRPO-style )	Trust-region constraint	Action distributions (e.g. Gaussians)	$\mathsf{W}_{2}$	Closed-form $\mathsf{W}_{2}$ for Gaussians under quadratic costs; dual optimisation over the multiplier.
BGPG, BGES ( Pacchiano et al., 2020 ) ( PG / ES )	Behaviour regulariser	Embedding distributions $\mu_{\pi_{\theta}}$ vs. $\mu_{\mathrm{ref}}$	Entropy-regularised $\mathsf{W}_{\varepsilon}$	Dual OT test functions on embedding space; OT signal used as shaping/weights in PG or ES updates.
WPR ( Na et al., 2026 ) ( RLHF )	Policy regularisation	learned policy $\pi_{\theta}$ vs. reference policy $\pi_{ref}$	Entropy-regularised $\mathsf{W}_{\varepsilon}$	Dual formulation with Sinkhorn-Knopp iterations yields a tractable semantic regulariser compatible with PPO-style updates.
Wasserstein geometry: gradient flows and natural gradients

Appendix

Table 4 : Representative OT methods for online policy optimisation and closely related objectives.

Method (setting)	OT role	Compared objects	OT form	Implementation / notes
Value-awareness
VAML ( Farahmand et al., 2017 ; Asadi et al., 2018 ) ( model-based )	Value-aware transition model loss	$P(\cdot\mid s,a)$ vs. $P_{\phi}(\cdot\mid s,a)$	$\mathsf{W}_{1}$	Equivalence between Lipschitz test-function losses and a $\mathsf{W}_{1}$ discrepancy on next-state distributions.
Link to bisimulation metric
( Calo et al., 2024 ) (dynamic programming)	Efficient OT distances for Markov processes	Stationary Markov processes (occupancy of state pairs)	Discounted $\mathsf{W}_{1}$	Sinkhorn value iteration; Markov-structured couplings improve computation and convergence.
Latent MDP training
WAE-MDP ( Delgrange et al., 2023a ) ( latent world model )	OT-based model learning	Trace/trajectory distributions (encoded/decoded)	$\mathsf{W}_{1}$ -type	WAE/WGAN-style objective; dual-critic realisation with gradient penalty in practice.

Appendix

Table 5 : Representative OT methods for model-based RL and Wasserstein world models.

Method (area)	OT role	Compared objects	OT form	Implementation / notes
BRAC ( Wu et al., 2019 ) ( offline RL )	Policy regularisation / value penalty	Action distributions $\pi(\cdot\mid s)$ vs. $\pi_{\beta}(\cdot\mid s)$	$\mathsf{W}_{1}$	Dual form with a discriminator (WGAN-style critic).
Q-DOT ( Omura et al., 2025 ) ( offline RL )	Policy regularisation via transport map	Action distributions (state-conditional)	$\mathsf{W}_{2}$	ICNN parameterisation of a convex potential / transport map.
VGF ( Xu et al., 2026 ) ( offline RL )	Gradient flow	Action distributions (state-conditional)	$\mathsf{W}_{2}$	JKO scheme, realised by an SVGD-style RKHS-restricted particle update rather than exact JKO minimisation; implicit behaviour control via a transport budget.
PPL ( Asadulaev et al., 2024 ) ( offline RL )	Partial behaviour cloning / selective matching	Action distributions under partial OT constraint	Partial OT	Maximin formulation with a tunable partiality parameter.
Givchi et al. (2021) ( RL (policy optimisation) )	Unbalanced OT formulation	Projected state and action marginals of occupancies	Unbalanced OT	Multiple convex constraints handled via Dykstra’s algorithm; Bregman (KL) divergences used in practice.
Fatras et al. (2021) ( Domain adaptation )	Domain-invariant representation learning	Source vs. target distributions	Unbalanced minibatch OT	POT and Geomloss packages

Appendix

Table 6 : Representative OT uses for offline policy optimisation.

Method (area)	OT role	Compared objects	OT form	Implementation / notes
Distributionally robust control
WR 2 L ( Abdullah et al., 2019 ) ( robust RL )	Minimax constraint	Transition kernels $P(\cdot\mid s,a)$ within a Wasserstein ball	$\mathsf{W}_{2}$	Approximates $\mathsf{W}_{2}$ via a Hessian-of-expectation surrogate; zero-order optimisation as proposed.
DRMDP ( Kordabad et al., 2022 ; Yang, 2020 ; Hou et al., 2020 ; Queeney et al., 2024 ; Yang, 2017 ) ( robust control / robust RL )	Ambiguity set (Wasserstein ball)	Disturbance or transition distributions	OT-cost uncertainty set (generalises $\mathsf{W}_{p}$ )	Kantorovich duality yields tractable robust Bellman operators / reformulations.
FOM for DRMDP ( Clement and Kroer, 2021 ) ( robust control / robust RL )	Ambiguity set (Wasserstein ball)	Disturbance or transition distributions	$\mathsf{W}_{2}$	First-order methods.
Safety
( Baheri, 2023a ; Shahrooei and Baheri, 2024 ) ( safe RL )	Visitation shaping / regularisation	Risk distribution vs. state visitation distribution	Wasserstein with squared Euclidean cost	OT-derived costs used as penalties in value-based updates (as proposed).

Appendix

Table 7 : Representative OT uses for control shift and transfer.

Method (area)	OT role	Compared objects	OT form	Implementation / notes
Diversity and hierarchical RL
WURL ( He et al., 2022 ) ( unsupervised RL )	Intrinsic reward for diversity	State marginal distributions across policies/skills	$\mathsf{W}_{1}$	Dual form and projected Wasserstein discrepancy amortise rewards from OT plans.
PWSEP ( Yang et al., 2025 ) ( unsupervised RL )	Skill separation / discovery	State distributions conditioned on skills	Wasserstein (sliced approximation)	Uses sliced Wasserstein distance for efficient estimation of separability objective.
WQDIL ( Yu et al., 2025 ) ( quality-diversity IL )	Stable reward learning in latent space	Latent distributions of demonstrations vs. policy trajectories	WAE-WGAN	Duality-based adversarial training within WAE latent space (as proposed).
WDER ( Li et al., 2023 ) ( hierarchical RL )	Diversity regulariser for subpolicies	Behavioural embeddings / action distributions	Smoothed $\mathsf{W}_{\varepsilon}$ with random-feature approximation	Duality with random-feature approximations enabling scalable updates.
HiPBOT ( Le et al., 2023 ) ( hierarchical control )	Policy blending (mixture weights)	Expert-policy and agent-policy priors (temperature-weighted)	Unbalanced entropic OT	Sinkhorn-like scaling algorithm ( Chizat et al., 2018 ) .

Appendix

Table 8 : Representative OT uses in multi-agent RL and other applications.

Study (role)	Evid.	Environment & setting	OT form & implementation	Within-study comparison & outcome	Limitation noted
OTR ( Luo et al., 2023b ) (reward labelling)	E	D4RL locomotion, AntMaze and Adroit; offline, reward-free, one expert demonstration	Entropic OT between learner/expert state marginals; Sinkhorn; cosine cost on raw states	OTR+IQL against oracle-reward IQL and prior reward-learning baselines on D4RL locomotion (HalfCheetah, Hopper, Walker2d): near-oracle aggregate performance from one demonstration, avoiding the severe drops seen in some alternatives; not equal or superior on every task	Labels are trajectory-context dependent; state-marginal OT ignores ordering
Dong et al. (2026) (critique)	E, T	32 benchmarks, offline and online (D4RL locomotion, AntMaze, Adroit); four learners: IQL, ReBRAC, TD3+BC, DrQ-v2	Ablates two axes: Wasserstein vs. a point-to-set distance (MinDist), and with vs. without temporal constraints (SegMatch)	Negative/mixed : offline, the gains from Wasserstein approximation and temporal constraints largely vanish against a well-tuned learner; online, temporal constraints are essential while the proximity approximation is secondary	Temporal alignment does matter offline as the number of demonstrations grows
SinkhornDRL ( Sun et al., 2024 ) (critic loss)	E, T	55 Atari 2600 games, online, fixed model capacity	Sinkhorn divergence between particle-based return distributions	Mixed : against QR-DQN and MMD-DQN, slower convergence early but better aggregate human-normalised score; largest gains with larger action spaces and multi-dimensional rewards	$\approx$ 20% more compute than MMD-DQN (a different overhead is reported against C51 and QR-DQN); sensitive to $\varepsilon$ , iteration count and particle number
TemporalOT ( Fu et al., 2024 ) (reward labelling)	E	Nine Meta-world image-based manipulation tasks; online imitation from expert video, DrQ-v2 backbone	Masked (banded) entropic OT with context-window cosine cost on encoder features	Against an online OTR variant, ADS, GAIfO, BC and an oracle task reward: higher success rate without task rewards; an ablation shows both the context cost and the mask contribute	Requires online interaction; a temporally informed proxy rather than a full dynamic-OT formulation
SMMOTIL ( Sebag et al., 2023 ) (reward labelling)	E	Pendulum-v0 and CartPole-v0, with five experts differing in pole length or mass	Sliced multi-marginal OT; closed-form 1D solvers	Against concatenating the demonstrations into a single empirical measure: higher mean episodic reward under both length and mass diversity	Two low-dimensional control tasks only; slicing approximates the geometry
WPO / SPO ( Song et al., 2023 ) (trust region)	E, T	Tabular MDPs; continuous actions via an IQN target	Per-state $\mathsf{W}_{1}$ / Sinkhorn trust region; 1D dual solve, no parametric form assumed	Against KL trust regions: monotonic improvement with exact advantages, and global convergence in the tabular setting under a decaying multiplier	Tabular guarantee needs finite spaces, non-negative rewards and full initial-state support; the bound degrades by $2\epsilon/(1-\gamma)$ under advantage error

Appendix

Table 9 : Experimental results for representative studies. Evidence type is denoted by E for empirical evaluation and T for explicit method-specific theoretical analysis. Labels refer to the cited study, not necessarily to each reported outcome.

Explore similar work

Reliable Modeling of Distribution Shifts via Displacement-Reshaped Optimal Transport
May 6, 2026Philip Naumann, Jacob Kauffmann, Klaus-Robert Müller +1Distribution Shifts Transport
BIFOLD – Berlin Institute for the Foundations of Learning and Data, 10587 Berlin, Germany · Machine Learning Group, Technische Universität Berlin, 10587 Berlin, Germany · Department of Artificial Intelligence, Korea University, Anam-dong, Seongbuk-gu, Seoul 02841, Korea +2
Sliced-Regularized Optimal Transport
Apr 27, 2026Khai NguyenGradient
The University of Texas at Austin
Simultaneous Neural Optimal Transport
Sep 29, 2026Milena Gazdieva, Kirill Sokolov, Jiawei Chen +2Optimal Transport Approach Image Restoration
Applied AI Institute AXXX Moscow, Russia · Lomonosov Moscow State University Moscow, Russia · MIRIAI Moscow, Russia

May 6, 2026, cs.LGJ/K move · Enter open · S save

Reliable Modeling of Distribution Shifts via Displacement-Reshaped Optimal Transport

Philip Naumann, Jacob Kauffmann, Klaus-Robert Müller, Grégoire Montavon

BIFOLD – Berlin Institute for the Foundations of Learning and Data, 10587 Berlin, Germany · Machine Learning Group, Technische Universität Berlin, 10587 Berlin, Germany · Department of Artificial Intelligence, Korea University, Anam-dong, Seongbuk-gu, Seoul 02841, Korea+2

Distribution Shifts Transport Distributionally-Robust Optimization Geometric Metrics

Open paper AskPDF arXiv

Optimal transport (OT) is a central framework for modeling distribution shifts. Because OT compares distributions directly in input space, a well-designed ground metric between observations is essential to ensure that the optimizer does not violate the true geometry of change. We propose Displacement-Reshaped Optimal Transport (ReshapeOT), a method that reshapes the ground metric by integrating observed sample displacements as an additional source of knowledge. Technically, ReshapeOT replaces the Euclidean metric with a Mahalanobis distance estimated from displacement second moments. This effectively carves expressways through the input space, inviting transport solutions that better align with observed displacements. Our method is computationally lightweight, integrates seamlessly into any OT solver that operates on a cost matrix, and can be kernelized for further flexibility. Experiments on synthetic and real-world data show that ReshapeOT achieves substantial gains in transport reliability. We further demonstrate our method's usefulness in two practical use cases.

Abstract

Explore similar work

Reliable Modeling of Distribution Shifts via Displacement-Reshaped Optimal Transport

Sliced-Regularized Optimal Transport

Simultaneous Neural Optimal Transport