Heterogenous Pre-trained Transformers

Lirui Wang, Xinlei Chen, Jialiang Zhao, Kaiming He

Neural Information Processing Systems (Spotlight), 2024

This is a Huggingface LeRobot implementation for pre-training Heterogenous Pre-trained Transformers (HPTs).

LeRobot Installation

Create a virtual environment with Python 3.10 and activate it, e.g. with miniconda:

conda create -y -n lerobot python=3.10
conda activate lerobot

Install 🤗 LeRobot with simulation environments:

pip install -e ".[aloha, pusht]

Code Modification Walkthrough

Check the following two folders for most of the modifications.

├── lerobot
|   ├── configs          # contains hydra yaml files with all options that you can override in the command line
|   |   ├── ...            # various sim environments and their datasets: aloha.yaml, pusht.yaml, xarm.yaml
|   |   └── policy         # including policies config for hpt.yaml
|   ├── common           # contains classes and utilities
|   |   ├── ...       # various datasets of human demonstrations: aloha, pusht, xarm
|   |   ├── ...            # various sim environments: aloha, pusht, xarm
|   |   ├── policies       # including modeling and configuration for hpt
|   ├── ...

Experiment Scripts

0. By default, the HPT model loads the x-large pre-trained trunk. Use these config parameters policy.embed_dim=256 policy.num_heads=8 policy.num_blocks=16 to switch to the hpt-base trunk for example.

1. Run the following scripts for aloha transfer cube experiments.

python lerobot/scripts/train.py \
policy=hpt_transformer env=aloha  env.task=AlohaTransferCube-v0 \
dataset_repo_id=lerobot/aloha_sim_transfer_cube_human \
wandb.enable=true

2. Run the following scripts for push-T experiments.

python lerobot/scripts/train.py \
policy=hpt_pusht  env=pusht  env.task=PushT-v0 \
dataset_repo_id=lerobot/pusht \
wandb.enable=true

3. Run the following scripts for real-world Koch experiments.

python lerobot/scripts/train.py policy=hpt_koch_real env=koch_real \
dataset_repo_id=lerobot/koch_pick_place_5_lego  \
wandb.enable=true

Citation

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

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

Acknowledgement

Our implementation is built upon the excellent LeRobot codebase.