Latent Wasserstein Adversarial Imitation Learning

This paper introduces Latent Wasserstein Adversarial Imitation Learning (LWAIL), a novel framework that leverages a pre-trained, dynamics-aware latent space to achieve expert-level performance using only one or a few state-only expert demonstrations, thereby overcoming the data and action requirements of traditional imitation learning methods.

The Gist

The Big Problem: Learning to Dance Without Seeing the Steps

Imagine you want to learn a complex dance routine.

• Traditional Reinforcement Learning (RL) is like trying to learn the dance by bumping into furniture and hoping a "ding" sound tells you when you did it right. It takes forever and requires a perfect "ding" (reward signal) every time.
• Imitation Learning (IL) is better: You watch a master dancer and try to copy them.
• The Catch: Usually, to copy the dance, you need to see both the dancer's moves (the state) and their footwork instructions (the action). But what if you only have a video of the dancer's body moving, and the audio (the specific footwork commands) is missing?
• The Real Problem: Even just watching the video is hard if the video is short, blurry, or if the dancer is doing something very specific that you've never seen before. Most AI methods fail if they don't have hours of high-quality video.

The Solution: LWAIL (The "Intuitive" Copycat)

The authors propose a new method called LWAIL. Think of it as teaching an AI to dance by giving it one single, short video of an expert, plus a tiny bit of "random flailing" data to help it understand the physics of the room.

Here is how it works, broken down into three simple steps:

1. The "Random Flailing" Phase (Pre-training)

Before the AI tries to copy the expert, it needs to understand the physics of the world.

• The Analogy: Imagine you are dropped into a dark room with a ball. You don't know the rules yet. So, you just throw the ball around randomly for a few seconds. You notice: "If I push the ball hard to the left, it hits the wall and bounces back. If I push it gently, it rolls slowly."
• In the Paper: The AI uses a tiny amount of random data (just 1% of what other methods need) to train a special "Intention Conditioned Value Function" (ICVF). This is like the AI building a mental map of how the world works. It learns that "State A" is close to "State B" not because they look similar on a screen, but because you can easily get from A to B in the real world.

2. The "New Language" (The Latent Space)

This is the paper's biggest innovation.

• The Problem: Most AI methods measure "distance" between states using a ruler (Euclidean distance).
  • Example: In a maze, two points might be 1 meter apart in a straight line (Euclidean distance). But if there is a wall between them, you have to walk 100 meters to get there. A ruler says they are close; reality says they are far.
• The Fix: LWAIL translates the world into a new language (a "Latent Space"). In this new language, the "distance" between two points isn't about how they look; it's about how hard it is to get from one to the other.
  • Analogy: Imagine a map where cities are placed not by their geographic location, but by how long it takes to drive between them. In this map, two cities separated by a mountain might look far apart, even if they are geographically close. This map understands the dynamics (the rules of movement).

3. The "Adversarial Dance-Off" (Imitation)

Now, the AI tries to copy the expert using this new, smart map.

• The Setup: The AI (the student) and a "Discriminator" (a strict judge) play a game.
  • The Judge looks at the expert's video and the student's video. It tries to tell them apart.
  • The Student tries to move in a way that makes the Judge think, "Hey, this looks just like the expert!"
• The Twist: The Judge doesn't just look at the pixels; it looks at the Latent Space. Because the map understands the physics (the walls, the gravity, the momentum), the student learns to move efficiently rather than just looking similar.

Why is this a Big Deal?

1. It works with almost no data: You only need one single video of an expert to learn a complex task. Other methods need dozens or hundreds.
2. It ignores the "noise": Even if the expert video is a bit shaky or the environment is noisy, LWAIL figures out the underlying rules of movement.
3. It solves the "Wall" problem: By using the ICVF map, the AI realizes that just because two states look close, it doesn't mean you can jump between them. It learns the true difficulty of the task.

The Summary Metaphor

Imagine you are trying to learn to drive a car in a city you've never visited.

• Old Methods: You are given a GPS that only shows straight-line distances. You try to drive from Point A to Point B, but you keep crashing into buildings because the GPS didn't tell you about the walls.
• LWAIL: Before you start driving, you spend 5 minutes walking around the block randomly. You learn where the walls are and how the streets connect. Then, you are given a single video of a pro driver. Because you already understand the city's layout (the dynamics), you can watch that one video and immediately start driving like a pro, avoiding all the walls.

In short: LWAIL teaches the AI to understand the rules of the game before trying to play the game, allowing it to learn complex skills from very few examples.

Technical Summary

1. Problem Statement

Imitation Learning (IL) allows agents to learn from expert demonstrations, bypassing the need for hand-crafted reward functions. However, traditional IL faces two major bottlenecks:

1. Data Scarcity: High-quality expert demonstrations (especially those including actions) are often expensive or impossible to acquire.
2. Observation-Only Limitations: Many real-world scenarios only provide state-only observations (Imitation Learning from Observations, or LfO), lacking expert actions.

Existing Adversarial Imitation Learning (AIL) methods for LfO often rely on $f$-divergences (e.g., KL, JS) or the Kantorovich-Rubinstein (KR) dual of the Wasserstein distance.

• $f$-divergences require "distribution coverage" (the learner's distribution must be supported on the same set as the expert's), which is theoretically restrictive and numerically unstable with limited data.
• Wasserstein-based methods (using the KR dual) are more robust to distribution shifts but typically rely on the Euclidean distance as the ground metric between states. The authors argue that Euclidean distance fails to capture the environment's dynamics (e.g., in a maze, two states might be spatially close but dynamically unreachable from one another), leading to suboptimal learning signals.

Core Challenge: How can we learn a distance metric that encodes environmental dynamics using only a minimal amount of low-quality, state-only data, enabling an agent to learn expert-level policies from just a single state-only trajectory?

2. Methodology: Latent Wasserstein Adversarial Imitation Learning (LWAIL)

LWAIL proposes a two-stage framework that replaces the naive Euclidean metric with a dynamics-aware latent space metric derived from an Intention Conditioned Value Function (ICVF).

Stage 1: Pre-training (Dynamics-Aware Embedding)

• Data Source: A small, unstructured dataset of random state transitions ($I$), which can be as small as 1% of the online rollout data. No expert actions or rewards are required.
• ICVF Training: The authors train an ICVF model (Ghosh et al., 2023) on this random data. ICVF learns a value function $V(s, s^+, z)$ representing the likelihood of reaching a future state $s^+$ (outcome) from a current state $s$ with a specific intention/goal $z$.
• Latent Representation: The value function is factorized as $V_\theta(s, s^+, z) = \phi_\theta(s)^T T_\theta(z) \psi_\theta(s^+)$. The term $\phi_\theta(s)$ serves as the state embedding.
• Theoretical Insight: The authors prove (Theorem 3.1) that in near-deterministic MDPs, the state-pair occupancy $d^\pi_{ss}(s, s')$ is approximately a linear combination of the embedding $\phi_\theta(s)$. This implies that the Euclidean distance in this latent space $\|\phi_\theta(s) - \phi_\theta(s')\|_2$ effectively captures the reachability and transition dynamics of the environment.

Stage 2: Imitation (Adversarial Training)

• Frozen Embedding: The learned embedding $\phi_\theta$ is frozen.
• Modified Objective: The standard Wasserstein AIL objective is modified to operate in the latent space. The discriminator $f$ and the policy $\pi$ now operate on $\phi(s)$ and $\phi(s')$ instead of raw states.
• Objective Function:
  $$\min_\pi \max_{\|f\|_L \leq 1} \left( \mathbb{E}_{(s,s') \sim d^\pi_{ss}} [f(\phi(s), \phi(s'))] - \mathbb{E}_{(s,s') \sim d^E_{ss}} [f(\phi(s), \phi(s'))] \right)$$
• Training Loop:
  1. The agent interacts with the environment using policy $\pi$.
  2. A pseudo-reward is generated: $r(s, s') = \sigma(-f(\phi(s), \phi(s')))$, where $\sigma$ is a sigmoid function to stabilize the reward scale.
  3. An off-policy RL algorithm (TD3) updates the policy using these pseudo-rewards.
  4. The discriminator $f$ is updated to distinguish between expert state pairs and agent state pairs in the latent space.

3. Key Contributions

1. Dynamics-Aware Ground Metric: The paper introduces the first direct remedy for the geometric limitations of Euclidean distance in Wasserstein occupancy matching. By leveraging ICVF pre-trained on minimal random data, it constructs a latent space where Euclidean distance reflects true environmental dynamics (reachability).
2. Data Efficiency: LWAIL achieves expert-level performance using only a single state-only expert trajectory. It does not require expert actions, nor does it require high-quality offline datasets for pre-training (random data suffices).
3. Theoretical and Empirical Validation:
   • Proves that ICVF embeddings align the structure of Wasserstein optimization with transition dynamics.
   • Demonstrates that the learned metric significantly outperforms vanilla Euclidean metrics and other contrastive learning embeddings (like CURL or PW-DICE) in imitation tasks.

4. Experimental Results

The method was evaluated on MuJoCo locomotion tasks (Hopper, HalfCheetah, Walker2D, Ant) and Maze navigation tasks (Maze2D, AntMaze) from the D4RL benchmark.

• Performance vs. Baselines: LWAIL consistently outperformed a wide range of baselines, including:
  • Classic IL (GAIL, AIRL, BC).
  • Wasserstein-based methods (WDAIL, IQ-learn, PWIL).
  • Observation-only methods (BCO, GAIfO, DIFO, OPOLO).
  • In many tasks, LWAIL achieved scores comparable to or exceeding methods that had access to expert actions, despite only using state-only data.
• Robustness to Noise:
  • Initial State Perturbation: In navigation tasks with Gaussian noise injected into initial states, LWAIL maintained high performance, whereas methods without ICVF embeddings failed catastrophically.
  • Transition Noise: LWAIL remained robust even when Gaussian noise was added to the environment's transition dynamics.
• Ablation Studies:
  • Removing the ICVF embedding (using raw Euclidean distance) caused a significant drop in performance.
  • Using different downstream RL algorithms (PPO, DDPG) resulted in lower performance compared to TD3.
  • The method is robust to the size and quality of the pre-training random dataset (performing well even with 10K random transitions).

5. Significance

LWAIL addresses a critical gap in Reinforcement Learning and Imitation Learning: learning complex behaviors from extremely sparse, state-only data.

• Democratization of RL: By removing the need for expert actions and high-quality offline datasets, LWAIL makes imitation learning feasible in robotics and real-world applications where collecting expert actions is difficult (e.g., human demonstrations via video) and where data is scarce.
• Geometric Understanding: The work highlights that the choice of distance metric in Wasserstein-based IL is not just a hyperparameter but a fundamental component of learning dynamics. It shifts the paradigm from "matching distributions in raw space" to "matching distributions in a dynamics-aware latent space."
• Practical Efficiency: The pre-training stage is computationally cheap and requires minimal data, making the overall pipeline highly efficient for real-world deployment.

In summary, LWAIL demonstrates that a simple two-stage process—pre-training a dynamics-aware embedding on random data followed by standard adversarial imitation in that latent space—can solve the state-only imitation problem with unprecedented data efficiency and robustness.