Add LoRA-aware JAX-to-PyTorch checkpoint conversion

[Ret1ehS]

Summary

This PR adds LoRA-aware handling to examples/convert_jax_model_to_pytorch.py.

Previously, converted checkpoints were loaded with load_state_dict(strict=False), so LoRA adapter tensors could appear as unexpected_keys and be silently ignored. This meant a LoRA-finetuned JAX checkpoint could be converted into a PyTorch checkpoint without the adapter updates being represented in the final weights.

This change folds LoRA adapter weights into the corresponding base weights before the existing slicing/load flow. For configs without LoRA adapters, the merge is a no-op and the previous conversion path is preserved.

Implementation notes

• Adds merge_lora_into_base() to merge LoRA weights before slicing/loading.
• Detects LoRA from the model config rather than from variant names.
• Uses the configured PaliGemma and action-expert variants instead of hardcoded model variant strings.
• Preserves the existing non-LoRA default precision behavior.
• Defaults to float32 when LoRA adapters are present, because merged LoRA weights showed noticeably lower drift in float32 than bfloat16.
• Raises an error if LoRA-related keys remain in unexpected_keys after loading, so failed merges do not silently produce divergent checkpoints.

OpenPI-specific LoRA details

openpi.models.lora.Einsum and openpi.models.lora.FeedForward apply LoRA slightly differently, so the merge follows the runtime behavior:

• attn_vec_einsum sums over the head dimension in the second LoRA einsum, so the merged delta must match that cross-head behavior rather than using an independent per-head outer product.
• FeedForward._dot() adds the LoRA path without applying scaling_value, while Einsum does apply scaling. The MLP merge therefore uses scale=1.0.

Validation

• pre-commit passed.
• ruff check . passed.
• ruff format . passed.
• python3 -m py_compile examples/convert_jax_model_to_pytorch.py passed.
• git diff --check passed.

Additional local validation on a LoRA-finetuned checkpoint:

• Before this change, LoRA adapter tensors appeared in unexpected_keys and were ignored.
• After merging, no LoRA keys remained in unexpected_keys.
• Corrected attn_vec_einsum merge matched the online LoRA computation closely:
  • max_abs = 0.000001669
  • mean_abs = 0.000000284
• Float32 converted checkpoint matched the JAX reference more closely than bfloat16:
  • float32: mean_abs = 0.0000929647, max_abs = 0.0016714334
  • bfloat16: mean_abs = 0.0011840940, max_abs = 0.0170528293

Fixes #958.
Related: #840, #729, #810.