Primal Wasserstein Imitation Learning

TL;DR

Imitation learning often hinges on reward design or adversarial training. PWIL reframes IL as distribution matching between the agent and expert state-action distributions using the primal Wasserstein distance, and derives an offline reward from a greedy-coupling upper bound. It delivers strong sample efficiency on MuJoCo locomotion tasks and even solves Humanoid with a single demonstration, while providing a true behavioral metric via $\mathcal{W}_1$. The method extends to vision-based imitation by learning a metric offline, avoiding the instability of min-max IL and requiring only a small number of hyperparameters. Overall, PWIL demonstrates that principled distance-based IL with offline rewards can yield robust, data-efficient imitation with broad practical impact.

Abstract

Imitation Learning (IL) methods seek to match the behavior of an agent with that of an expert. In the present work, we propose a new IL method based on a conceptually simple algorithm: Primal Wasserstein Imitation Learning (PWIL), which ties to the primal form of the Wasserstein distance between the expert and the agent state-action distributions. We present a reward function which is derived offline, as opposed to recent adversarial IL algorithms that learn a reward function through interactions with the environment, and which requires little fine-tuning. We show that we can recover expert behavior on a variety of continuous control tasks of the MuJoCo domain in a sample efficient manner in terms of agent interactions and of expert interactions with the environment. Finally, we show that the behavior of the agent we train matches the behavior of the expert with the Wasserstein distance, rather than the commonly used proxy of performance.

Figures (13)

• Figure 1: Illustration of the difference between the a) greedy coupling and the b) optimal coupling. We present an MDP where we drop the dependency on the action. The state space is $\mathbb{R}$ and the metric associated is the Euclidean distance. We note the states visited by the policy: $s^\pi_1, s^\pi_2, s^\pi_3$ and the states visited by the expert: $s^e_1, s^e_2, s^e_3$. When the policy encounters the state $s^\pi_2$, and because we do not know $s^\pi_3$ yet, the greedy coupling strategy consists in coupling it with $s_2^e$ although the optimal coupling strategy would be to couple it with $s^e_3$. Note that the total cost (omitting the constant coupling multiplication factor) with the greedy coupling is $7$ whereas the total cost with the optimal coupling is $5$ . This highlights that the optimal coupling needs knowledge about the whole policy's trajectory to be derived.
• Figure 2: Mean and standard deviation return of the evaluation policy over 10 rollouts and 10 seeds, reported every 10k environment steps. The return is in term of the environment's original reward. Top row: 1 demonstration, bottom row: 11 demonstrations.
• Figure 3: Mean of the Wasserstein distance between the state-action distribution of the evaluation policy and the state-action distribution of the expert over 10 rollouts and 10 seeds. Agents were trained with 11 demonstrations. For PWIL, we include the upper bound on the Wasserstein distance based on the greedy coupling defined in Equation equation \ref{['eq:upper_bound']}.
• Figure 4: Mean and standard deviation of the evaluation performance of PWIL and variants at the 1M environment interactions mark (2.5M for Humanoid). Results are computed over 10 seeds and 10 episodes for each seed.
• Figure 5: Visual rendering of the door opening task.