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)
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.