lerobot icon indicating copy to clipboard operation
lerobot copied to clipboard

Support for PI0+FAST

Open mshukor opened this issue 9 months ago • 0 comments

Overview

Support for PI0+FAST.

Paper: https://arxiv.org/abs/2501.09747 Original JAX model: https://github.com/Physical-Intelligence/openpi

  • This is a working version that achieve around 50% success rate on LIBERO. This is with careful selection of different hyperparameters (now default in configuration_pi0fast.py).

  • The implementation does not follow exactly the JAX model one (e.g. we don't use EMA). For instance, it's based on transformers and leverage several existing tools implemented there (e.g. .generate()).

  • Due to its size (~3B parameters), we suggest training the model on relatively large dataset, we also found that it's better to train with large batch sizes (e.g. 32)

The. model should achieve more than 80% SR as mentioned in the original paper, so any effort to improve this current version is welcome by the community.

Pretrained base checkpoint and instructions to train the model can be found here: https://huggingface.co/lerobot/pi0fast_base

This is a collaborative work with @danaaubakirova @molbap and @Cadene.

mshukor avatar Mar 31 '25 15:03 mshukor