CauchyFood / mvp

Masked Visual Pre-training for Motor Control

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Masked Visual Pre-training for Motor Control

This is a PyTorch implementation of the paper Masked Visual Pre-training for Motor Control. It contains the benchmark suite, pre-trained models, and the training code to reproduce the results from the paper.

Installation

Please see INSTALL.md for installation instructions.

Pre-trained visual enocoders

We provide pre-trained visual encoders used in the paper. The models are in the same format as mae and timm:

backbone objective data md5 download
ViT-S MAE in-the-wild model
ViT-S MAE ImageNet model
ViT-S Supervised ImageNet model

By default, the code assumes that the pre-trained encoders are placed in /tmp/pretrained directory.

Example training commands

Train FrankaPick from states:

python tools/train.py task=FrankaPick

Train FrankaPick from pixels:

python tools/train.py task=FrankaPickPixels

Train on 8 GPUs:

python tools/train_dist.py num_gpus=8

Test a policy after N iterations:

python tools/train.py test=True headless=False logdir=/path/to/job resume=N

Citation

If you find the code or pre-trained models useful in your research, please use the following BibTeX entry:

@article{Xiao2022
  title = {Masked Visual Pre-training for Motor Control},
  author = {Tete Xiao and Ilija Radosavovic and Trevor Darrell and Jitendra Malik},
  journal = {arXiv:2203.06173},
  year = {2022}
}

Acknowledgments

We thank NVIDIA IsaacGym and PhysX teams for making the simulator and preview code examples available.

About

Masked Visual Pre-training for Motor Control


Languages

Language:Python 99.8%Language:Batchfile 0.2%Language:CMake 0.1%