# Optimal Transport for Offline Imitation Learning

arXiv: https://arxiv.org/abs/2303.13971 · HTML: https://mlsanity.com/paper/2303.13971/

## Key facts

- Cited by 42 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 19 of its 37 references that are in the map.
- First posted on arXiv 24 Mar 2023.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 19 of 19)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Critic Regularized Regression](https://mlsanity.com/paper/2006.15134/) (2020)
- [Discriminator-Actor-Critic: Addressing Sample Inefficiency and Reward Bias in Adversarial Imitation Learning](https://mlsanity.com/paper/1809.02925/) (2018)
- [Imitation Learning via Off-Policy Distribution Matching](https://mlsanity.com/paper/1912.05032/) (2019)
- [Acme: A Research Framework for Distributed Reinforcement Learning](https://mlsanity.com/paper/2006.00979/) (2020)
- [Primal Wasserstein Imitation Learning](https://mlsanity.com/paper/2006.04678/) (2020)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Offline Learning from Demonstrations and Unlabeled Experience](https://mlsanity.com/paper/2011.13885/) (2020)
- [What Matters for Adversarial Imitation Learning?](https://mlsanity.com/paper/2106.00672/) (2021)
- [How to Leverage Unlabeled Data in Offline Reinforcement Learning](https://mlsanity.com/paper/2202.01741/) (2022)
- [Wasserstein Adversarial Imitation Learning](https://mlsanity.com/paper/1906.08113/) (2019)
- [Cross-Domain Imitation Learning via Optimal Transport](https://mlsanity.com/paper/2110.03684/) (2021)

## Built on by (most cited first, 30 of 42)

- [A Large Language Model-Driven Reward Design Framework via Dynamic Feedback for Reinforcement Learning](https://mlsanity.com/paper/2410.14660/) (2024)
- [What Matters to You? Towards Visual Representation Alignment for Robot Learning](https://mlsanity.com/paper/2310.07932/) (2023)
- [SEABO: A Simple Search-Based Method for Offline Imitation Learning](https://mlsanity.com/paper/2402.03807/) (2024)
- [CEIL: Generalized Contextual Imitation Learning](https://mlsanity.com/paper/2306.14534/) (2023)
- [CLUE: Calibrated Latent Guidance for Offline Reinforcement Learning](https://mlsanity.com/paper/2306.13412/) (2023)
- [Imitation Learning from Observation with Automatic Discount Scheduling](https://mlsanity.com/paper/2310.07433/) (2023)
- [PEARL: Zero-shot Cross-task Preference Alignment and Robust Reward Learning for Robotic Manipulation](https://mlsanity.com/paper/2306.03615/) (2023)
- [Rethinking Optimal Transport in Offline Reinforcement Learning](https://mlsanity.com/paper/2410.14069/) (2024)
- [Align Your Intents: Offline Imitation Learning via Optimal Transport](https://mlsanity.com/paper/2402.13037/) (2024)
- [Offline Imitation Learning with Model-based Reverse Augmentation](https://mlsanity.com/paper/2406.12550/) (2024)
- [Robot Policy Learning with Temporal Optimal Transport Reward](https://mlsanity.com/paper/2410.21795/) (2024)
- [Imitation Learning from a Single Temporally Misaligned Video](https://mlsanity.com/paper/2502.05397/) (2025)
- [Transport and Merge: Cross-Architecture Merging for Large Language Models](https://mlsanity.com/paper/2602.05495/) (2026)
- [Discovering Hierarchical Achievements in Reinforcement Learning via Contrastive Learning](https://mlsanity.com/paper/2307.03486/) (2023)
- [ENOT: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport](https://mlsanity.com/paper/2403.03777/) (2024)
- [Direct Post-Training Preference Alignment for Multi-Agent Motion Generation Models Using Implicit Feedback from Pre-training Demonstrations](https://mlsanity.com/paper/2503.20105/) (2025)
- [Offline Imitation from Observation via Primal Wasserstein State Occupancy Matching](https://mlsanity.com/paper/2311.01331/) (2023)
- [PINS: Proximal Iterations with Sparse Newton and Sinkhorn for Optimal Transport](https://mlsanity.com/paper/2502.03749/) (2025)
- [Leveraging Optimal Transport for Enhanced Offline Reinforcement Learning in Surgical Robotic Environments](https://mlsanity.com/paper/2310.08841/) (2023)
- [A Simple Solution for Offline Imitation from Observations and Examples with Possibly Incomplete Trajectories](https://mlsanity.com/paper/2311.01329/) (2023)
- [Zero-Shot Offline Imitation Learning via Optimal Transport](https://mlsanity.com/paper/2410.08751/) (2024)
- [PROF: An LLM-based Reward Code Preference Optimization Framework for Offline Imitation Learning](https://mlsanity.com/paper/2511.13765/) (2025)
- [Surgical Task Automation Using Actor-Critic Frameworks and Self-Supervised Imitation Learning](https://mlsanity.com/paper/2409.02724/) (2024)
- [Robust Offline Imitation Learning Through State-level Trajectory Stitching](https://mlsanity.com/paper/2503.22524/) (2025)
- [DPOT: A DeepParticle method for Computation of Optimal Transport with convergence guarantee](https://mlsanity.com/paper/2506.23429/) (2025)
- [Offline Reinforcement Learning with Wasserstein Regularization via Optimal Transport Maps](https://mlsanity.com/paper/2507.10843/) (2025)
- [From Novelty to Imitation: Self-Distilled Rewards for Offline Reinforcement Learning](https://mlsanity.com/paper/2507.12815/) (2025)
- [Re:Frame -- Retrieving Experience From Associative Memory](https://mlsanity.com/paper/2508.19344/) (2025)
- [On Discovering Algorithms for Adversarial Imitation Learning](https://mlsanity.com/paper/2510.00922/) (2025)
- [ReGIL: Retrieval-Guided Imitation Learning from a Single Demonstration](https://mlsanity.com/paper/2606.09381/) (2026)

## Abstract (arXiv)

With the advent of large datasets, offline reinforcement learning (RL) is a promising framework for learning good decision-making policies without the need to interact with the real environment. However, offline RL requires the dataset to be reward-annotated, which presents practical challenges when reward engineering is difficult or when obtaining reward annotations is labor-intensive. In this paper, we introduce Optimal Transport Reward labeling (OTR), an algorithm that assigns rewards to offline trajectories, with a few high-quality demonstrations. OTR's key idea is to use optimal transport to compute an optimal alignment between an unlabeled trajectory in the dataset and an expert demonstration to obtain a similarity measure that can be interpreted as a reward, which can then be used by an offline RL algorithm to learn the policy. OTR is easy to implement and computationally efficient. On D4RL benchmarks, we show that OTR with a single demonstration can consistently match the performance of offline RL with ground-truth rewards.