Wasserstein Adversarial Imitation Learning

TL;DR

This paper introduces Wasserstein Adversarial Imitation Learning, a method that uses optimal transport theory to improve imitation learning by enabling smooth reward functions and high sample efficiency, demonstrated in robotic tasks.

Contribution

It proposes a novel imitation learning approach leveraging optimal transport and Kantorovich potentials, enhancing reward function flexibility and sample efficiency.

Findings

Outperforms baselines in average cumulative rewards

Requires only one expert demonstration for high sample efficiency

Effective in large-scale robotic applications

Abstract

Imitation Learning describes the problem of recovering an expert policy from demonstrations. While inverse reinforcement learning approaches are known to be very sample-efficient in terms of expert demonstrations, they usually require problem-dependent reward functions or a (task-)specific reward-function regularization. In this paper, we show a natural connection between inverse reinforcement learning approaches and Optimal Transport, that enables more general reward functions with desirable properties (e.g., smoothness). Based on our observation, we propose a novel approach called Wasserstein Adversarial Imitation Learning. Our approach considers the Kantorovich potentials as a reward function and further leverages regularized optimal transport to enable large-scale applications. In several robotic experiments, our approach outperforms the baselines in terms of average cumulative rewards and shows a significant improvement in sample-efficiency, by requiring just one expert demonstration.