Visual Adversarial Imitation Learning using Variational Models (VMAIL)

This is the official implementation of the NeurIPS 2021 paper.

• Project website
• Research paper
• Datasets used in the paper

Method

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

Instructions

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 vmail.py --logdir .logdir --expert_datadir expert_datadir

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

tensorboard --logdir ./logdir

Citation

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

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