Visual Adversarial Imitation Learning using Variational Models (VMAIL)

This is the official implementation of the NeurIPS 2021 paper.



VMAIL simultaneously learns a variational dynamics model and trains an on-policy
adversarial imitation learning algorithm in the latent space using only model-based
rollouts. This allows for stable and sample efficient training, as well as zero-shot
imitation learning by transfering the learned dynamics model


Get dependencies:

conda env create -f vmail.yml
conda activate vmail
cd robel_claw/robel
pip install -e .

To train agents for each environmnet download the expert data from the provided link and run:

python3 -u --logdir .logdir --expert_datadir expert_datadir

The training will generate tensorabord plots and GIFs in the log folder:

tensorboard --logdir ./logdir


If you find this code useful, please reference in your paper:

      title={Visual Adversarial Imitation Learning using Variational Models}, 
      author={Rafael Rafailov and Tianhe Yu and Aravind Rajeswaran and Chelsea Finn},
      journal={Neural Information Processing Systems}


View Github