[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

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:

ATM

CoTracker

mixture-of-experts

st-moe-pytorch

RealRobot

LIBERO

RLbench