Original PyTorch implementation of MoDem: Accelerating Visual Model-Based Reinforcement Learning with Demonstrations by

Nicklas Hansen, Yixin Lin, Hao Su, Xiaolong Wang, Vikash Kumar, Aravind Rajeswaran (Meta AI, UC San Diego)

[Paper] [Website]

Our model-based method, MoDem, solves challenging visuo-motor control tasks with sparse rewards and high-dimensional action spaces in 100K interaction steps given only 5 demonstrations.

If you use this repo in your research, please consider citing the paper as follows:

@article{hansen2022modem,
  title={MoDem: Accelerating Visual Model-Based Reinforcement Learning with Demonstrations},
  author={Nicklas Hansen and Yixin Lin and Hao Su and Xiaolong Wang and Vikash Kumar and Aravind Rajeswaran},
  journal={arXiv preprint},
  year={2022}
}

We assume that your machine has a CUDA-enabled GPU, a local copy of MuJoCo 2.1.x installed (required for the Adroit/Meta-World domains), and at least 80GB of memory. Then, create a conda environment with conda env create -f environment.yml, and add /<path>/<to>/<your>/modem/tasks/mj_envs to your PYTHONPATH (required for the Adroit domain). No additional setup required for the DMControl domain. You will also need to configure wandb and your demonstration/logging directories in cfgs/config.yaml. Demonstrations are made available here. Once setup is complete, you should be able to run the following commands.

To train MoDem on a task from Adroit:

python train.py suite=adroit task=adroit-door

To train MoDem on a task from Meta-World:

python train.py suite=mw task=mw-assembly

To train MoDem on a task from DMControl:

python train.py suite=dmcontrol task=quadruped-run

This codebase is based on the original TD-MPC implementation. MoDem, TD-MPC and Meta-World are licensed under the MIT license. MuJoCo, DeepMind Control Suite, and mj_envs (Adroit) are licensed under the Apache 2.0 license. We thank the DrQv2 authors for their implementation of DMControl wrappers.