How to properly merge LoRA weights back into a Pi0.5 (pi05) base checkpoint?

[tqb-web]

Hi OpenPI team,

Thank you for open-sourcing OpenPI and the Pi0.5 models. I've been using the JAX LoRA fine-tuning pipeline with pi05_base and it has been working very well. The official OpenPI repository provides the pi05_base checkpoint for fine-tuning, and LoRA is the recommended approach for many setups.

I would like to export a standalone Pi0.5 model with the LoRA weights merged into the original model weights, so that the resulting checkpoint no longer depends on LoRA modules and can be used by downstream tools (for example LeRobot or PyTorch conversion workflows).

My understanding is:

1. Load pi05_base.
2. Restore the LoRA fine-tuned checkpoint.
3. Merge the LoRA weights into the corresponding base weights.
4. Save a new checkpoint containing only the merged base model weights.

However, I am not sure whether OpenPI already provides an official or recommended way to do this.