Sirius-Fleet: Multi-Task Interactive Robot Fleet Learning with Visual World Models

Huihan Liu, Yu Zhang, Vaarij Betala, Evan Zhang, James Liu, Crystal Ding, Yuke Zhu
UT Austin Robot Perception and Learning Lab
Conference on Robot Learning, 2024
[Paper]  [Project Website]

This codebase is build off open-source codebase robocasa and robomimic.

Installation

1. Install robocasa

Install robocasa with the following procedures. Refer to robocasa github and documentation for more details.

# create conda env
conda create -c conda-forge -n robocasa python=3.10
conda activate robocasa

# install robosuite dependency (use the master branch)
git clone https://github.com/huihanl/robosuite-sirius-fleet.git
cd robosuite-sirius-fleet
pip install -e .

# install robocasa
cd ..
git clone https://github.com/robocasa/robocasa
cd robocasa
pip install -e .
pip install pre-commit; pre-commit install           # Optional: set up code formatter.

# (optional: if running into issues with numba/numpy, run: conda install -c numba numba=0.56.4 -y)

# download assets
python robocasa/scripts/download_kitchen_assets.py   # Caution: Assets to be downloaded are around 5GB.
python robocasa/scripts/setup_macros.py              # Set up system variables.

3. Install Sirius Fleet

git clone https://github.com/UT-Austin-RPL/sirius-fleet
cd sirius-fleet
pip install -e .
pip install -r requirements.txt

4. Install diffusers

git clone git@github.com:huggingface/diffusers.git
pip install -e diffusers 

"""
(Optional) You might need to remove importing cached_download in src/diffusers/utils/dynamic_modules_utils.py:
 -from huggingface_hub import cached_download, hf_hub_download, model_info
 +from huggingface_hub import hf_hub_download, model_info
"""

Usage

Downloading data

Downloading original data

robocasa (simulation): Follow robocasa's documentation to download the data. Note that the world model training is using the mimicgen dataset.

mutex (real robot): Download mutex dataset here.

Create symlinks for the data such that robocasa dataset is at ~/robocasa, and mutex is at ~/mutex:

ln -s $ROBOCASA_DATASET_PATH ~/robocasa
ln -s $MUTEX_DATASET_PATH ~/mutex

Downloading human-in-the-loop data

robocasa (simulation): download dataset here

mutex (real robot): download dataset here

Create symlinks for the data so that the human-in-the-loop data will be saved at ~/robocasa-hitl and ~/mutex-hitl:

ln -s $HITL_ROBOCASA_DATASET_PATH ~/robocasa-hitl
ln -s $HITL_MUTEX_DATASET_PATH ~/mutex-hitl

Training Sirius-Fleet

World Model

We suggest that you use the config generation script to generate a json file that is customized to the environments and corresponding datasets. It also allows you to specify different hyperparameters using the flags, and also sweep different hyperparameter options. See gen_config_world_model.py and ../helper.py for details.

Robocasa:

python robomimic/scripts/config_gen/world_model.py --env robocasa --name world_model --batch_size 8

Mutex:

python robomimic/scripts/config_gen/world_model.py --env mutex --name world_model --batch_size 8

This training generates a checkpoint, which we refer to here as world_model.pth. It will be used to train the failure classifier below.

Failure Classifier

We train the failure classifier based on the world model checkpoint world_model.pth from the previous timestep.

Change your task's data path ($DATASET_PATH). Note that this dataset is assumed to contain the attribute intv_labels as classifier labels.

python robomimic/scripts/train.py --config robomimic/exps/templates/failure_classifier.json --dataset $DATASET_PATH --pretrained_world_model world_model.pth 

To finetune failure classifier on a previous failure classifier checkpoint:

python robomimic/scripts/train.py --config robomimic/exps/templates/failure_classifier.json --dataset $DATASET_PATH --pretrained_world_model world_model.pth --classifier_ckpt classifier.pth

Policy

We use BC-Transformer policy architecture.

Robocasa:

python robomimic/scripts/config_gen/bc.py --env robocasa --name bc --n_seeds 3

Mutex:

python robomimic/scripts/config_gen/bc.py --env mutex --name bc --n_seeds 3

To finetune the policy on a previous policy checkpoint:

python robomimic/scripts/config_gen/bc.py --env $ENV --name bc --n_seeds 3 --ckpt_path $POLICY_CKPT

Acknowledgements

This codebase is largely built off robocasa, robomimic and robosuite.

For real-robot experiments, we used Deoxys, a controller library for Franka Emika Panda developed by Yifeng Zhu.