Research map / Reinforcement learning / Imitation learning
Offline Imitation from Observation via Primal Wasserstein State Occupancy Matching
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PW-DICE, an offline learning-from-observations method that minimizes the primal 1-Wasserstein distance between learner and expert state occupancies with a contrastively learned distance metric, generalizing SMODICE.
Problem
Offline imitation from observations (LfO), where the agent has state-only expert demonstrations and task-agnostic state-action data; prior DICE methods are limited to f-divergences or to Wasserstein with Kantorovich-Rubinstein duality, which restricts the distance metric to Euclidean.
Approach
Optimize the primal Wasserstein distance between state occupancies, smoothed with KL pessimism regularizers so it becomes an unconstrained convex dual, allowing an arbitrary distance metric (a SMODICE-style reward plus a contrastive InfoNCE reachability embedding distance); the policy is recovered by weighted behavior cloning.
Key results
- In the tabular random MDP experiments, PW-DICE (with or without regularizer) achieves the lowest regret across expert and task-agnostic dataset sizes and noise levels, with the gap widening as task-agnostic data grows.
- On MuJoCo, antmaze and Franka kitchen testbeds, PW-DICE matches or beats baselines including SMODICE and LobsDICE, while OTR struggles on most environments.
- Ablations show distance metrics with a single component (e.g., Euclidean implied by Rubinstein duality) fail, while the combined reward plus contrastive distance works well.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace SMODICE or LobsDICE with PW-DICE | Tabular random MDP · Random MDP across expert dataset sizes, task-agnostic dataset sizes, and noise levels | regret | Helped |
| replace the baseline method with PW-DICE | MuJoCo locomotion benchmarks · Hopper, HalfCheetah, Ant, and Walker2d | normalized average reward | Mixed |
| use PW-DICE instead of OTR | Offline imitation learning · MuJoCo environments | average reward | Helped |
| use the proposed distance metric instead of a single-component distance metric (ablation) | PW-DICE · MuJoCo distance-metric ablation | average reward | Helped |
| use the combination of cosine similarity and R(s) instead of a single-component distance metric (ablation) | PW-DICE · MuJoCo distance-metric ablation | average reward | Helped |
| use only the Euclidean distance instead of the proposed distance metric (ablation) | PW-DICE · MuJoCo distance-metric ablation | average reward | Hurt |
| use only cosine similarity instead of the proposed distance metric (ablation) | PW-DICE · MuJoCo distance-metric ablation | average reward | Hurt |
| use only R(s) instead of the proposed distance metric (ablation) | PW-DICE · MuJoCo distance-metric ablation | average reward | Hurt |
| vary ε1 and ε2 across the tested settings (ablation) | PW-DICE · MuJoCo | average reward | Mixed |
| increase ε1 (ablation) | PW-DICE · MuJoCo | average reward | Hurt |
| use χ²-divergence instead of KL-divergence (ablation) | PW-DICE · MuJoCo | average reward | Hurt |
| use PW-DICE instead of SMODICE under distorted state representations | Offline imitation learning · MuJoCo environments with distorted state representations | average reward | Helped |
Surprising findings
- The Wasserstein-based OTR baseline, despite using primal Wasserstein distance, struggles on most MuJoCo environments.
- The choice of distance metric has a significant impact on Wasserstein-based performance, with Euclidean and cosine metrics giving very different results for OTR.
Limitations
- Logsumexp over minibatches yields biased gradient estimates, a potential instability in harder environments.
- Evaluation mostly on locomotion and manipulation testbeds; authors note potential misuse risks.
- Distance metric relies on a learned contrastive embedding and hyperparameters (alpha, beta, epsilon_1, epsilon_2) that need tuning; epsilon_1 should be small.
Experimental setup
Datasets: D4RL MuJoCo (hopper, halfcheetah, ant, walker2d) antmaze Franka kitchen random tabular MDPs (following LobsDICE)
Benchmarks: random tabular MDP regret D4RL MuJoCo antmaze Franka kitchen
Metrics: regret normalized average reward total variation distance of occupancies
Baselines: SMODICE LobsDICE ReCOIL ORIL OTR DWBC BC
Builds on: SMODICE LobsDICE OTR CURL InfoNCE
Closest papers by meaning
- ODICE: Revealing the Mystery of Distribution Correction Estimation via Orthogonal-gradient Update (ICLR 2024 · similarity 0.88)
- GO-DICE: Goal-Conditioned Option-Aware Offline Imitation Learning via Stationary Distribution Correction Estimation (AAAI 2024 · similarity 0.88)
- Behavior Cloning is Not All You Need: The Optimality of On-Policy Distillation for Noisy Expert Feedback (NeurIPS 2026 · similarity 0.86)
- Zero-Shot Offline Imitation Learning via Optimal Transport (ICML 2025 · similarity 0.86)
- Offline RL with Observation Histories: Analyzing and Improving Sample Complexity (ICLR 2024 · similarity 0.86)
- Imitation from Observations with Trajectory-Level Generative Embeddings (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.