🧸 Robots Pre-Train Robots: Manipulation-Centric Robotic Representation from Large-Scale Robot Datasets

Project Page | arXiv | Twitter | Dataset & Model

Guangqi Jiang*, Yifei Sun*, Tao Huang*, Huanyu Li, Yongyuan Liang, Huazhe Xu

In submission, 2024

🗞️ News

• 2024-10-31 Release code!!
• 2024-10-29 Release our paper on ArXiv.

🛠️ Installation

Clone this repository and create a conda environment:

git clone https://github.com/luccachiang/robots-pretrain-robots.git
cd robots-pretrain-robots
conda remove -n mcr --all
conda env create -f mcr/mcr.yaml
conda activate mcr

Install MCR:

pip install -e .

📚 Data and checkpoints

Our processed DROID subset (coming soon) and pre-trained model checkpoints are availble on our Huggingface repository. Our dataset has the following structure. As long as your dataset share the same, you can directly train MCR on your custom dataset.

/droid_processed
    /2023-02-28_Tue_Feb_28_20_25_47_2023
        /exterior_image_1_left
            /0.png
            /1.png
            ...
        /exterior_image_2_left
        other_data.pkl # proprioception and actions
    /2023-02-28_Tue_Feb_28_20_26_11_2023
    ...

💻 Usage

You can use this codebase for the following purposes:

1. Use our released pre-trained checkpoints.

# first, download our model checkpoint from Huggingface
# then get a torchvision.models.resnet50
import mcr
encoder = mcr.load_model(ckpt_path=<path_to_downloaded_ckpt>)

# please see more details in utils/example.py

2. Train MCR from scratch.

# first, download our pre-trained dataset from Huggingface TODO
# then run
cd mcr
bash train_mcr.sh
# you can get a full list of parameter helps in train_mcr.sh

3. Train MCR with custom dataset.

We also provide a guidance on how to train MCR on your own dataset. You can either process your data into the structure we provide above or modify the codebase to write your own dataloader. TODO Specifically, you need to modify code in xxx, xxx, and xxx.

🧭 Code Navigation

todo

🏷️ Licence

This repository is released under the MIT license. See LICENSE for additional details.

✉️ Acknowledgement & Contact

Our codebase is built upon R3M. We thank all these authors for their nicely open sourced code and their great contributions to the community.

Please contact Guangqi Jiang if you are interested in this project. Also feel free to open an issue or raise a pull request :)

📝 BibTeX

We will be glad if you find this work helpful. Please consider citing:

@article{jiang2024robots,
        title={Robots Pre-Train Robots: Manipulation-Centric Robotic Representation from Large-Scale Robot Datasets},
        author={Jiang, Guangqi and Sun, Yifei and Huang, Tao and Li, Huanyu and Liang, Yongyuan and Xu, Huazhe},
        journal={arXiv preprint arXiv:2410.22325},
        year={2024}
        }