Repository files navigation

Lirui Wang,Xinlei Chen,Jialiang Zhao,Kaiming He

Neural Information Processing Systems (Spotlight), 2024

This is a pytorch implementation for pre-training Heterogenous Pre-trained Transformers (HPTs). The pre-training procedure train on mixture of embodiment datasets with a supervised learning objective. The pre-training process can take some time, so we also provide pre-trained checkpoints below. You can find more details on ourproject page. An alternative clean implementation of HPT in Hugging Face can also be foundhere.

TL;DR: HPT aligns different embodiment to a shared latent space and investigates the scaling behaviors in policy learning. Put a scalable transformer in the middle of your policy and don’t train from scratch!

1. pip install -e .

mkdir ~/.mujococd ~/.mujocowget https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz  -O mujoco210.tar.gz --no-check-certificatetar -xvzf mujoco210.tar.gz# add the following line to ~/.bashrc if neededexport LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${HOME}/.mujoco/mujoco210/binexport LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidiaexport LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64export MUJOCO_GL=egl

0. Check outquickstart.ipynb for how to use the pretrained HPTs.
1. python -m hpt.run train policies on each environment. Add+mode=debug for debugging.
2. bash experiments/scripts/metaworld/train_test_metaworld_1task.sh test test 1 +mode=debug for example script.
3. Changetrain.pretrained_dir for loading pre-trained trunk transformer. The model can be loaded either from local checkpoint folder or huggingfacerepository.
4. Run the following scripts for mujoco experiments.

bash experiments/scripts/metaworld/train_test_metaworld_20task_finetune.sh hf://liruiw/hpt-base

0. For training, it requires a dataset conversionconvert_dataset function for packing your own datasets. Checkthis for example.
1. For evaluation, it requires arollout_runner.py file for each benchmark and alearner_trajectory_generator evaluation function that provides rollouts.
2. If needed, modify theconfig for changing the perception stem networks and action head networks in the models. Take a look atrealrobot_image.yaml for example script in the real world.
3. Adddataset.use_disk=True for saving and loading the dataset in disk.

You can find pretrained HPT checkpoints here. At the moment we provide the following model versions:

├── ...├── HPT|   ├── data            # cached datasets|   ├── output          # trained models and figures|   ├── env             # environment wrappers|   ├── hpt             # model training and dataset source code|   |   ├── models      # network models|   |   ├── datasets    # dataset related|   |   ├── run         # transfer learning main loop|   |   ├── run_eval    # evaluation main loop|   |   └── ...|   ├── experiments     # training configs|   |   ├── configs     # modular configs└── ...

If you find HPT useful in your research, please consider citing:

@inproceedings{wang2024hpt,author    = {Lirui Wang, Xinlei Chen, Jialiang Zhao, Kaiming He},title     = {Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers},booktitle = {Neurips},year      = {2024}}

If you have any questions, feel free to contact me through email (liruiw@mit.edu). Enjoy!