Latent Wasserstein Adversarial Imitation Learning

TL;DR

This work proposes Latent Wasserstein Adversarial Imitation Learning (LWAIL), a novel adversarial imitation learning framework that focuses on state-only distribution matching that outperforms prior Wasserstein-based IL methods and prior adversarial IL methods, achieving better results across various tasks.

Abstract

Imitation Learning (IL) enables agents to mimic expert behavior by learning from demonstrations. However, traditional IL methods require large amounts of medium-to-high-quality demonstrations as well as actions of expert demonstrations, both of which are often unavailable. To reduce this need, we propose Latent Wasserstein Adversarial Imitation Learning (LWAIL), a novel adversarial imitation learning framework that focuses on state-only distribution matching. It benefits from the Wasserstein distance computed in a dynamics-aware latent space. This dynamics-aware latent space differs from prior work and is obtained via a pre-training stage, where we train the Intention Conditioned Value Function (ICVF) to capture a dynamics-aware structure of the state space using a small set of randomly generated state-only data. We show that this enhances the policy's understanding of state transitions, enabling the learning process to use only one or a few state-only expert episodes to achieve expert-level performance. Through experiments on multiple MuJoCo environments, we demonstrate that our method outperforms prior Wasserstein-based IL methods and prior adversarial IL methods, achieving better results across various tasks.

Figures (12)

• Figure 1: Illustrating our motivation for a better distance metric and an outline of the algorithm. Panel (a) illustrates a case where the Euclidean distance between states is not a good metric: state B is closer to expert state C, but it is apparently less desirable than more distant state A, as it cannot reach C. To fix, we use random data and ICVF to find a more meaningful embedding space, as shown in the lower half. Panel (b), together with the lower half of panel (a), shows our pre-training stage: we first train ICVF to obtain $\phi(s)$, which serves as a reward for our agent in the online stage shown in panel (c). Fire indicates trainable modules, and snowflakes indicate frozen modules.
• Figure 2: t-SNE visualization of the same trajectory in the original state space and the embedding latent space. The color of the points represents the ground-truth reward of the state (greener is higher). States connected by lines are adjacent in the trajectory. We observe that an ICVF-trained embedding provides a much more dynamics-aware metric than the Euclidean distance. For Hopper and Ant environment embeddings, see Appendix \ref{['sec:morevis']}.
• Figure 3: Results on Maze2D. Panel (a) shows the environment setup with the start (green) and the goal (red). Panel (b) and (c) illustrate the reward distribution without and with ICVF embedding (orange represents the wall). Panel (d) shows the learning curves: LWAIL vs. TD3 with ground-truth rewards. We note that the reward is more dynamics-aware with LWAIL.
• Figure 4: Performance on the MuJoCo environments. Overall, our method consistently delivers strong performance across all tasks. For LWAIL, the number of gradient steps matches the number of online interaction samples. For clarity, we present training curves for main baselines here; please refer to Fig. \ref{['fig:main_result_full']} for the complete version.
• Figure 5: An illustration of "case 1" (green states) and "case 2" (orange states) in our proof, and the $z$ is the intention (i.e. goal). The transparent states are longer paths, which are unlikely be visited with a converged policy and $\gamma<1$ in a near-deterministic MDP.

Theorems & Definitions (7)

• Theorem 3.1
• Remark B.1
• Remark B.2
• Remark C.1
• Theorem D.1
• proof
• Remark D.2