Optimal Transport Meets RL: A Survey from Distribution Matching to Robust Decisions

1 Introduction

Reinforcement learning agents optimize long-term return while constantly comparing policy-induced distributions: imitation learning matches learner and expert visitation; trust-region methods limit policy change; offline RL constrains deviation from data; robust control accounts for transition shifts; curriculum learning adjusts task distributions. Standard divergences (KL, JS) compare densities pointwise and can fail when supports barely overlap — KL may become infinite, JS discriminators saturate. Optimal transport (OT) introduces task geometry via a ground cost function, measuring the minimum cost to move probability mass from one distribution to another. This allows meaningful distances even with disjoint supports. However, simply "using OT" does not guarantee correct geometry: wrong costs, excessive entropy regularization, unstable dual networks, or ignoring temporal structure can produce misleading signals. The survey organizes literature by OT's algorithmic role — reward signal, regularizer, trust-region constraint, distributional uncertainty set, model loss — and distinguishes core RL works from transfer and preference-learning applications. Four guiding questions run through the review: (1) What role does OT play in the objective? (2) Which distributions are compared? (3) Which transport form is adopted? (4) How is temporal structure handled?

2 Preliminaries

2.1 RL Foundations

Markov decision processes comprise states, actions, transitions, rewards, discount factor, and initial state distribution. A policy π(a|s) induces a full trajectory distribution. The discounted occupancy measure weights state-action visitation probabilities by discount and normalizes to a distribution, describing where the policy stays and what it does. It obeys flow-conservation constraints. Four key objects are distinguished: per-state action distributions for local policy constraints; occupancy measures for global visitation comparison; trajectory distributions or embeddings for behavioral comparison; return distributions for stochastic outcomes. Collapsing a trajectory into an empirical state distribution discards order — same length and frequency can represent different action sequences.

2.2 OT Foundations

Given source μ, target ν, and cost c(x,y), OT finds a coupling Γ with marginals μ, ν minimizing total transport cost. Couplings allow mass splitting; deterministic transport maps require single destinations. If cost is a distance raised to power p, the p-th root yields the p-Wasserstein distance. The cost function is a modeling choice — changing which point pairs are cheap changes the optimal plan. Pixel proximity, action proximity, dynamics proximity, and task-value proximity cannot be mixed arbitrarily. A simple example: two point masses at different locations — OT returns their distance, while pointwise density comparison struggles to express displacement magnitude.

Exact discrete OT solves a linear program but scales poorly. Entropic regularization trades smoothness for stable, solvable plans via Sinkhorn iterations; too large an entropy coefficient oversmooths the plan. Sinkhorn divergence subtracts self-transport bias so identical distributions have zero discrepancy, but it should not be universally called a metric. W₁ dual form trains Lipschitz potentials; W₂ can represent transport maps as gradients of convex potentials under conditions. Sliced Wasserstein projects high-dimensional comparisons onto multiple 1D lines to reduce cost. When no direct pointwise cost exists across domains, Gromov–Wasserstein (GW) compares intra-domain relational structures. For outliers, missing coverage, or mass mismatch, balanced OT's full-mass matching may be unsuitable: partial OT transports only a specified mass; unbalanced OT softens marginal constraints via penalty terms. Both require calibrating transported mass or penalty strength.

3 Imitation Learning and Reward Inference

3.1 Understanding Imitation via Occupancy Measures

Imitation learning minimizes distance between learner and expert state-action occupancy measures. If only expert observations (no actions) are available, the target shifts to state or transition representations. State visitation closeness does not universally guarantee action or return equivalence; conclusions depend on dynamics and assumptions.

3.2 Why OT Helps

Early in training, learner and expert visitation regions often barely overlap. OT provides geometric signals pointing toward expert states rather than merely indicating separability. Benefit depends on whether observation features preserve reachability and goal information; visually similar but unreachable states may appear close in a poor representation.

3.3 Online Adversarial Reward Shaping

Methods like WAIL replace discriminator-based divergences with a Wasserstein objective, alternating learner updates with dual potential training, then converting the potential into a reward. Others use raw entropic OT or Sinkhorn objectives. Advantages: geometric signal under weak overlap. Costs: extra optimization, Lipschitz enforcement, minibatch overfitting. The transport potential here differs from the value network estimating returns — they serve distinct roles.

3.4 Offline Reward Labeling

OTR, PWIL, and similar methods compute a coupling between learner/data and expert empirical measures, then distribute matching costs as per-step pseudo-rewards for standard RL pipelines. This avoids continuous adversarial training, but each state's label may depend on the whole trajectory match. The same state in different trajectories can receive different rewards — a natural consequence of this labeling scheme.

3.5 Temporal and Intent-Aware Extensions

Treating trajectories as unordered state sets can match reversed or differently paced trajectories. Two correction paths: (1) add near-diagonal bands, temporal masks, and context windows to the coupling to constrain matching order; (2) incorporate short-term dynamics, goal intent, or reachability features into the cost to redefine behavioral proximity. Too narrow windows restrict legitimate speed variations; too wide windows revert to erroneous matches.

3.6 Cross-Domain Imitation

When robot morphology or observation spaces differ, no natural Euclidean distance exists between source and target states. GW matches by preserving intra-domain relational structure, enabling heterogeneous imitation. If a reliable shared representation already exists, standard OT is often simpler; GW brings more complex optimization and risk of relational mismatch.

3.7 Section Summary

OT in imitation can provide online rewards or offline labels; comparing occupancy measures differs from comparing empirical trajectories. Cost design and temporal handling often affect reward quality more than the solver name. Multi-expert data requires careful handling of mode mixing.

4 Policy Optimization and Regularization

4.1 Classic Policy Optimization

TRPO uses KL trust regions; PPO clips probability ratios. They focus on how the probability of the same action changes. OT further considers how far probability mass moves in action space, suitable when actions have meaningful geometry.

4.2 OT Trust Regions

Constraining the Wasserstein distance between old and new action distributions at a given state makes moving probability between similar actions cheaper. In continuous control, small joint changes should cost less than direction reversals. However, small local action distances do not directly guarantee small trajectory-level changes; constraint radius and state sampling still affect stability. If actions lack useful geometry, adding OT may not help.

4.3 Latent-Space Behavioral Guidance

Some methods map trajectories to behavioral embeddings and constrain the induced embedding distributions to approach a target or maintain diversity. This captures whole-behavior similarity better than per-state-action constraints, but embeddings may lose critical temporal information or suffer distribution shift, requiring checks on which differences are preserved.

4.4 Gradient Flows and Natural Gradients

Wasserstein gradient flows view policy updates as proximal optimization in distribution space; Wasserstein natural gradients change the local geometry of parameter updates. Particle and neural transport maps approximate such updates, but restricted directions or numerical solutions do not equal exact continuous gradient flows. Stability and error analysis under function approximation remain limited.

4.5 Distributional Value Learning with Transport Distances

Distributional RL estimates the full return distribution , not just its expectation; 1D return comparison via Wasserstein leverages quantiles, with QR-DQN as a representative. SinkhornDRL extends to particle representations and multi-dimensional rewards. Another line, Wasserstein Q-learning , uses Wasserstein barycenters to update Q-value posterior distributions , expressing estimation uncertainty. The former captures return stochasticity; the latter captures epistemic uncertainty about values — they are not interchangeable. Contraction guarantees depend on specific operators, projections, and assumptions.

4.6 Section Summary

On the policy side, OT controls action or behavioral change; on the value side, it serves as a distributional loss. Before adopting, clarify the transport object; if only optimization stability is needed, classic KL/PPO baselines should be retained.

5 Model-Based RL and Wasserstein World Models

5.1 Model-Based RL

Learning transitions improves data efficiency, but multi-step prediction amplifies local errors. Model quality is not just low average prediction error; errors must be evaluated by their effect on planning and control. For example, the same position error may be harmless along a grasp tolerance direction but cause failure orthogonal to it.

5.2 Value-Aware Model Learning

Comparing true and learned next-state distributions, OT can penalize deviations according to control-relevant geometry. When the value function satisfies a corresponding Lipschitz condition, Wasserstein error bounds expected value error. This link requires the condition to hold; OT does not automatically discover grasp tolerances — directional weighting, reachability, or value-aware costs must encode them.

5.3 Bisimulation and Markov-Structured Transport

Bisimulation-style state distances combine immediate reward differences and next-state distribution differences to judge long-term decision equivalence. Embedding OT into the distance fixed point preserves dynamical structure. Compared to unordered state matching, Markov-structured couplings respect valid transitions but increase solution difficulty.

5.4 Latent World Models

Wasserstein autoencoders and OT latent models match states, transitions, or trajectories in compressed spaces, learning control-relevant abstractions. High-dimensional visual observations typically need feature costs; pure reconstruction may retain decision-irrelevant details, while over-compression loses planning information.

5.5 Section Summary

Model learning uses OT to fit transitions; robust control uses OT to guard against transition shifts — different purposes. Minibatch Sinkhorn, partial/unbalanced matching, and flow matching offer scaling ideas, but representation quality, coverage gaps, and multi-step error must be validated separately.

6 Offline Policy Optimization and Regularization

6.1 Offline RL Basics

Offline RL uses only a fixed dataset. When a policy selects actions rarely seen in data, value networks may extrapolate erroneously. Constraining the policy to stay close to behavior data mitigates this; OT defines "close" as transport cost in action space. Interpolating between 4 cm and 5 cm movements in the same direction is cheaper than reversing direction under suitable geometry, but proximity in distance does not prove action safety or high value .

6.2 Policy Constraints, Partial and Unbalanced Matching

BRAC modularizes behavioral divergence as a policy or value objective; BRAC-W uses a dual Wasserstein penalty. Q-DOT employs convex potentials / input-convex neural networks to construct W₂ transport maps, providing a discriminator-free route but requiring convexity and representation conditions. Other methods interpret updates as distributional gradient flows; their approximations and finite budgets still limit deviation magnitude.

With suboptimal or corrupted data, forcing full behavior replication hinders improvement. Partial Policy Learning (PPL) adopts a one-sided incomplete transport motivation, relaxing the target marginal constraint — not equivalent to simply keeping a fixed percentage of samples. Unbalanced methods soften marginals, allowing mass reweighting. Neither automatically identifies "good data"; selection still depends on cost, value estimates, and relaxation strength.

6.3 Section Summary

Offline OT methods play three roles: limiting actor deviation, generating rewards for data, and filtering/reweighting data. Minibatch objectives differ from full-distribution Wasserstein objectives; unbiased gradients for minibatch objectives do not imply unbiasedness for population OT. When data coverage is narrow and value extrapolation poor, distance tools alone cannot remove the fundamental limitation.

7 Robustness, Safety, and Transfer

7.1 Wasserstein Uncertainty Sets

Building a Wasserstein ball around the nominal transition distribution and optimizing for the worst case within the ball trains agents to handle plausible dynamics shifts. The cost defines which changes are plausible; the radius sets conservatism. Too large a radius harms nominal task return; too small fails to cover deployment shifts.

7.2 Occupancy Measures and Safe Control

Pushing the policy's state visitation distribution toward a safe reference distribution adds obstacle-avoidance preferences to the reward. This is a geometry-aware soft mechanism, not a hard safety constraint; an inaccurate reference distribution may push the policy into unsafe regions. Cited evidence comes from small discrete tasks; applicability depends on the original studies.

7.3 Action Uncertainty Scoring

Without a safety reference, one can compare the current Q-value distribution with the action distribution induced by the update target, using mass changes in the Sinkhorn coupling to identify unstable actions. The score adjusts exploration and selection, reflecting value estimation changes, not direct accident probability.

7.4 Preference Shift and Human Feedback

Preference data distribution shifts can be modeled as uncertainty sets, enabling distributionally robust optimization of DPO, REBEL, etc. Other methods use OT to align preferences, rewards, or token representations. The survey lists these as adjacent directions relevant to sequential decisions; their evidence does not directly translate to unified improvements in classic RL environments.

7.5 Reward Ambiguity

The same expert behavior can be explained by multiple reward functions. Mapping these rewards into a geometric space and measuring their Wasserstein distance quantifies reward interpretation uncertainty; it provides a metric perspective but does not eliminate the identifiability issue in inverse RL.

7.6 Transfer

OT can align shared feature distributions across source and target domains, or use matching scores to filter/weight data; GW is an option for spatially heterogeneous domains. Cross-domain geometry, coverage, and cost noise determine transfer quality. Hard threshold filtering risks discarding useful data due to cost errors; soft reweighting is a safer initial experiment.

7.7 Section Summary

Robustness compares environment changes; safety shapes visitation; preference alignment handles feedback shifts; transfer finds cross-domain correspondences. All use OT but with different objects and guarantees; radii, penalty weights, and reference geometries all need calibration.

8 Other Applications

8.1 Multi-Policy and Skills

Skill discovery can reward Wasserstein separation between different policies' visitation distributions, encouraging coverage of distinct regions — opposite direction to imitation's attraction. Whether diversity helps depends on features and tasks; merely maximizing distance may yield meaningless behaviors.

8.2 Multi-Agent RL

OT can coordinate policies, allocate resources, or pull visitation distributions toward a common barycenter. Core trade-off: comparing team joint distributions preserves interactions but scales poorly; comparing individual marginals is cheaper but may miss coordination. Even with decentralized policies, each agent's visitation is influenced by all policies through environment dynamics.

8.3 Curriculum Generation

Using OT interpolation, barycenters, or proximity constraints between easy and target tasks can generate curricula that simultaneously constrain "closeness to final goal" and "solvability by current policy". Geometrically interpolated tasks are not guaranteed physically feasible or learnable; task semantics and reachability must be considered.

8.4 Further Cross-Directions

The paper also discusses Sinkhorn assignment in combinatorial problems, generative trajectory alignment, exploration bonuses, learning efficiency evaluation, and connections to contextual bandits and causal flows. Some works belong to neighboring fields; the authors view them as potential inspirations for OT-RL, not yet mature universal solutions.

8.5 Section Summary

These directions are more scattered than imitation, offline constraints, and policy optimization. Evaluation should check extra statistical and computational costs, and whether the transport object sufficiently captures desired diversity, coordination, or task progress.

9 Discussion

9.1 Cross-Application Trade-offs

First, matching sharpness vs. stability : W₁ dual provides geometric signals but relies on Lipschitz control; entropic regularization eases solving but introduces smoothing and bias. Second, representation vs. time : feature spaces improve high-dimensional comparison, but bad embeddings distort task geometry; temporal masks avoid misalignment but restrict matching freedom. Third, hyperparameter calibration : entropy coefficients, uncertainty radii, partial transport mass, and marginal penalties lack cross-task data-driven rules.

9.2 When Is OT Suitable?

OT is especially suitable for tasks with weak support overlap, clear action/state geometry, need for explicit coupling to assign rewards, or heterogeneous space relational matching. When distributions already overlap well, geometry lacks task meaning, samples are extremely scarce, or compute budget is tight, the extra optimization may not pay off. The paper notes that certain proximity measures and lightweight temporal baselines can match OT reward labeling on broad benchmarks; comparisons should fix representation, reward transformation, and downstream algorithm to disentangle gains from representation vs. coupling.

9.3 Practical Guidelines

The original Table 1 offers initial experimental suggestions: for online imitation, try feature-space Sinkhorn divergence; for offline reward labeling, combine entropic regularization with temporal masks; for policy trust regions, compare per-state action distributions; for offline constraints, keep Wasserstein behavioral penalty as a baseline; for world models, compare control-relevant next-state features; for robust control, calibrate uncertainty radius. The authors emphasize these are starting points, not proven universal defaults.

Implementation: first decide transport object and units, then scale costs; set entropy coefficient relative to typical pairwise costs. Sinkhorn can use log-domain iterations with warm starts for slowly changing problems; dual network capacity and Lipschitz control must suit batch size. For long trajectories, try windows and masks first, and log sample size, solver cost, and task return. Full OT, sliced, entropic, minibatch, and constrained couplings may change the optimization problem — be explicit about approximations when comparing.

9.4 Open Problems

When is the coupling structure indispensable? Need to isolate representation, cost, temporal, and solver components to find conditions where coupling is truly necessary under multi-modality, mass mismatch, or severe temporal misalignment.

How to scale and amortize computation? Current methods often recompute OT per update. Learning reusable maps or potentials, using projection approximations, can reduce cost; structured transport for high-dimensional long-horizon control remains hard to scale.

How to treat mass mismatch as the norm? Real data is messy and incomplete; partial transport mass and marginal penalties should be calibrated based on coverage, value error, and uncertainty, not just by tuning.

How to build theory under realistic function approximation? Wasserstein trust regions, natural gradients, and proximal updates need more complete stability, sample complexity, and error analyses, distinguishing cost representation errors from numerical optimization errors.

The paper concludes that OT provides a geometric language for comparing distributions in RL. The next step is to verify when this geometry improves decisions and to develop calibratable, scalable methods. The appendix includes a notation table, per-section method cross-reference, and environment-level evidence tables; full definitions and references are available at the paper link.

Code example

来源：专知
本文
约7000字
，建议阅读
15
分钟
OT 在目标中起什么作用？比较哪些分布？采用哪种传输形式？如何处理时间结构？