PaperlayerView exact revision on arXiv

Source-linked AI summary

Primal Wasserstein Imitation Learning

Robert Dadashi, Léonard Hussenot, Matthieu Geist, Olivier Pietquin

arXiv:2006.04678v2cs.LGstat.ML

TL;DR

Imitation Learning seeks to match expert behavior when task rewards are difficult to specify. The paper introduces PWIL, which uses an offline reward derived from the primal Wasserstein distance between expert and agent state-action distributions. PWIL recovers expert behavior across continuous-control tasks with little tuning and evaluates similarity directly with Wasserstein distance rather than task performance.

  • Problem

    Imitation Learning needs effective reward-free ways to match expert behavior when task rewards are difficult to specify, while Wasserstein metrics require a nontrivial state-action metric.

  • Method

    PWIL derives an offline reward from an upper bound on the primal Wasserstein distance between expert and agent state-action distributions, avoiding an inner minimization problem.

  • Results

    PWIL recovers expert behavior comparably to existing state-of-the-art methods across continuous-control environments, including with a single demonstration in the considered settings.

  • Takeaways & Limitations

    Wasserstein distance provides a direct measure of expert-agent behavioral similarity for imitation-learning settings where the task reward is unavailable or cannot be specified.

  • Takeaways & Limitations

    The approach assumes that a metric over state-action pairs can be defined, although the paper notes that constructing such a metric in an MDP is nontrivial.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Imitation Learning addresses tasks where rewards are difficult to specify by learning from expert demonstrations. PWIL minimizes a primal Wasserstein objective with an offline reward, reducing optimization complexity while recovering expert behavior across continuous-control settings.

  • Imitation Learning learns policies from fixed expert demonstrations when reinforcement-learning rewards are hard to specify or too sparse for practical use.Behavioral Cloning directly matches expert behavior, whereas Inverse Reinforcement Learning first seeks to recover an unknown reward.
  • Iterative and adversarial Inverse Reinforcement Learning methods can reduce sample efficiency by repeatedly estimating rewards and training agents.Earlier methods require multiple Markov decision process solver calls, while adversarial methods interleave reward learning with agent learning.
  • PWIL measures the Wasserstein distance between expert and agent state-action distributions using its primal formulation.The primal formulation avoids the minmax optimization problem and requires little fine-tuning.
  • PWIL introduces an offline reward based on an upper bound of the primal Wasserstein distance and recovers expert behavior comparably to state-of-the-art methods.The method operates in low-data regimes and is reported to make Humanoid run with one subsampled demonstration.

2 BACKGROUND AND NOTATIONS

The paper formulates imitation learning in episodic Markov decision processes using empirical state-action distributions for policies and demonstrations. Wasserstein comparison then requires a metric over state-action pairs, whose construction is nontrivial and may be learned from demonstrations.

  • Markov decision processes: An episodic Markov decision process is specified by state and action spaces, transitions, rewards, discounting, initial-state distribution, and a finite horizon.Policies map states to action distributions, and reinforcement learning seeks to maximize expected discounted return.
  • State action distributions: A policy’s empirical state-action distribution records the successive state-action pairs visited during an episode.
  • State action distributions: Expert demonstrations define a corresponding empirical expert state-action distribution for distribution matching.
  • Wasserstein distance: The p-th order Wasserstein distance compares distributions over a metric space through couplings between their supports.For finite supports, a coupling is represented as a doubly stochastic matrix.
  • Wasserstein distance: Applying Wasserstein distance to state-action distributions requires a metric on the combined state-action space, which is nontrivial in an MDP.The paper assumes such a metric and later gives an example learned from demonstrations.

3 METHOD

PWIL minimizes the Wasserstein distance between expert and policy state-action distributions using an offline reward derived from an upper bound based on greedy coupling. The greedy coupling makes this cost computable online while retaining a direct connection to the primal Wasserstein objective.

  • 3.1 WASSERSTEIN DISTANCE MINIMIZATION: PWIL minimizes the Wasserstein distance between the policy and expert state-action distributions using its primal formulation.
  • 3.1 WASSERSTEIN DISTANCE MINIMIZATION: The method restricts the analysis to the 1-Wasserstein distance and interprets couplings as transport strategies between policy state-action pairs and expert pairs.
  • 3.2 GREEDY COUPLING: Because greedy coupling is suboptimal, its cost provides an upper bound on the Wasserstein distance that can be computed online.
  • 3.2 GREEDY COUPLING: The optimal coupling requires knowledge of the whole policy trajectory, whereas greedy coupling matches each newly observed policy pair to the closest remaining expert pair.In the illustrated example, greedy coupling has total cost 7 versus 5 for optimal coupling.
  • 3.2 GREEDY COUPLING: PWIL converts the greedy-coupling cost into a monotonically decreasing, episodic history-dependent reward derived offline from expert demonstrations.The reward is not re-optimized through an inner minimization during agent learning.
  • 3.2 GREEDY COUPLING: The resulting algorithm supports generic agents and computes episode-wide rewards with complexity O((|S| + |A|)DT + D^2).When T = D, greedy coupling reduces to selecting and removing the nearest expert state-action pair.

4 EXPERIMENTS

PWIL is evaluated on MuJoCo locomotion and visual door-opening tasks, with comparisons covering performance, sample efficiency, Wasserstein distance, and component ablations.

  • Experimental setup: PWIL is evaluated against DAC and BC on six MuJoCo locomotion environments using D4PG agents and subsampled expert demonstrations.Training uses 1M environment interactions, or 2.5M for Humanoid, with evaluation over multiple seeds and episodes.
  • Locomotion results: PWIL improves final performance over DAC on Hopper, Ant, HalfCheetah, and Humanoid, matches Walker2d, and underperforms on Reacher.On Humanoid, PWIL achieves near-optimal performance even with a single demonstration.
  • Sample efficiency: DAC is more sample-efficient than PWIL on HalfCheetah, Ant, and Reacher, while the methods perform similarly on Hopper and Walker2d.The authors attribute DAC’s interaction cost partly to reward-function tuning and experimental cycles.
  • Wasserstein evaluation: PWIL produces smaller Wasserstein distances than DAC on all environments except HalfCheetah, where they are similar, and Reacher, where PWIL is worse.The measured distance decreases during learning, and PWIL’s upper bound is reported as empirically tight.
  • Ablation study: Removing standardized distance weighting causes a significant performance drop across environments except Hopper, while omitting replay-buffer expert transitions significantly harms Walker and Ant.These ablations indicate dependence on the state-action metric and expert-transition replay initialization.
  • Visual task: PWIL extends to door opening from pixels by learning an embedding-space metric offline with Temporal Cycle-Consistency Learning.The visual task uses human demonstrations and evaluates imitation returns, embedding-space Wasserstein distance, and door-opening success rate.

5 RELATED WORK

Prior imitation-learning approaches match expert behavior through adversarial divergences, expert-support rewards, trajectories, offline reward estimation, or moment matching, but differ in assumptions and sample efficiency.

  • Adversarial IL methods match agent and expert state-action distributions using divergences such as Jensen-Shannon or Kullback-Leibler.
  • Adversarial approaches can require carefully tuned discriminators, including architecture, regularization, and learning-rate schedules, while PWIL uses two hyperparameters.
  • Expert-support methods encourage agents to remain on the expert support but may not recover the expert state-action distribution.
  • Trajectory-based methods can rely on a single demonstration or temporal alignment, whereas PWIL generalizes to randomly varying door and handle locations.
  • Offline reward-estimation methods often assume structured rewards, while PWIL uses no structural assumptions but produces a non-stationary reward dependent on visited state-actions.
  • MMD imitation learning minimizes Maximum Mean Discrepancy but requires complete agent rollouts, lacks a greedy-coupling relaxation, and is less sample-efficient because it is on-policy.

6 CONCLUSION

The paper formulates imitation learning as distribution matching and introduces an offline reward based on an upper bound of the Wasserstein distance. PWIL provides a direct expert-agent similarity measure for settings where task rewards cannot be specified, while recovering near-expert performance with little tuning.

  • PWIL frames imitation learning as matching agent and expert state-action distributions with a reward based on an upper bound of their Wasserstein distance.
  • The method provides a direct expert-agent similarity measure for imitation-learning settings where the task reward cannot be specified.
  • PWIL computes its reward offline, uses two hyperparameters, and recovers near-expert performance with one demonstration across the considered environments, including Humanoid.
  • The method also extends to visual observations by learning a metric offline from expert demonstrations with self-supervised learning.

A.1 PWIL IMPLEMENTATION

PWIL uses D4PG with specified actor, critic, optimization, replay, and reward-search settings for its implementation.

  • PWIL uses D4PG with the default ACME architecture, including four-layer actor and critic networks.
  • The actor uses 256-unit hidden layers, while the critic uses 512-, 512-, and 256-unit hidden layers with a 201-category distributional output.
  • Adam optimization uses λa = 5×10^-5 and λc = 7×10^-5, batch size 256, and gradient clipping at L2 norm 40.
  • The replay buffer has size 10^6, uses discount factor θ = 0.99 and five-step returns, is prefilled with 50000 demonstration state-action pairs, and updates every four environment interactions.
  • The DDPG hyperparameter search is documented in Table 1, while PWIL reward parameters α and β are searched over {1, 5, 10}.

A.2 DAC IMPLEMENTATION

The DAC implementation uses extensive Humanoid tuning, while evaluation reports imitation, original-environment, and Wasserstein-distance learning curves across demonstrations.

  • DAC Humanoid tuning evaluated 729 experiments, of which 7 exceeded average performance 1000 at 2.5M training environment steps.
  • Table 2 reports the tested and best hyperparameters for the Humanoid experiment with 11 demonstrations.
  • The DAC implementation uses a three-layer network with 128 and 64 hidden units, tanh output, expert-observation normalization, mean squared error, and Adam optimization.
  • Learning curves report imitation reward and original environment return over 10 rollouts and 10 seeds, evaluated every 10k steps through 1M steps.
  • Wasserstein-distance curves compare evaluation-policy and expert state-action distributions and include the greedy-coupling upper bound across one, four, and eleven demonstrations.

B ABLATION STUDY

The ablation study presents learning curves for different PWIL variants, evaluated across varying numbers of demonstrations and environment steps.

  • PWIL variants are compared through learning curves in the ablation study.
  • Evaluation performance is reported every 10k environment steps over 1M steps.
  • The evaluation varies the demonstration set size across 1, 4, and 11 demonstrations.

C INFLUENCE OF THE DIRECT RL ALGORITHM

The study compares PWIL with D4PG, TD3, and SAC using the same reward function. PWIL-SAC and PWIL-TD3 recover expert-like performance on several locomotion tasks, while SAC performs better than D4PG overall.

  • PWIL is evaluated with D4PG, TD3, and SAC on locomotion tasks using the same reward function.
  • PWIL-SAC and PWIL-TD3 recover expert-like performance on Hopper, Ant, HalfCheetah, and Reacher.
  • Walker2d shows large performance variability because some seeds approach expert performance while others do not.

D DOOR OPENING EXPERIMENTS

The door-opening experiments extend PWIL to early-terminating tasks and describe visual demonstration encoding, SAC implementation, runtime, and an example execution. Early termination states are treated as absorbing states.

  • Door opening: The door-opening environment adds early termination when the door opens, removing any incentive for survival after success.
  • Door opening: Early termination states are treated as absorbing, so demonstrations ending early assign substantial state-distribution weight to those states.
  • Visual demonstrations: Visual demonstrations at 84 × 84 resolution are mapped into a lower-dimensional latent space using the self-supervised TCC method.
  • Visual demonstrations: The encoder is selected by validation Kendall’s tau, with the best reported scores reaching 0.98 on training and 0.91 on validation.
  • Direct RL implementation: In the direct-RL comparison, SAC outperforms D4PG, using an actor, twin critic, adaptive temperature, and tuned optimization settings.
  • Runtime: PWIL’s episode-level reward computation has complexity O((|S| + |A|)DT + D^2).
Loading 2006.04678v2…