PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimal Transport Meets Reinforcement Learning: A Survey

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

arXiv:2610.01413v1stat.MLcs.AIcs.LG

TL;DR

RL often requires comparing distributions that may have weakly overlapping supports, limiting pointwise divergences when the differences matter for control. This survey organizes OT’s roles and formulations in RL, synthesizes recurring trade-offs, and identifies practical guidance and open problems.

  • Problem

    RL needs discrepancies that reflect control-relevant differences between distributions, while weak support overlap can make pointwise comparisons problematic.

    Source: §1 · p. 2, §2.2 · p. 6

  • Method

    The survey organizes OT in RL by its role in objectives and algorithms, the distributions compared, the OT formulation, and the handling of temporal structure.

    Source: §1 · p. 3

  • Results

    The survey synthesizes recurring trade-offs, practical guidance, and open problems across OT methods in reinforcement learning.

    Source: §9 · p. 47

  • Takeaways & Limitations

    The four-question framework helps compare OT methods usually discussed in separate literatures.

    Source: §1 · p. 3

    Dual-W1 approaches practically require stable enforcement or relaxation of 1-Lipschitzness, commonly through weight clipping or gradient penalties.

    Source: §3.3 · p. 16

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames many RL problems as comparing distributions that matter for control, while noting that pointwise divergences can falter when those distributions overlap weakly. It positions OT as a geometry-aware alternative and organizes the survey around its roles, compared distributions, formulations, and treatment of temporal structure.

1 Introduction

  • RL settings compare policy, expert, data, or transition distributions, so useful discrepancies must capture how distributions differ in ways that matter for control.Examples include imitation, policy optimisation, offline RL, robust control, and curriculum learning.Source: §1 · pp. 1–2
  • When supports overlap weakly, pointwise KL or JS comparisons can become problematic, with JS potentially saturating and KL becoming undefined or dominated by near-zero-density regions.The introduction identifies early imitation learning as a setting where learner states may be absent from expert data.Source: §1 · p. 2
  • OT measures the effort to move probability mass under a ground cost, remaining informative across disjoint supports while letting task geometry shape which states, actions, or trajectories count as similar.The ground cost can reflect kinematic or geodesic distance, or distance in a learned feature embedding.Source: §1 · p. 2
  • Its scope centers on work where OT substantively shapes RL or imitation-learning methods, while retaining selected enabling or adjacent topics and related sequential decision problems without aiming for an exhaustive catalogue.Examples of retained adjacent topics include goal transfer and domain adaptation; RLHF and contextual bandits receive brief discussion.Source: §1 · pp. 2–3
  • The survey compares methods by OT’s role, the distributions compared, the formulation used, and how temporal structure is handled; recurring roles include reward, regulariser, trust-region constraint, ambiguity set, and model loss.These questions are intended to make methods from separate literatures easier to compare.Source: §1 · p. 3
  • The survey provides a broad methodological overview across RL settings, unlike prior surveys that treat RL as one application domain or focus on multi-agent RL.It emphasizes how OT enters algorithms and why formulations are chosen, rather than empirical comparisons that depend on environments, tasks, and experimental settings.Source: §1 · p. 3

2 Preliminaries

The preliminaries define policy-induced distributions and OT variants, notation, and computational choices for comparing distributions in RL. They explain how formulation, cost design, and temporal structure relate to the task, and introduce tools for domain mismatch, incomplete mass, and high-dimensional transport.

2.1 Reinforcement learning in brief

  • OT compares policy-generated measures, especially discounted occupancies that encode where an agent goes and what it does while obeying a dynamics-induced mass-flow constraint.Occupancies can be compared directly or used to constrain policy updates in occupancy space.Source: §2.1 · p. 5
  • An empirical trajectory state marginal is a bag of states that discards order, so temporal structure must enter through the transported object, cost, or admissible couplings.Banded or masked couplings and windowed or intent-aware costs are ways to reintroduce temporal information.Source: §2.1 · p. 6

2.2 Optimal transport in brief

  • When supports overlap weakly, OT can provide informative gradients by measuring the effort to move mass, unlike KL, which may be undefined or infinite; the ground cost changes what counts as a good match.For point masses, W1(δx, δy) = d(x, y), while KL is infinite when x ≠ y.Source: §2.2 · p. 6, Figure 2 · p. 7
  • Kantorovich transport optimizes over couplings that can split mass, relaxing Monge’s map-based assignment when a deterministic destination for each source point is infeasible.In discrete settings, a coupling is a nonnegative matrix with row and column sums equal to the source and target histograms.Source: §2.2 · p. 7
  • HiRef’s reported log-linear time and linear space apply to a low-rank bijective Monge assignment scheme, with exact recovery conditional on optimal low-rank subproblem solutions.The cited low-rank problems are non-convex in practice, so numerical solvers are not guaranteed to find global optima; continuous OT is approximated by empirical measures to obtain discrete problems.Source: §2.2 · p. 8
  • The partial-OT dual maximin form uses an unconstrained measurable selector, not necessarily an optimal transport map, and its measurable-map interpretation requires selection and integrability conditions.The c-transform identity is the safer general statement; the maximin form is used in dual formulations of partial OT for offline RL.Source: §2.2 · p. 9
  • W1’s dual form replaces coupling optimization with a 1-Lipschitz critic that can be trained from samples, while W2 can yield unique gradient transport maps under absolute-continuity and regularity conditions.The W2 map conditions do not apply to the Kantorovich or Kantorovich-Rubinstein dualities used for empirical atomic measures.Source: §2.2 · p. 9
  • Entropic regularization makes discrete OT strictly convex and solvable by Sinkhorn iterations, each costing O(N^2) time and storage, instead of exact OT’s stated O(N^3 log N) complexity.The cited accuracy-dependent complexity guarantee requires choosing the entropic coefficient and marginal stopping tolerance as functions of the target accuracy.Source: §2.2 · pp. 8–10
  • The survey distinguishes genuine Wp from entropy-regularised Wε and debiased Sε, then selects variants by support overlap, desired smoothness, domain comparability, mass mismatch, dimension, and temporal structure.Wp denotes the genuine p-Wasserstein distance; regularised costs are not p-rooted, and Sε removes entropic self-transport bias without being a metric.Source: §2.2 · pp. 10, 12
  • Entropic regularisation produces smoother plans and facilitates repeated minibatch computation, but ε that is too small slows and destabilises Sinkhorn while ε that is too large washes out task geometry.Choose ε on the scale of typical ground costs; log-domain Sinkhorn helps stabilise small-ε computation.Source: §2.2 · pp. 10, 13, Figure 4 · p. 11
  • Gromov-Wasserstein compares within-domain pairwise relations when heterogeneous spaces lack a meaningful pointwise cross-domain cost, supporting structural alignment in transfer and imitation.It matches neighbourhood geometry rather than requiring direct correspondence between individual observations.Source: §2.2 · p. 10
  • Partial OT transports a prescribed fraction of mass, whereas unbalanced OT permits marginal deviations at a penalty, addressing incomplete, contaminated, or systematically mismatched data.Partial OT can leave outliers or suboptimal demonstrations unmatched; unbalanced OT penalises marginal differences, often with KL-type penalties.Source: §2.2 · p. 11, Figure 5 · p. 12
  • Wasserstein barycentres provide geometry-aware distributional averages used for task-distribution curricula and some multi-agent consensus formulations.They also underlie displacement interpolation techniques.Source: §2.2 · p. 11
  • Sliced Wasserstein approximates high-dimensional transport by averaging inexpensive one-dimensional projections, with generalised variants and learned approximations refining the surrogate.For N observations in d dimensions and L projection directions, the stated complexity is O(LdN+LN log N).Source: §2.2 · p. 12
  • Ground-cost design is the main modelling choice: use task-aware or learned feature costs where appropriate, and normalise heterogeneous units so no coordinate dominates.Compared objects can be states, actions, state-action pairs, or learned trajectory and behaviour embeddings.Source: §2.2 · p. 13

3 Imitation learning and reward inference

This section surveys how OT supports task-geometry-aware matching of learner and expert behavior distributions in imitation learning and reward inference, including online reward shaping, offline reward labelling, and coupling-derived rewards. It covers temporal constraints, intent representations, and relational matching to address alignment and cross-domain mismatch.

3.1 Imitation as occupancy matching

  • Imitation can match normalized discounted occupancy measures rather than actions one state at a time; OT changes the discrepancy central to that match.Behaviour cloning is brittle under covariate shift, while max-entropy IRL and adversarial imitation formulate distribution matching; GAIL uses a JS f-divergence (Ho & Ermon, 2016).Source: §3.1 · pp. 13–14

3.2 Why optimal transport helps

  • OT remains informative when learner and expert occupancies overlap weakly, and its ground cost defines behavioral similarity over state-action, state-only, or learned feature spaces.Unlike geometry-blind divergences, transport measures how far probability mass must move under task-dependent geometry.Source: §3.2 · pp. 14–15

3.3 Reward shaping in the online adversarial setting

  • Dual-W1 methods use a learned Kantorovich critic as a shaping reward, but their regularized Lipschitz constraint is an approximation requiring stable enforcement (Xiao et al., 2019).WAIL uses a smooth penalty in place of the hard 1-Lipschitz constraint; common controls include weight clipping and gradient penalties.Source: §3.3 · pp. 15–16
  • Primal methods derive shaping rewards from transport couplings: SIL uses entropic Sinkhorn OT, while PWIL uses greedy matching to avoid adversarial training (Dadashi et al., 2021a).SIL provides an explicit, smooth minibatch coupling but depends on entropic-parameter and feature-cost scaling; PWIL yields a history-dependent reward.Source: §3.3 · p. 16

3.4 Offline reward labelling

  • OTR converts trajectory couplings into offline pseudo-rewards and reports robust performance without degradation; transition-based labels can be reused with an upper bound under monotonic improvement (Luo et al., 2023b; Chang et al., 2024).OTR uses entropy-regularized OT with Sinkhorn iterations; the reuse guarantee concerns pseudo-rewards from previously collected trajectories.Source: §3.4 · p. 17
  • PW-DICE uses primal state-occupancy transport with flexible ground costs and unifies Wasserstein and f-divergence matching; SMMOTIL aggregates diverse experts in OT geometry (Sebag et al., 2023).SMMOTIL’s sliced multi-marginal approach reduces instability and improves performance with diverse demonstrations compared with naive concatenation.Source: §3.4 · pp. 17–18

3.5 Temporal and intent-aware extensions

  • Trajectory-level matching can ignore order and produce noisy, context-dependent rewards; TemporalOT uses windowed context costs and near-diagonal coupling constraints to reduce spurious alignments (Fu et al., 2024).The band constraint permits modest speed variation while discouraging matches between temporally distant states.Source: §3.5 · pp. 18–19, Figure 7 · p. 18
  • Temporal masks discourage cross-time matches as a proxy reward, while intermediate k_m and k_c reduce spurious matches and tolerate modest speed variation (Fu et al., 2024).The mask is not a full dynamic OT formulation.Source: §3.5 · p. 19
  • AILOT learns intent embeddings and uses fixed-lookahead paired states in entropic OT, converting the plan into dense pseudo-rewards that align trajectories by intent (Bobrin et al., 2024).Its tail restriction anchors alignment to the closest expert prefix for the agent’s start state, and the paired-window cost reflects near-future evolution rather than appearance alone.Source: §3.5 · p. 19

3.6 Cross-domain imitation with Gromov-Wasserstein

  • GWIL matches expert and learner state-action occupancies across different spaces by preserving intra-domain distances, then converts the coupling into pseudo-rewards (Fickinger et al., 2022).GW is computationally heavier than standard Wasserstein OT; entropic regularisation offers scalable approximations, while keypoint-guided OT uses partial cross-domain anchors.Source: §3.6 · p. 20

3.7 Summary

  • The section groups imitation OT into online occupancy matching, offline trajectory-based reward labelling, temporal or intent-aware alignment, and cross-domain relational matching.Online methods use Wasserstein or entropic OT objectives; offline methods convert trajectory couplings into per-step pseudo-rewards for standard offline RL.Source: §3.7 · p. 20
  • Poor ground costs or weak temporal alignment can produce misleading rewards; representational problems, unstable Lipschitz enforcement, and overly diffuse entropic regularisation are dominant failure modes.Source: §3.7 · p. 20

4 Policy optimisation and regularisation

This section examines how optimal-transport geometry shapes policy optimisation and regularisation, including direct policy updates and critic learning. It covers transport between policy-induced measures and the modelling or regularisation of return and Q-value distributions, with transport costs defining structured changes.

4.1 Classical policy optimisation

  • Trust regions stabilize local surrogate policy updates; TRPO’s average-KL constraint lacks a monotonicity guarantee, and PPO clipping imposes no hard likelihood-ratio bound.The monotonic-improvement result relies on a calibrated penalty on the maximum state-wise KL divergence; average-KL updates are described as a heuristic.Source: §4.1 · p. 22

4.2 Wasserstein trust regions for policies

  • Wasserstein trust regions replace KL with action-cost-aware transport; Song et al. (2023) prove WPO monotonic improvement with exact advantages and bound degradation from inaccurate estimates.The bound includes -2ϵ/(1-γ); SPO converges to WPO as entropy regularization vanishes, and both converge to tabular optimal value under the stated finite-space, reward, support, and decaying-multiplier assumptions. OT-TRPO handles Gaussian continuous actions through scalar dual optimization, while WPR anchors RLHF policies to a reference using an entropy-regularized Wasserstein penalty.Source: §4.2 · pp. 22–24

4.3 Behaviour-guided regularisation in latent spaces

  • Pacchiano et al. (2020) regularize distributions over behavioral embeddings with smoothed Wasserstein distance, accommodating non-injective summaries that make likelihood-based divergences difficult to apply.BGPG uses OT-dual test functions to shape rewards, while BGES uses an entropic Wasserstein term in its fitness function.Source: §4.3 · pp. 24–25

4.4 Wasserstein gradient flows and natural gradients

  • Wasserstein gradient flows and natural gradients make transport geometry the update metric itself, through proximal steps or approximated Wasserstein information matrices.Zhang et al. (2018) apply gradient flows to policy-related measures; Moskovitz et al. (2021) propose low-rank and kernel approximations for WNPG and WNES. Richemond & Maginnis (2017) provide a continuous-time interpretation restricted to bandits under strong regularity assumptions.Source: §4.4 · pp. 25–26
  • WPO/SPO use Wasserstein or Sinkhorn trust regions, while Pfau et al. (2025) derive a Wasserstein trust region but implement it with KL regularisation.The survey frames trust regions, behavioural regularisers, and Wasserstein natural gradients as ways to choose policy-space geometry and keep updates small while improving return.Source: §4.4 · p. 26

4.5 Wasserstein-trained critics

  • Wasserstein-trained critics affect policy updates indirectly by changing the advantage estimates used in actor-critic surrogate objectives.They model return distributions, Q-value posteriors, or changes in empirical Q-value distributions rather than modifying the policy update rule itself.Source: §4.5 · pp. 26, 28
  • QR-DQN (Dabney et al., 2018b) approximates W1 with quantile regression, while SinkhornDRL uses Sinkhorn divergence for particle-based returns.The passage reports contraction guarantees and empirical improvements over non-distributional baselines for both methods; SinkhornDRL supports multi-dimensional rewards, with guarantees tied to specific assumptions, costs, and operators.Source: §4.5 · p. 27
  • Wasserstein Q-learning updates Q-value posteriors by 2-Wasserstein barycentres, and WAC (Likmeta et al., 2023) uses posterior uncertainty to guide exploration.WAC updates Gaussian Q-posteriors toward target distributions and uses an optimistic estimate from their upper tail; Roy et al. (2025) propose MMD barycentres as an alternative for model-free uncertainty propagation.Source: §4.5 · pp. 27–28
  • WAVE adds an adaptive entropy-regularised Wasserstein term between successive empirical Q-value distributions to discourage abrupt critic changes.The transport term smooths critic and policy dynamics while retaining the standard TD loss.Source: §4.5 · p. 28

4.6 Summary

  • OT in policy optimisation covers trust regions, trajectory-embedding regularisation, Wasserstein gradient-flow updates, and critics trained on return or Q-value distributions.These methods redefine a small update through transport geometry; the summary notes that KL/PPO-style updates are usually simpler when action geometry is weak or uninformative.Source: §4.6 · p. 28

5 Model-based RL and Wasserstein world models

Model-based RL uses OT to measure transition and trajectory errors in ways that reflect control relevance, while the chosen ground cost determines which errors matter. The survey covers value-aware model fitting, Markov-structured bisimulation comparisons, and OT-trained latent world models, alongside practical extensions.

5.5 Summary

  • The survey distinguishes model fitting from robust ambiguity sets: the former fits learned dynamics to data, while the latter hedges against transition uncertainty.It notes that these modelling choices also recur in offline settings exposed to distribution shift and compounding error.Source: §5 · p. 29, §5.5 · p. 32

5.1 Model-based RL in brief

  • Model errors that compound through planning can distort values, so OT replaces purely pointwise losses with discrepancies shaped by state geometry and control relevance.An isotropic Euclidean cost may treat equally sized errors alike even when one is harmless and the other makes a grasp fail; task-sensitive geometry must be encoded in the cost.Source: §5 · p. 28, §5.1 · p. 29

5.2 Value-aware model learning and Wasserstein losses

  • A Wasserstein loss with a value-relevant state metric links model fitting to worst-case squared value error over a Lipschitz class, rather than prediction accuracy alone.The equivalence requires a Lipschitz-MDP assumption; for a smaller test-function class, the relation is an upper bound.Source: §5.2 · p. 30

5.3 Bisimulation metrics and Markov-structured transport

  • Bisimulation metrics use transport between next-state laws inside a fixed-point state metric, comparing states by rewards and their future transition distributions.Markovian couplings preserve process structure by making each next-step pairing depend only on the current paired states, enabling Bellman-style OT computations.Source: §5.3 · pp. 30–31

5.4 Wasserstein auto-encoded and OT-based latent world models

  • WAE-MDPs and OT-MDPs train latent models by aligning real and latent trajectory distributions, with local reward and transition losses supporting bisimulation-like properties.OT-MDP costs can use features or value differences, and the framework gives conditions for bisimilarity and planning in latent space.Source: §5.4 · pp. 31–32

5.5 Summary

  • Across model-based methods, the central choice is the geometry for comparing next-state laws or trajectory traces, often defined in learned features for high-dimensional observations.The survey also identifies entropic OT/Sinkhorn as a scalable minibatch option and partial or unbalanced OT as a possible fit for missing outcomes, outliers, or mass mismatch.Source: §5.5 · p. 32

6 Offline policy optimisation and regularisation

Offline RL learns from fixed data, where policy departures can lead critics into unsupported regions and biased value estimates. The survey examines OT behaviour constraints and relaxed transport formulations as geometric ways to stay near data while accommodating suboptimal or mismatched mass.

6 Offline policy optimisation and regularisation

  • Offline policy shifts into poorly covered regions force critic extrapolation and can badly bias value estimates; OT constraints measure policy deviation using action-space geometry.Unlike density-ratio penalties, OT remains informative under support mismatch when a ground cost is defined, assigning lower cost to similar actions and higher cost to dissimilar ones.Source: §6 · p. 33

6.1 Standard offline RL framework

  • Behaviour regularisation constrains the learned policy relative to the dataset-generating policy, while pessimistic value estimation instead makes out-of-distribution actions receive conservative values.The survey describes policy constraints and pessimism as two broad offline-RL approaches; OT-based policy constraints can also induce implicit pessimism.Source: §6.1 · pp. 33–34

6.2 OT-based policy constraints, partiality, and unbalanced formulations

  • BRAC uses a modular divergence penalty in the actor objective, critic target, or both; its Wasserstein variant compares per-state action laws with a Lipschitz-constrained dual critic.The survey reports that simple behaviour regularisation is the essential ingredient for stable offline learning, while several design complexities prove unnecessary.Source: §6.2 · pp. 34–35
  • Minibatch Sinkhorn losses are useful surrogates, but their gradients are not unbiased estimates of population Wasserstein gradients and the loss may lack distance properties.Averaging per-minibatch gradients estimates the minibatch objective; it does not turn that objective into the full-distribution transport loss.Source: §6.2 · p. 35
  • Partial and unbalanced OT relax exact mass matching so offline learning need not copy all suboptimal behaviour; unbalanced formulations penalize mass creation or destruction.The survey notes that minibatch unbalanced losses can yield more robust gradient estimates than balanced OT, while excessive relaxation can hurt when demonstrations are near-optimal.Source: §6.2 · p. 36

6.3 Summary

  • Open directions include data-driven choices of costs and partiality or unbalanced parameters, and theory comparable to KL-regularised offline frameworks.The survey also suggests combining transition ambiguity sets or OT world models with offline learning, and using OT-based data filtering or behavioural embeddings.Source: §6.3 · p. 37

7 Robustness, safety, and transfer under distribution shift

OT frames robustness, safety, and transfer as geometric questions about which distribution shifts or correspondences matter. Across these applications, the ground metric and scale parameters shape the resulting method and its conservatism.

7.1 Distributionally robust control with Wasserstein ambiguity sets

  • Wasserstein ambiguity sets let nature replace nominal transitions with nearby models, while the radius δ and ground metric determine the adversarial shifts and conservatism of minimax control.Kantorovich duality yields a finite-dimensional convex program for finite-support nominal distributions; first-order methods target improved scalability.Source: §7.1 · p. 38

7.2 Safety control via visitation regularisation

  • Safety shaping transports policy visitation toward or away from a reference risk distribution, but a poorly specified safety prior can push learning in the wrong direction.The transport penalty acts as soft reward shaping rather than a hard constraint.Source: §7.2 · p. 39

7.3 Action uncertainty scores via entropic OT

  • Entropic OT measures action uncertainty through transport between current- and target-value-induced action distributions, penalising actions whose value estimates shift erratically across updates.Sinkhorn iterations compute the coupling, and β controls the uncertainty penalty.Source: §7.3 · pp. 39–40

7.4 Preference shift and Wasserstein DRO in RLHF

  • Wasserstein DRO treats empirical human preferences as nominal and optimises against plausible preference shifts; OT alignment also transports preference or reward distributions directly.An open direction is extending these mechanisms to temporally extended preferences and trajectory-level feedback.Source: §7.4 · p. 40

7.5 Reward ambiguity and Wasserstein structure

  • Wasserstein geometry can quantify dispersion among reward functions consistent with expert behaviour, but the result depends on the reward representation and ground metric.The ambiguity measure is not invariant to reward shaping unless the representation map encodes shaping equivalence.Source: §7.5 · p. 41

7.6 Transfer

  • For preference transfer across tasks with a shared action space, Liu et al. (2024) use Gromov–Wasserstein alignment between trajectory sets to transfer source preference labels to target pairs.GW compares intra-trajectory relations when state observations differ across tasks.Source: §7.6 · p. 42
  • Transfer methods also align source and target feature distributions with Wasserstein critics, filter offline source transitions by transport compatibility, or match state occupancies to goal sequences.For goal alignment, full occupancy patterns reduce the myopia of earlier goal-conditioned methods; unbalanced OT showed no significant gain in that setting.Source: §7.6 · pp. 41–42

7.7 Summary

  • Across these applications, OT turns robustness and transfer into geometric comparisons, with the ground metric and scale parameters determining plausible shifts and conservatism.Poor calibration of the ambiguity radius, safety prior, or feature metric can undermine these methods.Source: §7.7 · p. 43

8 Other applications

Beyond the survey’s main RL applications, OT supplies transport-based measures for diversity, coordination, curricula, and other settings. These applications compare distributions using geometries chosen for the task, but remain less settled in some settings.

8.1 Multiple policies or skills

  • For skill discovery, Wasserstein separation rewards policies for visiting state-space regions not claimed by other skills, while WSEP also promotes novelty against mixtures of previously learned skills.Projected Wasserstein discrepancies can mitigate dimensionality and make costs task-aware; sliced Wasserstein improves computational efficiency for WSEP.Source: §8.1 · pp. 43–44
  • WQDIL uses Wasserstein adversarial training in a Wasserstein auto-encoder latent space to improve reward-learning stability for diverse imitation from limited demonstrations.Its latent distribution matching is adapted to align expert and policy latent distributions.Source: §8.1 · p. 44

8.2 Multi-agent RL

  • In multi-agent RL, OT can act on joint occupancies, per-agent marginals, or learned embeddings; Sinkhorn barycentres provide a shared visitation target for decentralised agents.Joint distributions capture coordination but scale poorly, whereas marginals are cheaper but may miss agent interactions; differing action spaces require a common space or explicit embeddings for the barycentre.Source: §8.2 · p. 45, §8.5 · p. 47

8.3 Curricula via distributional interpolation or constraints

  • Curriculum methods use OT to choose intermediate task distributions that stay close to targets while remaining solvable, trading target proximity against current-policy feasibility.Some formulations rely on a performance-Lipschitz assumption with an unknown constant treated as a hyperparameter.Source: §8.3 · pp. 45–46

8.4 Further OT-RL intersections

  • Other applications use OT within generative models, to assess learning efficiency from successive occupancy changes, and as an intrinsic control signal for wide-coverage behaviour.The section also describes Wasserstein-robust contextual bandits, dynamic-OT counterfactual identifiability, and OT-based fusion of task-specific models.Source: §8.4 · p. 47

8.5 Summary

  • These specialised applications remain less settled, and multi-agent methods face a key choice between joint distributions that capture coordination but scale poorly and cheaper marginals that may miss interactions.Some methods trade statistical efficiency or wall-clock cost for more structured notions of diversity, alignment, or evaluation.Source: §8.5 · p. 47

9 Discussion

The discussion synthesizes recurring OT-RL design trade-offs, practical guidance, and open problems across applications, rather than recataloguing method-level details.

9.1 Cross-cutting trade-offs

  • W1 critics stay informative under weak overlap but are fragile, while entropic OT is smoother yet biased; geometry, temporal granularity, and uncalibrated hyperparameters remain consequential choices.Feature shift can degrade performance; no unified calibration procedure exists, and calibration with dependent, adaptive high-dimensional data and approximate models remains underdeveloped.Source: §9.1 · pp. 47–48

9.2 When OT helps—and when it does not

  • OT is most useful with weak overlap, task-relevant geometry, explicit transport plans, or heterogeneous domains, but may add overhead when simpler divergences or proximity baselines suffice.When demonstrations overlap well or action geometry is uninformative, KL- or JS-based methods may be cheaper and adequate.Source: §9.2 · pp. 48–49

9.3 Practitioner’s guide

  • Treat application-specific defaults as first experiments; prioritize task-relevant feature costs, normalized and scaled units, and trajectory masks or context windows over raw observations or unconstrained couplings.Finite-support linear programming scales rapidly; sample-based and entropic or sliced approximations trade tractability for sampling error or modified objectives.Source: §9.3 · pp. 49–50

9.4 Open problems

  • Open priorities are controlled tests isolating coupling, scalable long-horizon transport, principled calibration for mass mismatch, and theory separating representation from optimization error.The broader goal is to determine when transport geometry changes the problem for the better rather than mainly adding machinery.Source: §9.4 · pp. 50–51

Appendix: notation and summary tables

This appendix defines notation used throughout the survey and summarizes representative OT methods across RL applications, organizing them by the distributions compared, OT formulation, role, and implementation notes while collating reported outcomes and limitations.

Notation and summary tables

  • The notation table defines entropy and IRL regularization terms, embedding and pushforward maps, empirical state marginals, return distributions, the distributional Bellman operator, trajectory space, and latent behaviour space.For example, µπ = Φ#Pπ is the behavioural distribution, and Zπ(s, a) denotes a return distribution.Source: Table 2 · p. 64, Appendix: notation and summary tables · p. 65
  • Table 3 presents representative OT uses for imitation learning and reward inference, including primal W1, weighted BC solved by OT, sliced multi-marginal W2, visual embeddings with temporal mask/context costs, and state-action occupancy measures via relations.Source: Table 3 · p. 65, Appendix: notation and summary tables · p. 66
  • Policy optimization uses W2 gradient-flow and natural-gradient views, while Wasserstein policy optimization discretizes a gradient flow and projects updates to a parametric policy class.WGF treats policy-induced measures such as trajectory laws; WNG (Moskovitz et al., 2021) uses kernel or low-rank approximations, and the implemented policy-optimization algorithm retains KL regularization.Source: Table 4 · p. 66, Appendix: notation and summary tables · p. 67
  • Distributional RL and value-learning methods apply W1 quantile regression or entropic Sinkhorn OT to return distributions, and use Wasserstein distances or barycentres in critic updates.QR-DQN (Dabney et al., 2018b) uses fixed quantiles; SinkhornDRL uses particle-set Sinkhorn iterations, while WQL uses Wasserstein TD Q-posteriors and target barycentres.Source: Appendix: notation and summary tables · p. 67
  • For safe RL and preference alignment, OT supplies visitation-based uncertainty scores, action-uncertainty scoring, and Wasserstein ambiguity sets over preference distributions.Safe-RL examples use Sinkhorn OT maps to define safer behavioural policies and OT-derived costs in value-based updates; Xu et al. (2025) and Sahu & Wells (2025) use Wasserstein DRO, while OTPO (Li et al., 2025) weights tokens according to OT couplings.Source: Appendix: notation and summary tables · p. 69
  • For transfer and preference alignment, the cited methods use a negated goal-conditioned value function as a possibly asymmetric ground cost, GW matching of intra-trajectory structure, or WGAN-style W1 adversarial alignment.The goal-conditioned cost may be asymmetric and therefore is not a Wasserstein metric; WDGRL is identified as domain adaptation.Source: Appendix: notation and summary tables · p. 69
  • Control-shift methods use Wasserstein feature alignment, intermediate-domain reward shaping, or OT-based data filtering and reweighting.WAPPO (Roy & Konidaris, 2021) minimizes feature-distribution W1 during policy learning; Liu et al. (2025) uses a Wasserstein reward to select intermediate domains, and OTDF (Lyu et al., 2025) filters or reweights state-action-next-state triples. The appendix also tabulates offline policy optimization and control-shift methods.Source: Table 7 · p. 68, Appendix: notation and summary tables · p. 69

Experimental results

  • SinkhornDRL (Sun et al., 2024) converges more slowly early but improves aggregate human-normalised score over QR-DQN and MMD-DQN, with largest gains for larger action spaces and multidimensional rewards.The study reports approximately 20% more compute than MMD-DQN and sensitivity to ε, iteration count, and particle number.Source: Appendix: notation and summary tables · p. 72
  • TemporalOT (Fu et al., 2024) achieves higher success than online OTR, ADS, GAIfO, BC, and an oracle task reward without task rewards; ablations support both its context cost and mask.It requires online interaction and uses a temporally informed proxy rather than a full dynamic-OT formulation.Source: Appendix: notation and summary tables · p. 72
  • Sliced multi-marginal OT yields higher mean episodic reward than concatenated demonstrations under both pole-length and mass diversity.Evaluation is limited to Pendulum-v0 and CartPole-v0, and slicing approximates the geometry.Source: Appendix: notation and summary tables · p. 72
  • WPO/SPO reports monotonic improvement with exact advantages against KL trust regions and global convergence in tabular settings with a decaying multiplier.The tabular guarantee requires finite spaces, non-negative rewards, and full initial-state support; advantage error weakens the bound by 2ϵ/(1 −γ).Source: Appendix: notation and summary tables · p. 72
  • An explicit W2 Brenier map via an ICNN removes the inner min-max loop of adversarial Wasserstein penalties while retaining W2 geometry.Expressiveness is limited by the finite ICNN parameterisation, and minibatch OT introduces bias.Source: Appendix: notation and summary tables · p. 72
  • GWIL recovers optimal behaviour in two tests, but mismatched cheetah-to-walker spaces produce forward or backward locomotion depending on the seed.Its identification is only up to isometries, which can yield suboptimal behaviour; reported head-to-head baseline comparisons are limited.Source: Appendix: notation and summary tables · p. 72
  • WR2L (Abdullah et al., 2019) improves worst-case return over non-robust training under the same physical-parameter perturbations.The evaluation covers MuJoCo tasks with perturbed physical parameters and uses a curvature-based approximation around a reference model.Source: Appendix: notation and summary tables · p. 72

Related papers

Loading 2610.01413v1…