Invariant Adversarial Imitation Learning from Visual Inputs

Generalization across environments is critical when using imitation learning algorithms in real-world applications. In this paper, we propose an invariant model-based adversarial imitation learning (IMAIL) method to improve generalization. IMAIL develops a variational dynamics model providing rich auxiliary objectives for efficiently learning compact state representations. The latent representations are then regularized using mutual information constraints, guaranteeing that they are insensitive to environmental changes. Based on such representations, we utilize model-based adversarial imitation learning to mimic expert behavior in the latent space. As a result, the learned policies are well generalized in unseen environments. We conduct experiments with several vision-based control tasks to demonstrate the performance of IMAIL. Experimental results show that IMAIL significantly outperforms existing baselines and successfully achieves expert-level performance in all unseen test environments.