admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 09:46:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(pi052): split pi05_backbone into pi_gemma + modeling_pi052

Browse Source 
Eliminate the standalone pi052/pi05_backbone.py by distributing its contents:
- Generic dual-expert transformer machinery -> lerobot/policies/pi_gemma.py
  (sdpa_attention_forward, compute_layer_complete, PaliGemmaWithExpertModel,
  get_gemma_config; the openpi width/depth config is renamed GemmaConfig ->
  GemmaVariantConfig to avoid clashing with transformers' GemmaConfig). These
  sit next to the existing PiGemma layer code they already depend on.
- pi052-specific model + helpers -> pi052/modeling_pi052.py (PI05Pytorch,
  ActionSelectKwargs, make_att_2d_masks, pad_vector, resize_with_pad_torch,
  create_sinusoidal_pos_embedding, sample_beta, get_safe_dtype).

DEFAULT_IMAGE_SIZE is duplicated as a plain constant in pi_gemma to avoid a
pi_gemma -> pi05 import cycle. Additive to pi_gemma; pi0/pi05 unaffected.
Verified bit-exact on pepijn223/pi052_robocasa_full (embed/predict/forward
identical) and all 34 pi052 tests pass.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
pepijn223
2026-06-04 20:18:18 +02:00
parent 4cbd91a04e
commit c020c0d053
4 changed files with 933 additions and 963 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/policies/pi052/test_pi052_sdpa_attention.py
+2 -4
View File
@@ -33,10 +33,8 @@ pytest.importorskip("transformers")
from transformers.models.gemma import modeling_gemma # noqa: E402
from lerobot.policies.pi052.pi05_backbone import ( # noqa: E402
make_att_2d_masks,
sdpa_attention_forward,
)
from lerobot.policies.pi052.modeling_pi052 import make_att_2d_masks # noqa: E402
from lerobot.policies.pi_gemma import sdpa_attention_forward # noqa: E402
from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE # noqa: E402