# Offline Imitation from Observation via Primal Wasserstein State Occupancy Matching

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 43 of its 61 references that are in the map.
- First posted on arXiv 2 Nov 2023.

## Builds on (most cited first, 30 of 43)

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [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)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [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)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Deep Reinforcement Learning for Autonomous Driving: A Survey](https://mlsanity.com/paper/2002.00444/) (2020)
- [Behavioral Cloning from Observation](https://mlsanity.com/paper/1805.01954/) (2018)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Is Pessimism Provably Efficient for Offline RL?](https://mlsanity.com/paper/2012.15085/) (2020)
- [DualDICE: Behavior-Agnostic Estimation of Discounted Stationary Distribution Corrections](https://mlsanity.com/paper/1906.04733/) (2019)
- [Discriminator-Actor-Critic: Addressing Sample Inefficiency and Reward Bias in Adversarial Imitation Learning](https://mlsanity.com/paper/1809.02925/) (2018)
- [IQ-Learn: Inverse soft-Q Learning for Imitation](https://mlsanity.com/paper/2106.12142/) (2021)
- [A Divergence Minimization Perspective on Imitation Learning Methods](https://mlsanity.com/paper/1911.02256/) (2019)
- [Generative Adversarial Imitation from Observation](https://mlsanity.com/paper/1807.06158/) (2018)
- [Imitation Learning via Off-Policy Distribution Matching](https://mlsanity.com/paper/1912.05032/) (2019)
- [The Surprising Effectiveness of Representation Learning for Visual Imitation](https://mlsanity.com/paper/2112.01511/) (2021)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Primal Wasserstein Imitation Learning](https://mlsanity.com/paper/2006.04678/) (2020)
- [OptiDICE: Offline Policy Optimization via Stationary Distribution Correction Estimation](https://mlsanity.com/paper/2106.10783/) (2021)
- [Learning Generalizable Robotic Reward Functions from "In-The-Wild" Human Videos](https://mlsanity.com/paper/2103.16817/) (2021)
- [Discriminator-Weighted Offline Imitation Learning from Suboptimal Demonstrations](https://mlsanity.com/paper/2207.10050/) (2022)
- [State Alignment-based Imitation Learning](https://mlsanity.com/paper/1911.10947/) (2019)
- [Offline Learning from Demonstrations and Unlabeled Experience](https://mlsanity.com/paper/2011.13885/) (2020)
- [Off-Policy Imitation Learning from Observations](https://mlsanity.com/paper/2102.13185/) (2021)
- [Generative Adversarial User Model for Reinforcement Learning Based Recommendation System](https://mlsanity.com/paper/1812.10613/) (2018)

## Built on by (most cited first, 5 of 5)

- [Domain-Invariant Per-Frame Feature Extraction for Cross-Domain Imitation Learning with Visual Observations](https://mlsanity.com/paper/2502.02867/) (2025)
- [Learning from Observation: A Survey of Recent Advances](https://mlsanity.com/paper/2509.19379/) (2025)
- [Imitation from Observations with Trajectory-Level Generative Embeddings](https://mlsanity.com/paper/2601.00452/) (2026)
- [Latent Wasserstein Adversarial Imitation Learning](https://mlsanity.com/paper/2603.05440/) (2026)
- [Differentiating Bisimulation Metrics: A Framework for Parametric Markov Chain Fitting via Bicausal Optimal Transport](https://mlsanity.com/paper/2609.37239/) (2026)

## Abstract (arXiv)

In real-world scenarios, arbitrary interactions with the environment can often be costly, and actions of expert demonstrations are not always available. To reduce the need for both, offline Learning from Observations (LfO) is extensively studied: the agent learns to solve a task given only expert states and task-agnostic non-expert state-action pairs. The state-of-the-art DIstribution Correction Estimation (DICE) methods, as exemplified by SMODICE, minimize the state occupancy divergence between the learner's and empirical expert policies. However, such methods are limited to either f-divergences (KL and chi2) or Wasserstein distance with Rubinstein duality, the latter of which constrains the underlying distance metric crucial to the performance of Wasserstein-based solutions. To enable more flexible distance metrics, we propose Primal Wasserstein DICE (PW-DICE). It minimizes the primal Wasserstein distance between the learner and expert state occupancies and leverages a contrastively learned distance metric. Theoretically, our framework is a generalization of SMODICE, and is the first work that unifies f-divergence and Wasserstein minimization. Empirically, we find that PW-DICE improves upon several state-of-the-art methods. The code is available at https://github.com/KaiYan289/PW-DICE.