# Wasserstein Distance guided Adversarial Imitation Learning with Reward Shape Exploration

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

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Placed in Imitation learning methods, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 9 of its 23 references that are in the map.
- First posted on arXiv 5 Jun 2020.

## Placed in

- Offline RL and imitation learning from data / [Imitation learning methods](https://mlsanity.com/offlinerl/imitation-learning-methods/)

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

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Learning Robust Rewards with Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/1710.11248/) (2017)
- [Guided Cost Learning: Deep Inverse Optimal Control via Policy Optimization](https://mlsanity.com/paper/1603.00448/) (2016)
- [Variational Discriminator Bottleneck: Improving Imitation Learning, Inverse RL, and GANs by Constraining Information Flow](https://mlsanity.com/paper/1810.00821/) (2018)
- [Multi-Agent Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1807.09936/) (2018)
- [Dialogue Generation: From Imitation Learning to Inverse Reinforcement Learning](https://mlsanity.com/paper/1812.03509/) (2018)

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

- [Reward is enough for convex MDPs](https://mlsanity.com/paper/2106.00661/) (2021)
- [Imitation Learning with Sinkhorn Distances](https://mlsanity.com/paper/2008.09167/) (2020)
- [Online Apprenticeship Learning](https://mlsanity.com/paper/2102.06924/) (2021)
- [Hierarchical Generative Adversarial Imitation Learning with Mid-level Input Generation for Autonomous Driving on Urban Environments](https://mlsanity.com/paper/2302.04823/) (2023)
- [Imitation Learning from Observation through Optimal Transport](https://mlsanity.com/paper/2310.01632/) (2023)
- [Wasserstein Distance Maximizing Intrinsic Control](https://mlsanity.com/paper/2110.15331/) (2021)
- [Offline Imitation from Observation via Primal Wasserstein State Occupancy Matching](https://mlsanity.com/paper/2311.01331/) (2023)
- [Generative Adversarial Imitation Learning for End-to-End Autonomous Driving on Urban Environments](https://mlsanity.com/paper/2110.08586/) (2021)
- [Reward function shape exploration in adversarial imitation learning: an empirical study](https://mlsanity.com/paper/2104.06687/) (2021)
- [How does Your RL Agent Explore? An Optimal Transport Analysis of Occupancy Measure Trajectories](https://mlsanity.com/paper/2402.09113/) (2024)
- [Unsupervised Behavioral Compression: Learning Low-Dimensional Policy Manifolds through State-Occupancy Matching](https://mlsanity.com/paper/2603.27044/) (2026)
- [Latent Wasserstein Adversarial Imitation Learning](https://mlsanity.com/paper/2603.05440/) (2026)
- [Boosting Reinforcement Learning with Verifiable Rewards via Randomly Selected Few-Shot Guidance](https://mlsanity.com/paper/2605.15012/) (2026)

## Abstract (arXiv)

The generative adversarial imitation learning (GAIL) has provided an adversarial learning framework for imitating expert policy from demonstrations in high-dimensional continuous tasks. However, almost all GAIL and its extensions only design a kind of reward function of logarithmic form in the adversarial training strategy with the Jensen-Shannon (JS) divergence for all complex environments. The fixed logarithmic type of reward function may be difficult to solve all complex tasks, and the vanishing gradients problem caused by the JS divergence will harm the adversarial learning process. In this paper, we propose a new algorithm named Wasserstein Distance guided Adversarial Imitation Learning (WDAIL) for promoting the performance of imitation learning (IL). There are three improvements in our method: (a) introducing the Wasserstein distance to obtain more appropriate measure in the adversarial training process, (b) using proximal policy optimization (PPO) in the reinforcement learning stage which is much simpler to implement and makes the algorithm more efficient, and (c) exploring different reward function shapes to suit different tasks for improving the performance. The experiment results show that the learning procedure remains remarkably stable, and achieves significant performance in the complex continuous control tasks of MuJoCo.