Latent Wasserstein Adversarial Imitation Learning
This work is significant for researchers and practitioners in reinforcement learning and robotics who face challenges with data availability and quality for imitation learning, offering a method to achieve expert-level performance with minimal state-only demonstrations.
This paper introduces Latent Wasserstein Adversarial Imitation Learning (LWAIL), a new framework for imitation learning that addresses the need for large, high-quality action-labeled demonstrations by focusing on state-only distribution matching. LWAIL uses a dynamics-aware latent space, learned via a pre-training stage with a small set of randomly generated state-only data, to enable expert-level performance with only one or a few state-only expert episodes.
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.