feat(pi0): add train_expert_only and freeze_vision_encoder flags to pi0 and pi0.5

[pxe3]

Title

Add train_expert_only and freeze_vision_encoder flags to PI0 and PI0.5

Type / Scope

• Type: FEATURE
• Scope: PI0 and PI0.5

Summary / Motivation

Adds two config options for pi0/0.5 for finetuning purposes: freeze_vision_encoder (only freeze vision encoder) and train_expert_only (freeze entire VLM, only train action expert and projection layers). This addresses #2153.

Related issues

• Closes: Why can’t I find something like train_expert_only in the latest version of pi0? Do the current versions of pi0 and pi0.5 only support full-parameter training? #2153

What changed

• Added freeze_vision_encoder and train_expert_only config options to PI0Config and PI05Config
• Added _set_requires_grad() to PaliGemmaWithExpertModel in both pi0 and pi0.5 to freeze params based on config
• Added train() override to keep frozen modules in eval mode during training

Both options default to False for backwards compatibility.

How was this tested

• Tests added: None (existing tests cover config and policy instantiation - consistent with testing for xvla and smolvla freezing)
• Manual checks / dataset runs performed: pytest tests/policies/pi0_pi05/test_pi0.py tests/policies/pi0_pi05/test_pi05.py -v

Testing with DEVICE='cuda'

tests/policies/pi0_pi05/test_pi0.py::test_policy_instantiation PASSED                                                                                              [ 25%]
tests/policies/pi0_pi05/test_pi0.py::test_config_creation PASSED                                                                                                   [ 50%]
tests/policies/pi0_pi05/test_pi05.py::test_policy_instantiation PASSED                                                                                             [ 75%]
tests/policies/pi0_pi05/test_pi05.py::test_config_creation PASSED                                                                                                  [100%]

===================================================================== 4 passed in 126.45s (0:02:06) ======================================================================

• Pre-commit: all passed

How to run locally (reviewer)

• Run the relevant tests: pytest tests/policies/pi0_pi05/test_pi0.py tests/policies/pi0_pi05/test_pi05.py -v

Checklist (required before merge)

• Linting/formatting run (pre-commit run -a)
• All tests pass locally (pytest)
• Documentation updated
• CI is green

Reviewer notes

Hello! This is my first contribution to LeRobot and additionally, one of my first open source contributions. I appreciate and welcome your guidance, and I apologize if there were any best practices not followed.

Happy to make any changes, thanks in advance for the feedback!