README.md

April 1, 2025 · View on GitHub

[CVPR 2025] Tra-MoE: Learning Trajectory Prediction Model from Multiple Domains for Adaptive Policy Conditioning

Jiange Yang, Haoyi Zhu, Yating Wang, Gangshan Wu, Tong He, Liming Wang

caps

Low-Cost Dual-Arm Robot Demos

The videos are all done automatically by learned policy (Learn from human and robot data).

Fold	Pick and Pass	Pour

Pull out	Push

Prepare

Following ATM, install environment (including robomimic and robosuite).

conda env create -f environment.yml
conda activate atm
mkdir third_party & cd third_party
git clone https://github.com/ARISE-Initiative/robomimic.git
git clone https://github.com/ARISE-Initiative/robosuite.git
pip install -e third_party/robosuite/
pip install -e third_party/robomimic/

Downloading and processing libero data as well as using Cotracker to get trajectory labels.

mkdir data
python -m scripts.download_libero_datasets
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
python -m scripts.split_libero_dataset

Training

Stage 1: Training trajectory prediction models with actionless large-scale out-of-domain video data and small-scale in-domain video data.

USE_BFLOAT16=true python -m scripts.train_libero_track_transformer --suite $SUITE_NAME

Stage 2: Training trajectory-guided policy with small-scale in-domain robot data.

USE_BFLOAT16=false python -m scripts.train_libero_policy_atm --suite $SUITE_NAME --tt $PATH_TO_TT

Evaluation

USE_BFLOAT16=false python -m scripts.eval_libero_policy --suite $SUITE_NAME --exp-dir $PATH_TO_EXP

Checkpoints

You can download our trained checkpoints.

Citation

Please cite the following paper if you feel this repository useful for your research.

@article{yang2024tra,
  title={Tra-MoE: Learning Trajectory Prediction Model from Multiple Domains for Adaptive Policy Conditioning},
  author={Yang, Jiange and Zhu, Haoyi and Wang, Yating and Wu, Gangshan and He, Tong and Wang, Limin},
  journal={arXiv preprint arXiv:2411.14519},
  year={2024}
}

Acknowledges

Thanks to the open source of the following projects: