Code release for Model-Based Visual Planning with Self-Supervised Functional Distances.
This repository includes the distance learning component of the MBOLD method. To learn video prediction models, we use the training code located at RoboNet, and to run control, we use the visual foresight codebase. In order to run MBOLD, you will need to use all three components. If you only want to perform distance learning, this repository will be sufficient.
To install the package, clone the repository and run pip install -e . from the main directory containing setup.py.
Install Meta-World: Please see the instructions here to install Meta-World to use the simulated environments. Note: The latest version of Meta-World is incompatible with this codebase. Please use a commit with hash 56e7ef2c52b275325f2409df7f6c07ad09b1cc59. As an example:
pip install git+https://github.com/rlworkgroup/metaworld.git@56e7ef2c52b275325f2409df7f6c07ad09b1cc59#egg=metaworld
Installing FAISS: To perform GPU-enabled negative training example sampling, this code uses the FAISS library. Please see this document for information about how to install the GPU version of FAISS.
To train a distance function, run
python classifier_control/train.py <path_to_experiment_config>
For example, to train the MBOLD distance function on the three object tabletop pushing task data, run
python classifier_control/train.py experiments/distfunc_training/q_func_training/tabletop_3obj/conf.py
To run control experiments or perform data collection, run
python run_control_experiment.py <path_to_experiment_config>
For example, to run MBOLD on the three object tabletop pushing task, run
python run_control_experiment.py experiments/control/tabletop_3obj/mbold/hparams.py
Note that running control experiments requires a trained video prediction model, a trained distance function, as well as a set of evaluation tasks (possibly downloaded from the below section).
We provide training datasets for each simulated environment (collected using a random policy) as well as evaluation tasks
used in comparisons in the paper here.
To train on these datasets, download the .zip file and unzip it in your preferred location. Then, set the parameter 'data_dir'
in the training configuration file to point to this directory.
With the exception of the Frank door environment, the code is licensed under the terms of the MIT license. The assets from the Franka door environment are from the the playroom environment, licensed by Google Inc. under a Creative Commons Attribution 4.0 International License. The Franka Panda model is by Vikash Kumar and licensed under the terms of the Apache 2.0 license.
If you find this useful, consider citing:
@inproceedings{tian2021mbold,
title={Model-Based Visual Planning with Self-Supervised Functional Distances},
author={Tian, Stephen and Nair, Suraj and Ebert, Frederik and Dasari, Sudeep and Eysenbach, Benjamin and Finn, Chelsea and Levine, Sergey},
booktitle={International Conference on Learning Representations},
year={2021}
}