TACO-RL

Latent Plans for Task-Agnostic Offline Reinforcement Learning

Erick Rosete-Beas, Oier Mees, Gabriel Kalweit, Joschka Boedecker, Wolfram Burgard

We present TACO-RL (Task-AgnostiC Offline Reinforcement Learning), a hierarchical general-purpose agent that absorbs uncurated, unlabeled, highly diverse, offline data. TACO-RL makes sense of this data by combining the strengths of imitation learning and offline RL. We combine a low-level policy that learns latent skills via imitation and a high-level policy learned from offline RL for skill-chaining the latent behaviors. We learn a single multi-task visuomotor policy for over 25 tasks in the real world that outperforms state-of-the-art baselines by an order of magnitude.

💻 Quick Start

The package was tested using python 3.7 and Ubuntu 20.04 LTS.

# At root of project
git clone https://github.com/ErickRosete/tacorl.git
export TACORL_ROOT=$(pwd)/tacorl
conda create -n tacorl_venv python=3.7 
conda activate tacorl_venv
sh install.sh

This script install all the required dependencies for this repository. If you want to download the dataset follow

Download

Task-Agnostic Real World Robot Play Dataset

We host the multimodal 9 hours of human teleoperated play dataset on kaggle.

cd $TACORL_ROOT/dataset
kaggle datasets download -d oiermees/taco-robot

CALVIN Dataset

If you want to train on the simulated CALVIN dataset, we augment the original dataset in order to balance the task distribution:

cd $TACORL_ROOT/dataset
wget http://tacorl.cs.uni-freiburg.de/dataset/taco_rl_calvin.zip && unzip taco_rl_calvin.zip && rm taco_rl_calvin.zip

🏋️‍♂️ Train TACO-RL Agent

Simulation

To run the LMP training

python scripts/train.py experiment=play_lmp_for_rl data_dir="dataset/calvin"

To run the TACO-RL training

python scripts/train.py experiment=tacorl data_dir="dataset/calvin" module.play_lmp_dir="models/lmp_calvin"

Real world

To run the LMP training

python scripts/train.py experiment=play_lmp_real_world data_dir="dataset/real_world" 

To run the TACO-RL training

python scripts/train.py experiment=tacorl_real_world data_dir="dataset/real_world" module.play_lmp_dir="models/lmp_real_world"

The training will start automatically with the last checkpoint of the LMP model if you want to run another checkpoint for further experimentation, you can specify it as follows

python scripts/train.py experiment=tacorl_real_world data_dir="dataset/real_world" module.play_lmp_dir="models/lmp_real_world" +module.lmp_epoch_to_load=30

🏆 Evaluation

To run our evaluation on the CALVIN environment you can use the following script

python scripts/evaluate.py evaluation=tacorl_easy module_path="PATH_MODEL_TO_EVAL" data_dir="dataset/calvin"

You can choose between the following options

• tacorl_easy: Single goal tasks where the goal image contains the end effector achieving the task.
• tacorl_hard: Single goal tasks where the goal image does not contain the end effector achieving the task.
• tacorl_lh_easy: Perform two tasks in a row using a single goal image.
• tacorl_lh_seq_easy: Perform five tasks in a row using intermediate goal images.

Analogous evaluation configurations are available for CQL, LMP and RIL.

Additional evaluation instructions

The evaluation will use the last epoch of the TACO-RL training, if you want to run another epoch you can also explicitly state it

python scripts/evaluate.py evaluation=tacorl_easy module_path="PATH_MODEL_TO_EVAL" data_dir="dataset/calvin" +epoch_to_load=30

In case you want to change something from the module you are evaluating, e.g. changing the play_lmp_dir path to load, you can run the following command

python scripts/evaluate.py evaluation=tacorl_easy module_path="PATH_MODEL_TO_EVAL" data_dir="dataset/calvin" +overwrite_module_cfg.play_lmp_dir="LMP_PATH_TO_LOAD"

🧑‍🎓 Pre-trained Models

Download the TACO-RL model checkpoint trained on the static camera rgb images on CALVIN environment D. Additionally, you can find checkpoints for the baselines CQL, RIL and PlayLMP inside the zip file.

$ wget http://tacorl.cs.uni-freiburg.de/model_weights/tacorl_calvin.zip
$ unzip tacorl_calvin.zip

Acknowledgements

This work uses code from the following open-source projects and datasets:

CALVIN

Original: https://github.com/mees/calvin License: MIT

✍️ Citation

If you find the dataset or code useful, please cite:

@inproceedings{rosete2022tacorl,
author = {Erick Rosete-Beas and Oier Mees and Gabriel Kalweit and Joschka Boedecker and Wolfram Burgard},
title = {Latent Plans for Task Agnostic Offline Reinforcement Learning},
journal = {Proceedings of the 6th Conference on Robot Learning (CoRL)},
year = {2022}
}

Contributing

If you want to contribute to the repo it is important to install git hooks in your .git/ directory, such that your code complies with the black code style and Flake8 style guidelines.
In the root directory run:
pre-commit install