diff --git a/docs/source/pi052.mdx b/docs/source/pi052.mdx index dbea9bd69..ee697b3a7 100644 --- a/docs/source/pi052.mdx +++ b/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 \ diff --git a/src/lerobot/policies/pi05/modeling_pi05.py b/src/lerobot/policies/pi05/modeling_pi05.py index 9ce2d80f1..4df7c5f53 100644 --- a/src/lerobot/policies/pi05/modeling_pi05.py +++ b/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() diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index 47d7af249..abfafefe8 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/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. diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 90d894e30..cf800bf69 100644 --- a/src/lerobot/policies/pretrained.py +++ b/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", } diff --git a/tests/policies/pi052/test_pi052_bucketed_ce.py b/tests/policies/pi052/test_pi052_bucketed_ce.py index 717874783..0cb16fbb3 100644 --- a/tests/policies/pi052/test_pi052_bucketed_ce.py +++ b/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): diff --git a/tests/policies/pi0_pi05/test_pi05.py b/tests/policies/pi0_pi05/test_pi05.py index 44327fb6c..04b27eb5c 100644 --- a/tests/policies/pi0_pi05/test_pi05.py +++ b/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():