Chuan Wen*, Xingyu Lin*, John So*, Kai Chen, Dou Qi, Yang Gao, Pieter Abbeel
Robotics: Science and Systems (RSS) 2024
git clone --recursive https://github.com/Large-Trajectory-Model/ATM.git
cd ATM/
conda env create -f environment.yml
conda activate atm
pip install -e third_party/robosuite/
pip install -e third_party/robomimic/
We first need to download the raw LIBERO datasets:
mkdir data
python -m scripts.download_libero_datasets
and then preprocess them with Cotracker:
python -m scripts.preprocess_libero --suite libero_spatial
python -m scripts.preprocess_libero --suite libero_object
python -m scripts.preprocess_libero --suite libero_goal
python -m scripts.preprocess_libero --suite libero_10
python -m scripts.preprocess_libero --suite libero_90
After preprocessing, split the datasets into training and validation sets:
python -m scripts.split_libero_dataset
To reproduce the experimental results in our paper, we provide the checkpoints trained by us. Please download the zip file from here and put it in the current folder. Then,
mkdir results
unzip -o atm_release_checkpoints.zip -d results/
rm atm_release_checkpoints.zip
As shown in Figure 2 in our paper, the training of our Trajectory Modeling framework includes two stages: Track Transformer Pretraining and Trajectory-guided Policy Training.
The Track Transformer training can be executed by this command, where SUITE_NAME can be libero_spatial, libero_object, libero_goal, or libero_100:
python -m scripts.train_libero_track_transformer --suite $SUITE_NAME
The vanilla BC baseline can be trained by the following command, where $SUITE_NAME can be libero_spatial, libero_object, libero_goal, or libero_10 (i.e., LIBERO-Long in our paper):
python -m scripts.train_libero_policy_bc --suite $SUITE_NAME
Our Track-guided policy can be trained with:
python -m scripts.train_libero_policy_atm --suite $SUITE_NAME --tt $PATH_TO_TT
where $PATH_TO_TT is the path to the folder of Track Transformer pretrained in Stage 1. We have provided the pretrained checkpoints in results/track_transformers/. For example,
python -m scripts.train_libero_policy_atm --suite libero_spatial --tt results/track_transformer/libero_track_transformer_libero-spatial/
python -m scripts.train_libero_policy_atm --suite libero_object --tt results/track_transformer/libero_track_transformer_libero-object/
python -m scripts.train_libero_policy_atm --suite libero_goal --tt results/track_transformer/libero_track_transformer_libero-goal/
python -m scripts.train_libero_policy_atm --suite libero_10 --tt results/track_transformer/libero_track_transformer_libero-100/
The evaluation can be executed by this command, where $SUITE_NAME is the desired suite name and $PATH_TO_EXP is the path to the your trained policy folder in results/policy/. The success rate and evaluation videos will be saved in $PATH_TO_EXP/eval_results/.
python -m scripts.eval_libero_policy --suite $SUITE_NAME --exp-dir $PATH_TO_EXP
For example, you can evaluate the provided checkpoints by:
python -m scripts.eval_libero_policy --suite libero_spatial --exp-dir results/policy/atm-policy_libero-spatial_demo10
python -m scripts.eval_libero_policy --suite libero_object --exp-dir results/policy/atm-policy_libero-object_demo10
python -m scripts.eval_libero_policy --suite libero_goal --exp-dir results/policy/atm-policy_libero-goal_demo10
python -m scripts.eval_libero_policy --suite libero_10 --exp-dir results/policy/atm-policy_libero-10_demo10
If you find our codebase is useful for your research, please cite our paper with this bibtex:
@article{wen2023atm,
title={Any-point trajectory modeling for policy learning},
author={Wen, Chuan and Lin, Xingyu and So, John and Chen, Kai and Dou, Qi and Gao, Yang and Abbeel, Pieter},
journal={arXiv preprint arXiv:2401.00025},
year={2023}
}