refactor(pi052): load native base checkpoint

docs/source/pi052.mdx

@@ -92,14 +92,14 @@ recipe resolver.
 
 ## Train Pi052
 
-This example initializes Pi052 from the public Pi05 base checkpoint and trains
-the default subtask-and-memory recipe:
+This example initializes Pi052 from the native Pi052 initialization checkpoint
+and trains the default subtask-and-memory recipe:
 
 ```bash
 lerobot-train \
     --dataset.repo_id=${HF_USER}/my_language_annotated_dataset \
     --policy.type=pi052 \
-    --policy.pretrained_path=lerobot/pi05_base \
+    --policy.pretrained_path=lerobot/pi052_base \
     --policy.recipe_path=recipes/subtask_mem.yaml \
     --policy.dtype=bfloat16 \
     --policy.device=cuda \

src/lerobot/policies/pi05/modeling_pi05.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 
 import builtins
-import json
 import logging
 from collections import deque
 from pathlib import Path
@@ -76,81 +75,6 @@ class ActionSelectKwargs(TypedDict, total=False):
 
 
 _SAFETENSORS_FILE = "model.safetensors"
-_SAFETENSORS_INDEX = "model.safetensors.index.json"
-
-
-def _resolve_weight_files(
-    pretrained_name_or_path: str | Path,
-    *,
-    force_download: bool,
-    resume_download: bool | None,
-    proxies: dict | None,
-    token: str | bool | None,
-    cache_dir: str | Path | None,
-    local_files_only: bool,
-    revision: str | None,
-) -> list[Path]:
-    model_id = str(pretrained_name_or_path)
-    local_dir = Path(model_id)
-    load_kwargs = {
-        "revision": revision,
-        "cache_dir": cache_dir,
-        "force_download": force_download,
-        "resume_download": resume_download,
-        "proxies": proxies,
-        "token": token,
-        "local_files_only": local_files_only,
-    }
-
-    if local_dir.is_dir():
-        index_path = local_dir / _SAFETENSORS_INDEX
-        single_path = local_dir / _SAFETENSORS_FILE
-    else:
-        resolved_index = cached_file(
-            model_id,
-            _SAFETENSORS_INDEX,
-            _raise_exceptions_for_missing_entries=False,
-            **load_kwargs,
-        )
-        index_path = Path(resolved_index) if resolved_index is not None else None
-        single_path = None
-        if index_path is None:
-            resolved_file = cached_file(model_id, _SAFETENSORS_FILE, **load_kwargs)
-            single_path = Path(resolved_file) if resolved_file is not None else None
-
-    if index_path is None or not index_path.is_file():
-        if single_path is None or not single_path.is_file():
-            raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
-        return [single_path]
-
-    index = json.loads(index_path.read_text())
-    shard_names = sorted(set(index.get("weight_map", {}).values()))
-    if not shard_names:
-        raise ValueError(f"Invalid safetensors index without a weight_map: {index_path}")
-    if local_dir.is_dir():
-        files = [local_dir / name for name in shard_names]
-    else:
-        files = []
-        for name in shard_names:
-            resolved_file = cached_file(model_id, name, **load_kwargs)
-            if resolved_file is None:
-                raise FileNotFoundError(f"Checkpoint shard {name!r} not found in {model_id!r}.")
-            files.append(Path(resolved_file))
-    missing = [str(path) for path in files if not path.is_file()]
-    if missing:
-        raise FileNotFoundError(f"Missing checkpoint shards: {missing}")
-    return files
-
-
-def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
-    state_dict: dict[str, Tensor] = {}
-    for path in files:
-        shard = load_file(path)
-        overlap = state_dict.keys() & shard.keys()
-        if overlap:
-            raise ValueError(f"Duplicate checkpoint keys in {path}: {sorted(overlap)[:5]}")
-        state_dict.update(shard)
-    return state_dict
 
 
 # Define the complete layer computation function for gradient checkpointing
@@ -862,6 +786,7 @@ class PI05Policy(PreTrainedPolicy):
     model_class = PI05Pytorch
     eval_after_pretrained_load = False
     show_openpi_disclaimer = True
+    use_native_pretrained_loader = False
 
     def __init__(
         self,
@@ -905,7 +830,22 @@ class PI05Policy(PreTrainedPolicy):
         strict: bool = True,
         **kwargs,
     ) -> T:
-        """Load PI05-compatible single-file or sharded safetensors checkpoints."""
+        """Load a native LeRobot checkpoint or convert the PI05 base checkpoint."""
+        if cls.use_native_pretrained_loader:
+            return super().from_pretrained(
+                pretrained_name_or_path,
+                config=config,
+                force_download=force_download,
+                resume_download=resume_download,
+                proxies=proxies,
+                token=token,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+                revision=revision,
+                strict=strict,
+                **kwargs,
+            )
+
         if cls.show_openpi_disclaimer:
             print(
                 "The PI05 model is a direct port of the OpenPI implementation. \n"
@@ -929,8 +869,11 @@ class PI05Policy(PreTrainedPolicy):
             )
 
         model = cls(config, **kwargs)
-        files = _resolve_weight_files(
-            pretrained_name_or_path,
+        model_id = str(pretrained_name_or_path)
+        resolved_file = cached_file(
+            model_id,
+            _SAFETENSORS_FILE,
+            _raise_exceptions_for_missing_entries=False,
             force_download=force_download,
             resume_download=resume_download,
             proxies=proxies,
@@ -939,7 +882,10 @@ class PI05Policy(PreTrainedPolicy):
             local_files_only=local_files_only,
             revision=revision,
         )
-        fixed_state_dict = model._fix_pytorch_state_dict_keys(_load_weight_files(files), model.config)
+        if resolved_file is None:
+            raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
+
+        fixed_state_dict = model._fix_pytorch_state_dict_keys(load_file(resolved_file), model.config)
         remapped_state_dict = {
             key if key.startswith("model.") else f"model.{key}": value
             for key, value in fixed_state_dict.items()

src/lerobot/policies/pi052/modeling_pi052.py

@@ -878,6 +878,7 @@ class PI052Policy(PI05Policy):
     model_class = PI05Pytorch
     eval_after_pretrained_load = True
     show_openpi_disclaimer = False
+    use_native_pretrained_loader = True
 
     def __init__(self, config: PI052Config, **kwargs: Any) -> None:
         # Patch before constructing Gemma/SigLIP layers; the operation is optional and idempotent.

src/lerobot/policies/pretrained.py

@@ -338,6 +338,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             "smolvla": "lerobot/smolvla_base",
             "pi0": "lerobot/pi0_base",
             "pi05": "lerobot/pi05_base",
+            "pi052": "lerobot/pi052_base",
             "pi0_fast": "lerobot/pi0fast-base",
             "xvla": "lerobot/xvla-base",
         }

tests/policies/pi052/test_pi052_bucketed_ce.py

@@ -36,56 +36,6 @@ def test_shifted_ce_none_retains_distinct_per_sample_losses():
     assert losses[0] < losses[1]
 
 
-def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_path):
-    import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
-
-    checkpoint = tmp_path / "model.safetensors"
-    checkpoint.touch()
-    calls = []
-
-    def fake_cached_file(model_id, filename, **kwargs):
-        calls.append((model_id, filename, kwargs))
-        return None if filename.endswith("index.json") else str(checkpoint)
-
-    monkeypatch.setattr(modeling_pi05, "cached_file", fake_cached_file)
-    files = modeling_pi05._resolve_weight_files(
-        "org/model",
-        force_download=True,
-        resume_download=True,
-        proxies={"https": "proxy"},
-        token="secret",
-        cache_dir=tmp_path / "cache",
-        local_files_only=True,
-        revision="commit",
-    )
-
-    assert files == [checkpoint]
-    for _model_id, _filename, kwargs in calls:
-        assert kwargs["revision"] == "commit"
-        assert kwargs["cache_dir"] == tmp_path / "cache"
-        assert kwargs["force_download"] is True
-        assert kwargs["resume_download"] is True
-        assert kwargs["proxies"] == {"https": "proxy"}
-        assert kwargs["token"] == "secret"
-        assert kwargs["local_files_only"] is True
-
-
-def test_checkpoint_resolution_rejects_local_directory_without_weights(tmp_path):
-    import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
-
-    with pytest.raises(FileNotFoundError, match="model.safetensors"):
-        modeling_pi05._resolve_weight_files(
-            tmp_path,
-            force_download=False,
-            resume_download=None,
-            proxies=None,
-            token=None,
-            cache_dir=None,
-            local_files_only=False,
-            revision=None,
-        )
-
-
 @pytest.mark.parametrize("z_loss_weight", [0.0, 1e-4])
 @pytest.mark.parametrize("rows,valid_rows", [(24, 9), (48, 25)])
 def test_bucketed_ce_matches_dense_loss_and_gradients(z_loss_weight, rows, valid_rows):

tests/policies/pi0_pi05/test_pi05.py

@@ -46,6 +46,15 @@ class _CheckpointPolicy(PI05Policy):
         return [], []
 
 
+class _NativeCheckpointPolicy(PI05Policy):
+    use_native_pretrained_loader = True
+
+    def __init__(self, config, **kwargs):
+        nn.Module.__init__(self)
+        self.config = config
+        self.weight = nn.Parameter(torch.zeros(1))
+
+
 def test_from_pretrained_loads_existing_single_file_checkpoint(tmp_path):
     save_file({"weight": torch.tensor([1.0])}, tmp_path / "model.safetensors")
 
@@ -55,6 +64,64 @@ def test_from_pretrained_loads_existing_single_file_checkpoint(tmp_path):
     torch.testing.assert_close(policy.loaded_state_dict["model.weight"], torch.tensor([1.0]))
 
 
+def test_pi05_checkpoint_loader_forwards_hub_options(monkeypatch, tmp_path):
+    import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
+
+    checkpoint = tmp_path / "model.safetensors"
+    save_file({"weight": torch.tensor([1.0])}, checkpoint)
+    calls = []
+
+    def fake_cached_file(model_id, filename, **kwargs):
+        calls.append((model_id, filename, kwargs))
+        return str(checkpoint)
+
+    monkeypatch.setattr(modeling_pi05, "cached_file", fake_cached_file)
+    _CheckpointPolicy.from_pretrained(
+        "org/model",
+        config=SimpleNamespace(),
+        force_download=True,
+        resume_download=True,
+        proxies={"https": "proxy"},
+        token="secret",
+        cache_dir=tmp_path / "cache",
+        local_files_only=True,
+        revision="commit",
+    )
+
+    assert len(calls) == 1
+    model_id, filename, kwargs = calls[0]
+    assert model_id == "org/model"
+    assert filename == "model.safetensors"
+    assert kwargs["revision"] == "commit"
+    assert kwargs["cache_dir"] == tmp_path / "cache"
+    assert kwargs["force_download"] is True
+    assert kwargs["resume_download"] is True
+    assert kwargs["proxies"] == {"https": "proxy"}
+    assert kwargs["token"] == "secret"
+    assert kwargs["local_files_only"] is True
+
+
+def test_pi05_checkpoint_loader_rejects_missing_weights(tmp_path):
+    with pytest.raises(FileNotFoundError, match="model.safetensors"):
+        _CheckpointPolicy.from_pretrained(tmp_path, config=SimpleNamespace())
+
+
+def test_native_checkpoint_uses_standard_lerobot_loader(tmp_path):
+    save_file({"weight": torch.tensor([2.0])}, tmp_path / "model.safetensors")
+
+    policy = _NativeCheckpointPolicy.from_pretrained(
+        tmp_path, config=SimpleNamespace(device="cpu"), strict=True
+    )
+
+    torch.testing.assert_close(policy.weight, torch.tensor([2.0]))
+
+
+def test_pi052_uses_native_checkpoint_loader():
+    from lerobot.policies.pi052.modeling_pi052 import PI052Policy
+
+    assert PI052Policy.use_native_pretrained_loader
+
+
 @require_cuda
 @require_hf_token
 def test_policy_instantiation():