feat: rewrite Pi0, enable loading pretrained weights and validate the performance

[kwangneuraco]

Features

This is a total rewrite of PI0 based on the official code from Physical Intelligence https://github.com/Physical-Intelligence/openpi. The official code has a lot of difference with our previous version ( How VLM and action expert interact, learning scheduler, gradient checkpoint to save VRAM etc) so I did a total rewrite based on the official code. At the same time I made the code interface same as our other algorithms and made several improvements.

Main features from the original code beside the main pi0 model structure:

• Modification of transformers library and gemma model (transformers_replace and gemma_torch.py)
• Can load pretrained model and do finetune
• Gradient checkpoint to save VRAM (finetuning full model only needs around 14 GB VRAM, 4090 can do it)

Features I added:

• Add patching code at __init__.py to automatically replace transformers library, instead of manually copy and paste as suggested by PI
• Modified code and updated the library version in requirements.txt to use numpy >= 2.0 instead of 1.0 from the official code
• Add flags so we can only fineture action experts, finetune action_experts + vision backbone (most cases lauguage part is not needed for finetuning) and full parameter finetune.
• • Use torch.compile() for inference speedup
• Some libraries requires specific version (transformers etc) so added them in requirements.txt

Questions from me:

• As PI0 model is complex, I hide most model code from pi0.py and put it in modlue.py and created a class call PI0Policy and PI0PolicyConfig. This makes the code cleaner and better structured (otherwise pi0.py may have 2k lines) but created a bit of repeated code.
• As we discussed PI05 will be separated from PI0 (which i agree, because users can modify pi0 without affecting any pi05 and vice versa). So i put trainsformers_replace and gemma_pytorch.py inside PI0, although these parts are same for PI05.

Validation

Finetune with vx300 for 3 epochs and achieves 70% success rate

Ticket

https://neuracore.atlassian.net/browse/NCRL-33

[github-actions]

Consider updating changelogs/pending-changelog.md with a summary of this change for the release notes. This is optional and non-blocking.

Stephen is on holiday and we already discussed how to improve the code and get the code in

[stepjam]

Implementation is good (pending one small comment), but some final modification needed to tests.

[stepjam]

LGTM