Support for PI0+FAST
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
transformersand 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.