refactor pi052 to reuse pi05

tests/policies/pi052/test_pi052_bucketed_ce.py

@@ -37,7 +37,7 @@ def test_shifted_ce_none_retains_distinct_per_sample_losses():
 
 
 def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_path):
-    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+    import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
 
     checkpoint = tmp_path / "model.safetensors"
     checkpoint.touch()
@@ -47,8 +47,8 @@ def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_pa
         calls.append((model_id, filename, kwargs))
         return None if filename.endswith("index.json") else str(checkpoint)
 
-    monkeypatch.setattr(modeling_pi052, "cached_file", fake_cached_file)
-    files = modeling_pi052._resolve_weight_files(
+    monkeypatch.setattr(modeling_pi05, "cached_file", fake_cached_file)
+    files = modeling_pi05._resolve_weight_files(
         "org/model",
         force_download=True,
         resume_download=True,
@@ -71,10 +71,10 @@ def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_pa
 
 
 def test_checkpoint_resolution_rejects_local_directory_without_weights(tmp_path):
-    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+    import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
 
     with pytest.raises(FileNotFoundError, match="model.safetensors"):
-        modeling_pi052._resolve_weight_files(
+        modeling_pi05._resolve_weight_files(
             tmp_path,
             force_download=False,
             resume_download=None,

tests/policies/pi052/test_pi052_vision_checkpointing.py

@@ -41,14 +41,12 @@ def _checkpoint_model():
     tower = _MockVisionTower()
     language_model = SimpleNamespace(gradient_checkpointing=False)
     expert_model = SimpleNamespace(gradient_checkpointing=False)
-    model = SimpleNamespace(
-        gradient_checkpointing_enabled=False,
-        paligemma_with_expert=SimpleNamespace(
-            paligemma=SimpleNamespace(
-                model=SimpleNamespace(language_model=language_model, vision_tower=tower)
-            ),
-            gemma_expert=SimpleNamespace(model=expert_model),
-        ),
+    model = PI05Pytorch.__new__(PI05Pytorch)
+    nn.Module.__init__(model)
+    model.gradient_checkpointing_enabled = False
+    model.paligemma_with_expert = SimpleNamespace(
+        paligemma=SimpleNamespace(model=SimpleNamespace(language_model=language_model, vision_tower=tower)),
+        gemma_expert=SimpleNamespace(model=expert_model),
     )
     return model, tower, language_model, expert_model
 

tests/policies/pi0_pi05/test_pi05.py

@@ -16,8 +16,12 @@
 
 """Test script to verify PI0.5 (pi05) support in PI0 policy"""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
+from safetensors.torch import save_file
+from torch import nn
 
 pytest.importorskip("transformers")
 
@@ -31,6 +35,26 @@ from lerobot.utils.random_utils import set_seed
 from tests.utils import require_cuda, require_hf_token  # noqa: E402
 
 
+class _CheckpointPolicy(PI05Policy):
+    def __init__(self, config, **kwargs):
+        nn.Module.__init__(self)
+        self.config = config
+        self.loaded_state_dict = None
+
+    def load_state_dict(self, state_dict, strict=True, assign=False):
+        self.loaded_state_dict = state_dict
+        return [], []
+
+
+def test_from_pretrained_loads_existing_single_file_checkpoint(tmp_path):
+    save_file({"weight": torch.tensor([1.0])}, tmp_path / "model.safetensors")
+
+    policy = _CheckpointPolicy.from_pretrained(tmp_path, config=SimpleNamespace())
+
+    assert policy.loaded_state_dict is not None
+    torch.testing.assert_close(policy.loaded_state_dict["model.weight"], torch.tensor([1.0]))
+
+
 @require_cuda
 @require_hf_token
 def test_policy_instantiation():