From 9d00293f4980a10c22a89e3adf408f601f5a680b Mon Sep 17 00:00:00 2001 From: pepijn Date: Fri, 24 Jul 2026 09:57:10 +0000 Subject: [PATCH] fix(pi052): make fitted FAST checkpoints portable Package fitted tokenizer artifacts with processor pipelines so pretrained PI052 checkpoints restore their saved recipe and normalization state without refitting. Co-authored-by: Cursor --- scripts/backfill_pi052_action_tokenizer.py | 227 ++++++++++++++++++ src/lerobot/policies/factory.py | 27 +-- .../policies/pi052/configuration_pi052.py | 13 + .../policies/pi052/fit_fast_tokenizer.py | 107 ++++++++- src/lerobot/processor/pipeline.py | 89 ++++++- src/lerobot/processor/tokenizer_processor.py | 9 + tests/envs/test_robocasa_ordering.py | 48 ++++ .../test_pi052_checkpoint_portability.py | 152 ++++++++++++ .../pi052/test_pi052_fit_fast_tokenizer.py | 30 +++ .../test_backfill_pi052_action_tokenizer.py | 67 ++++++ 10 files changed, 734 insertions(+), 35 deletions(-) create mode 100644 scripts/backfill_pi052_action_tokenizer.py create mode 100644 tests/envs/test_robocasa_ordering.py create mode 100644 tests/policies/pi052/test_pi052_checkpoint_portability.py create mode 100644 tests/scripts/test_backfill_pi052_action_tokenizer.py diff --git a/scripts/backfill_pi052_action_tokenizer.py b/scripts/backfill_pi052_action_tokenizer.py new file mode 100644 index 000000000..d97f62d37 --- /dev/null +++ b/scripts/backfill_pi052_action_tokenizer.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Audit or backfill checkpoint-local FAST artifacts for PI052 model repositories.""" + +from __future__ import annotations + +import argparse +import hashlib +import io +import json +from pathlib import Path, PurePosixPath +from typing import Any + +from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download + +DEFAULT_REPOSITORIES = ( + "pepijn223/pi052_atomic4_01_baseline", + "pepijn223/pi052_atomic4_02_lr_1e5", + "pepijn223/pi052_atomic4_03_recipe_50_50", + "pepijn223/pi052_atomic4_04_flow_weight_10", + "pepijn223/pi052_atomic4_05_flow_repeat_1", + "pepijn223/pi052_atomic4_06_ki_off", +) +CHECKPOINT_DIRECTORIES = ( + "", + "checkpoints/003000/pretrained_model", + "checkpoints/006000/pretrained_model", + "checkpoints/009000/pretrained_model", + "checkpoints/012000/pretrained_model", +) +TOKENIZER_DIRECTORY = "action_tokenizer" + + +def artifact_fingerprint(files: list[tuple[str, bytes]]) -> str: + digest = hashlib.sha256() + for relative_path, content in sorted(files): + encoded_path = relative_path.encode() + digest.update(len(encoded_path).to_bytes(8, "big")) + digest.update(encoded_path) + digest.update(len(content).to_bytes(8, "big")) + digest.update(content) + return digest.hexdigest() + + +def tokenizer_files(tokenizer_path: Path) -> list[tuple[str, Path]]: + return [ + (path.relative_to(tokenizer_path).as_posix(), path) + for path in sorted(tokenizer_path.rglob("*")) + if path.is_file() + ] + + +def _repo_path(directory: str, filename: str) -> str: + return (PurePosixPath(directory) / filename).as_posix() if directory else filename + + +def _download_json(repo_id: str, path_in_repo: str, revision: str | None = None) -> dict[str, Any]: + path = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision) + return json.loads(Path(path).read_text()) + + +def make_portable_preprocessor(config: dict[str, Any]) -> dict[str, Any]: + config = json.loads(json.dumps(config)) + action_steps = [ + step for step in config["steps"] if step.get("registry_name") == "action_tokenizer_processor" + ] + if len(action_steps) != 1: + raise ValueError(f"Expected one action tokenizer step, found {len(action_steps)}") + action_step = action_steps[0] + action_step["config"]["action_tokenizer_name"] = TOKENIZER_DIRECTORY + action_step["artifacts"] = {"action_tokenizer_name": TOKENIZER_DIRECTORY} + + recipe_steps = [ + step for step in config["steps"] if step.get("registry_name") == "render_messages_processor" + ] + if len(recipe_steps) != 1 or not recipe_steps[0].get("config", {}).get("recipe"): + raise ValueError("PI052 preprocessor does not contain an embedded training recipe") + return config + + +def _json_operation(path_in_repo: str, content: dict[str, Any]) -> CommitOperationAdd: + serialized = (json.dumps(content, indent=2) + "\n").encode() + return CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=io.BytesIO(serialized)) + + +def prepare_operations( + repo_id: str, + tokenizer_path: Path, + revision: str | None = None, +) -> list[CommitOperationAdd]: + operations: list[CommitOperationAdd] = [] + files = tokenizer_files(tokenizer_path) + for directory in CHECKPOINT_DIRECTORIES: + preprocessor_path = _repo_path(directory, "policy_preprocessor.json") + operations.append( + _json_operation( + preprocessor_path, + make_portable_preprocessor(_download_json(repo_id, preprocessor_path, revision)), + ) + ) + for relative_path, local_path in files: + operations.append( + CommitOperationAdd( + path_in_repo=_repo_path( + directory, + f"{TOKENIZER_DIRECTORY}/{relative_path}", + ), + path_or_fileobj=str(local_path), + ) + ) + return operations + + +def audit_repository( + api: HfApi, + repo_id: str, + expected_tokenizer_fingerprint: str, + revision: str | None = None, +) -> None: + info = api.model_info(repo_id, revision=revision) + repository_files = {sibling.rfilename for sibling in info.siblings or []} + + for directory in CHECKPOINT_DIRECTORIES: + preprocessor_path = _repo_path(directory, "policy_preprocessor.json") + policy_config_path = _repo_path(directory, "config.json") + postprocessor_path = _repo_path(directory, "policy_postprocessor.json") + for required_path in (preprocessor_path, policy_config_path, postprocessor_path): + if required_path not in repository_files: + raise FileNotFoundError(f"{repo_id}@{revision or 'main'} is missing {required_path}") + + preprocessor = _download_json(repo_id, preprocessor_path, revision) + portable_preprocessor = make_portable_preprocessor(preprocessor) + if preprocessor != portable_preprocessor: + raise ValueError(f"{repo_id}:{preprocessor_path} is not portable") + + normalizer_steps = [ + step for step in preprocessor["steps"] if step.get("registry_name") == "normalizer_processor" + ] + if len(normalizer_steps) != 1 or "state_file" not in normalizer_steps[0]: + raise ValueError(f"{repo_id}:{preprocessor_path} is missing normalizer state metadata") + normalizer_path = _repo_path(directory, normalizer_steps[0]["state_file"]) + if normalizer_path not in repository_files: + raise FileNotFoundError(f"{repo_id} is missing {normalizer_path}") + + remote_tokenizer_files: list[tuple[str, bytes]] = [] + for relative_path in _tokenizer_relative_paths(repository_files, directory): + path_in_repo = _repo_path(directory, f"{TOKENIZER_DIRECTORY}/{relative_path}") + downloaded = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision) + remote_tokenizer_files.append((relative_path, Path(downloaded).read_bytes())) + fingerprint = artifact_fingerprint(remote_tokenizer_files) + if fingerprint != expected_tokenizer_fingerprint: + raise ValueError( + f"{repo_id}:{_repo_path(directory, TOKENIZER_DIRECTORY)} fingerprint " + f"{fingerprint} != {expected_tokenizer_fingerprint}" + ) + + +def _tokenizer_relative_paths(repository_files: set[str], directory: str) -> list[str]: + prefix = _repo_path(directory, TOKENIZER_DIRECTORY).rstrip("/") + "/" + paths = sorted(path.removeprefix(prefix) for path in repository_files if path.startswith(prefix)) + if not paths: + raise FileNotFoundError(f"Missing tokenizer artifact directory {prefix.rstrip('/')}") + return paths + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokenizer-path", type=Path, required=True) + parser.add_argument("--repo-id", action="append", dest="repo_ids") + parser.add_argument("--revision") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--audit-only", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + tokenizer_path = args.tokenizer_path.resolve() + if not tokenizer_path.is_dir(): + raise FileNotFoundError(f"Tokenizer directory does not exist: {tokenizer_path}") + + files = tokenizer_files(tokenizer_path) + fingerprint = artifact_fingerprint([(relative_path, path.read_bytes()) for relative_path, path in files]) + api = HfApi() + repositories = tuple(args.repo_ids or DEFAULT_REPOSITORIES) + print(f"Tokenizer fingerprint: {fingerprint}") + + for repo_id in repositories: + if args.audit_only: + audit_repository(api, repo_id, fingerprint, args.revision) + print(f"AUDIT OK {repo_id}@{args.revision or 'main'}") + continue + + operations = prepare_operations(repo_id, tokenizer_path, args.revision) + if args.dry_run: + print(f"DRY RUN {repo_id}: {len(operations)} files") + for operation in operations: + print(f" {operation.path_in_repo}") + continue + + commit = api.create_commit( + repo_id=repo_id, + repo_type="model", + operations=operations, + commit_message="Embed fitted FAST tokenizer for portable PI052 checkpoints", + revision=args.revision, + ) + audit_repository(api, repo_id, fingerprint, commit.oid) + print(f"BACKFILLED {repo_id}@{commit.oid}") + + +if __name__ == "__main__": + main() diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 315812530..a7c778da7 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -315,31 +315,10 @@ def make_pre_post_processors( NotImplementedError: If a processor factory is not implemented for the given policy configuration type. """ - if ( - pretrained_path - and getattr(policy_cfg, "type", None) in {"pi0_fast", "pi052"} - and getattr(policy_cfg, "auto_fit_fast_tokenizer", False) - and kwargs.get("dataset_repo_id") is not None - ): - from .pi052.fit_fast_tokenizer import resolve_fast_tokenizer - - overrides = dict(kwargs.get("preprocessor_overrides") or {}) - action_tokenizer_override = { - **overrides.get("action_tokenizer_processor", {}), - "action_tokenizer_name": resolve_fast_tokenizer( - policy_cfg, - kwargs.get("dataset_repo_id"), - kwargs.get("dataset_root"), - kwargs.get("dataset_stats"), - kwargs.get("dataset_revision"), - kwargs.get("dataset_episodes"), - kwargs.get("dataset_exclude_episodes"), - ), - } - overrides["action_tokenizer_processor"] = action_tokenizer_override - kwargs["preprocessor_overrides"] = overrides - if pretrained_path: + if policy_cfg.type == "pi052": + from .pi052 import processor_pi052 as _processor_pi052 # noqa: F401 + if isinstance(policy_cfg, GrootConfig): from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index c7fe2d260..01acd3b72 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -103,6 +103,15 @@ class PI052Config(PI05Config): fast_tokenizer_fit_samples: int = 1024 """Number of action chunks sampled when fitting FAST.""" + fast_tokenizer_validation_samples: int = 256 + """Held-out action chunks used to validate tokenizer reconstruction.""" + + fast_tokenizer_max_reconstruction_rmse: float = 0.10 + """Maximum normalized RMSE allowed across held-out action chunks.""" + + fast_tokenizer_max_dim_rmse: float = 0.20 + """Maximum normalized RMSE allowed for any nonconstant action dimension.""" + # Knowledge insulation detaches VLM K/V from action-loss gradients (paper §III.B). knowledge_insulation: bool = True """Block action-loss gradients through VLM keys and values.""" @@ -168,6 +177,10 @@ class PI052Config(PI05Config): self.train_expert_only = False if self.flow_num_repeats < 1: raise ValueError(f"flow_num_repeats must be >= 1, got {self.flow_num_repeats}") + if self.fast_tokenizer_validation_samples < 1: + raise ValueError("fast_tokenizer_validation_samples must be >= 1") + if self.fast_tokenizer_max_reconstruction_rmse <= 0 or self.fast_tokenizer_max_dim_rmse <= 0: + raise ValueError("FAST tokenizer reconstruction thresholds must be positive") if self.manual_attention_scope not in {"all", "action"}: raise ValueError( f"manual_attention_scope must be 'all' or 'action', got {self.manual_attention_scope!r}" diff --git a/src/lerobot/policies/pi052/fit_fast_tokenizer.py b/src/lerobot/policies/pi052/fit_fast_tokenizer.py index 2098deaed..b6736258b 100644 --- a/src/lerobot/policies/pi052/fit_fast_tokenizer.py +++ b/src/lerobot/policies/pi052/fit_fast_tokenizer.py @@ -64,6 +64,9 @@ def _dataset_signature( action_stats: dict | None = None, use_relative_actions: bool = False, relative_action_mask: list[bool] | None = None, + validation_samples: int = 256, + max_reconstruction_rmse: float = 0.10, + max_dim_rmse: float = 0.20, ) -> str: """Hash every input that changes the fitted action distribution.""" payload = { @@ -78,6 +81,9 @@ def _dataset_signature( "action_stats": action_stats, "use_relative_actions": use_relative_actions, "relative_action_mask": relative_action_mask, + "validation_samples": validation_samples, + "max_reconstruction_rmse": max_reconstruction_rmse, + "max_dim_rmse": max_dim_rmse, } encoded = json.dumps(_jsonable(payload), sort_keys=True, separators=(",", ":")).encode() return hashlib.sha256(encoded).hexdigest()[:16] @@ -148,6 +154,64 @@ def _normalize_actions( return (2.0 * (actions - low) / np.where(high == low, 1e-8, high - low) - 1.0).astype(np.float32) +def _validate_fast_reconstruction( + tokenizer: Any, + actions: np.ndarray, + max_reconstruction_rmse: float, + max_dim_rmse: float, +) -> tuple[dict[str, Any], np.ndarray]: + """Decode held-out chunks and reject tokenizers with excessive quantization error.""" + decoded = np.asarray(tokenizer.decode(tokenizer(actions)), dtype=np.float32) + if decoded.shape != actions.shape: + raise RuntimeError( + f"FAST tokenizer reconstruction shape mismatch: expected {actions.shape}, got {decoded.shape}." + ) + if not np.isfinite(decoded).all(): + raise RuntimeError("FAST tokenizer reconstruction contains non-finite values.") + + squared_error = np.square(decoded - actions) + rmse = float(np.sqrt(squared_error.mean())) + dim_rmse = np.sqrt(squared_error.mean(axis=(0, 1))) + nonconstant_dims = np.ptp(actions, axis=(0, 1)) > 1e-8 + max_observed_dim_rmse = float(dim_rmse[nonconstant_dims].max(initial=0.0)) + report = { + "num_validation_chunks": int(actions.shape[0]), + "reconstruction_rmse": rmse, + "max_dim_rmse": max_observed_dim_rmse, + "dim_rmse": dim_rmse.tolist(), + "max_reconstruction_rmse": max_reconstruction_rmse, + "max_allowed_dim_rmse": max_dim_rmse, + } + if rmse > max_reconstruction_rmse or max_observed_dim_rmse > max_dim_rmse: + raise RuntimeError( + "FAST tokenizer reconstruction error exceeds the configured limit: " + f"rmse={rmse:.4f} (max {max_reconstruction_rmse:.4f}), " + f"max_dim_rmse={max_observed_dim_rmse:.4f} (max {max_dim_rmse:.4f})." + ) + return report, decoded + + +def _load_fast_fitter(base_tokenizer_name: str) -> Any: + """Load FAST's fitting implementation without requiring its universal BPE weights.""" + from transformers import AutoProcessor # noqa: PLC0415 + + try: + return AutoProcessor.from_pretrained(base_tokenizer_name, trust_remote_code=True) + except ValueError as error: + if base_tokenizer_name != "physical-intelligence/fast": + raise + logger.warning( + "Could not load the universal FAST tokenizer backend; loading its fitting class directly: %s", + error, + ) + from transformers.dynamic_module_utils import get_class_from_dynamic_module # noqa: PLC0415 + + return get_class_from_dynamic_module( + "processing_action_tokenizer.UniversalActionProcessor", + base_tokenizer_name, + ) + + def fit_fast_tokenizer( *, dataset_repo_id: str, @@ -164,6 +228,9 @@ def fit_fast_tokenizer( action_stats: dict | None = None, use_relative_actions: bool = False, relative_action_mask: list[bool] | None = None, + validation_samples: int = 256, + max_reconstruction_rmse: float = 0.10, + max_dim_rmse: float = 0.20, ) -> str: """Fit a FAST tokenizer on a LeRobot dataset's action distribution. @@ -208,6 +275,9 @@ def fit_fast_tokenizer( action_stats, use_relative_actions, relative_action_mask, + validation_samples, + max_reconstruction_rmse, + max_dim_rmse, ) out_dir = cache_dir / sig @@ -246,8 +316,6 @@ def fit_fast_tokenizer( out_dir, ) - from transformers import AutoProcessor # noqa: PLC0415 - # Read action columns directly to avoid video decoding and bound memory to sampled episodes. rng = np.random.default_rng(seed) actions_buf: list[np.ndarray] = [] @@ -313,13 +381,14 @@ def fit_fast_tokenizer( ep_indices = _select_episode_indices(list(ep_to_slice), episodes, exclude_episodes) if not ep_indices: raise RuntimeError("FAST fit: episode selection is empty after applying exclusions.") - samples_per_episode = max(1, n_samples // len(ep_indices)) + total_samples = n_samples + validation_samples + samples_per_episode = max(1, (total_samples + len(ep_indices) - 1) // len(ep_indices)) collected = 0 eps_visited = 0 short_episodes = 0 states_buf: list[np.ndarray] = [] for ep_idx in rng.permutation(ep_indices): - if collected >= n_samples: + if collected >= total_samples: break start, stop = ep_to_slice[int(ep_idx)] ep_actions = acts[start:stop] @@ -332,7 +401,7 @@ def fit_fast_tokenizer( if states is not None: states_buf.append(states[start + int(s)]) collected += 1 - if collected >= n_samples: + if collected >= total_samples: break eps_visited += 1 @@ -357,7 +426,7 @@ def fit_fast_tokenizer( actions = _normalize_actions(actions, normalization_mode, action_stats) - base = AutoProcessor.from_pretrained(base_tokenizer_name, trust_remote_code=True) + base = _load_fast_fitter(base_tokenizer_name) if not hasattr(base, "fit"): raise ImportError( f"Base FAST tokenizer {base_tokenizer_name!r} has no ``.fit()`` " @@ -365,11 +434,32 @@ def fit_fast_tokenizer( "to the current ``physical-intelligence/fast`` revision." ) - fitted = base.fit(actions) + if actions.shape[0] < total_samples: + raise RuntimeError( + f"FAST fit collected {actions.shape[0]} chunks, but {total_samples} are required " + f"for {n_samples} fit and {validation_samples} validation chunks." + ) + fit_actions = actions[:n_samples] + validation_actions = actions[n_samples:total_samples] + fitted = base.fit(fit_actions) + validation_report, decoded_actions = _validate_fast_reconstruction( + fitted, + validation_actions, + max_reconstruction_rmse, + max_dim_rmse, + ) cache_dir.mkdir(parents=True, exist_ok=True) staging_dir = cache_dir / f".{sig}.tmp-{os.getpid()}" shutil.rmtree(staging_dir, ignore_errors=True) fitted.save_pretrained(str(staging_dir)) + (staging_dir / "reconstruction_validation.json").write_text( + json.dumps(validation_report, indent=2) + "\n" + ) + np.savez_compressed( + staging_dir / "reconstruction_examples.npz", + original=validation_actions[:8], + decoded=decoded_actions[:8], + ) if out_dir.exists(): shutil.rmtree(out_dir) staging_dir.replace(out_dir) @@ -416,4 +506,7 @@ def resolve_fast_tokenizer( action_stats=(dataset_stats or {}).get("action"), use_relative_actions=getattr(config, "use_relative_actions", False), relative_action_mask=relative_action_mask, + validation_samples=config.fast_tokenizer_validation_samples, + max_reconstruction_rmse=config.fast_tokenizer_max_reconstruction_rmse, + max_dim_rmse=config.fast_tokenizer_max_dim_rmse, ) diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index b9b9c6c43..1e44ef634 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -41,7 +41,7 @@ from pathlib import Path from typing import Any, TypedDict, TypeVar, cast import torch -from huggingface_hub import hf_hub_download +from huggingface_hub import hf_hub_download, snapshot_download from safetensors.torch import load_file, save_file from lerobot.configs import PipelineFeatureType, PolicyFeature @@ -205,6 +205,10 @@ class ProcessorStep(ABC): """ return None + def save_artifacts(self, save_directory: Path) -> dict[str, str]: + """Save non-tensor assets and map constructor arguments to relative paths.""" + return {} + def reset(self) -> None: """Resets the internal state of the processor step, if any.""" return None @@ -549,6 +553,22 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): pipeline_config = self.get_config() pipeline_state_dict = self.state_dict() + for processor_step, step_entry in zip(self.steps, pipeline_config["steps"], strict=True): + artifacts = processor_step.save_artifacts(save_directory) + if artifacts: + for config_key, relative_path in artifacts.items(): + artifact_path = Path(relative_path) + if artifact_path.is_absolute() or ".." in artifact_path.parts: + raise ValueError( + f"Processor artifact path must be relative to the checkpoint: {relative_path!r}" + ) + if not (save_directory / artifact_path).exists(): + raise FileNotFoundError( + f"Processor step did not save declared artifact '{relative_path}'" + ) + step_entry["config"][config_key] = artifact_path.as_posix() + step_entry["artifacts"] = artifacts + for state_key, step_state_dict in pipeline_state_dict.items(): state_filename = f"{state_key}.safetensors" save_file(step_state_dict, save_directory / state_filename) @@ -731,7 +751,12 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): # 3. Build steps with overrides steps, validated_overrides = cls._build_steps_with_overrides( - loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs + loaded_config, + overrides or {}, + model_id, + base_path, + config_filename, + hub_download_kwargs, ) # 4. Validate that all overrides were used @@ -920,6 +945,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): overrides: dict[str, Any], model_id: str, base_path: Path | None, + config_filename: str, hub_download_kwargs: dict[str, Any], ) -> tuple[list[ProcessorStep], set[str]]: """Build all processor steps with overrides and state loading. @@ -972,13 +998,67 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): ImportError: If a step class cannot be imported or found in registry ValueError: If a step cannot be instantiated with its configuration """ + loaded_config = deepcopy(loaded_config) + cls._resolve_artifact_paths( + loaded_config, + model_id, + base_path, + config_filename, + hub_download_kwargs, + ) steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides) for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True): - cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs) + cls._load_step_state( + step_instance, + step_entry, + model_id, + base_path, + config_filename, + hub_download_kwargs, + ) return steps, remaining_override_keys + @classmethod + def _resolve_artifact_paths( + cls, + loaded_config: dict[str, Any], + model_id: str, + base_path: Path | None, + config_filename: str, + hub_download_kwargs: dict[str, Any], + ) -> None: + """Resolve declared relative processor artifacts before step construction.""" + is_local = Path(model_id).is_dir() or Path(model_id).is_file() + + for step_entry in loaded_config["steps"]: + artifacts = step_entry.get("artifacts", {}) + for config_key, relative_path in artifacts.items(): + artifact_path = Path(relative_path) + if artifact_path.is_absolute() or ".." in artifact_path.parts: + raise ValueError( + f"Processor artifact path must be relative to the checkpoint: {relative_path!r}" + ) + + resolved_path = base_path / artifact_path if base_path is not None else artifact_path + if not resolved_path.exists() and not is_local: + repository_path = Path(config_filename).parent / artifact_path + snapshot_download( + repo_id=model_id, + repo_type="model", + allow_patterns=f"{repository_path.as_posix()}/**", + **hub_download_kwargs, + ) + + if not resolved_path.exists(): + step_name = step_entry.get("registry_name", step_entry.get("class", "unknown")) + raise FileNotFoundError( + f"Missing processor artifact '{relative_path}' for step '{step_name}' " + f"next to '{config_filename}'. Checkpoint artifacts are incomplete." + ) + step_entry["config"][config_key] = str(resolved_path) + @classmethod def _build_steps_from_config( cls, @@ -1138,6 +1218,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): step_entry: dict[str, Any], model_id: str, base_path: Path | None, + config_filename: str, hub_download_kwargs: dict[str, Any], ) -> None: """Load state dictionary for a processor step if available. @@ -1195,7 +1276,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): # Download from Hub state_path = hf_hub_download( repo_id=model_id, - filename=state_filename, + filename=(Path(config_filename).parent / state_filename).as_posix(), repo_type="model", **hub_download_kwargs, ) diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py index ff4c2d77c..45a4c77f5 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -25,6 +25,7 @@ from __future__ import annotations import logging from dataclasses import dataclass, field +from pathlib import Path from typing import TYPE_CHECKING, Any import torch @@ -582,6 +583,14 @@ class ActionTokenizerProcessorStep(ActionProcessorStep): return config + def save_artifacts(self, save_directory: Path) -> dict[str, str]: + artifact_path = Path("action_tokenizer") + save_pretrained = getattr(self.action_tokenizer, "save_pretrained", None) + if save_pretrained is None: + raise TypeError("Action tokenizer must implement save_pretrained() to save a portable pipeline.") + save_pretrained(save_directory / artifact_path) + return {"action_tokenizer_name": artifact_path.as_posix()} + def transform_features( self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: diff --git a/tests/envs/test_robocasa_ordering.py b/tests/envs/test_robocasa_ordering.py new file mode 100644 index 000000000..c4fbe9041 --- /dev/null +++ b/tests/envs/test_robocasa_ordering.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np + +from lerobot.envs.robocasa import RoboCasaEnv, convert_action + + +def test_robocasa_action_uses_openpi_checkpoint_order(): + action = np.arange(12, dtype=np.float32) + + converted = convert_action(action) + + np.testing.assert_array_equal(converted["action.end_effector_position"], [0, 1, 2]) + np.testing.assert_array_equal(converted["action.end_effector_rotation"], [3, 4, 5]) + np.testing.assert_array_equal(converted["action.gripper_close"], [6]) + np.testing.assert_array_equal(converted["action.base_motion"], [7, 8, 9, 10]) + np.testing.assert_array_equal(converted["action.control_mode"], [11]) + + +def test_robocasa_state_uses_openpi_checkpoint_order(): + env = object.__new__(RoboCasaEnv) + env.obs_type = "pixels_agent_pos" + env.camera_name = [] + raw_observation = { + "state.end_effector_position_relative": np.arange(0, 3), + "state.end_effector_rotation_relative": np.arange(3, 7), + "state.base_position": np.arange(7, 10), + "state.base_rotation": np.arange(10, 14), + "state.gripper_qpos": np.arange(14, 16), + } + + observation = env._format_raw_obs(raw_observation) + + np.testing.assert_array_equal(observation["agent_pos"], np.arange(16, dtype=np.float32)) diff --git a/tests/policies/pi052/test_pi052_checkpoint_portability.py b/tests/policies/pi052/test_pi052_checkpoint_portability.py new file mode 100644 index 000000000..824dc5263 --- /dev/null +++ b/tests/policies/pi052/test_pi052_checkpoint_portability.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import shutil +from dataclasses import asdict +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature +from lerobot.configs.recipe import MessageTurn, TrainingRecipe +from lerobot.policies import make_pre_post_processors +from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, NormalizerProcessorStep +from lerobot.processor.converters import identity_transition +from lerobot.processor.render_messages_processor import RenderMessagesStep +from lerobot.utils.constants import ACTION + + +class _ActionTokenizer: + def __call__(self, actions): + return np.asarray(actions).round().astype(np.int64) + + def save_pretrained(self, path): + path.mkdir(parents=True) + (path / "processor_config.json").write_text('{"processor_class": "_ActionTokenizer"}\n') + + +class _PaligemmaTokenizer: + vocab_size = 4096 + bos_token_id = 2 + + def encode(self, text, **kwargs): + return [10, 11] if text == "Action: " else [12] + + +def _make_pipeline(action_tokenizer_path): + recipe = TrainingRecipe( + messages=[ + MessageTurn(role="user", content="${task}", stream="high_level"), + MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True), + ] + ) + stats = {ACTION: {"min": torch.tensor([-1.0, -2.0]), "max": torch.tensor([1.0, 2.0])}} + normalizer = NormalizerProcessorStep( + features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,))}, + norm_map={FeatureType.ACTION: NormalizationMode.MIN_MAX}, + stats=stats, + ) + action_tokenizer = ActionTokenizerProcessorStep( + action_tokenizer_name=str(action_tokenizer_path), + max_action_tokens=16, + fast_skip_tokens=128, + ) + return DataProcessorPipeline( + [normalizer, RenderMessagesStep(recipe), action_tokenizer], + name="policy_preprocessor", + to_transition=identity_transition, + to_output=identity_transition, + ) + + +def test_pi052_pipeline_embeds_and_loads_fitted_action_tokenizer(tmp_path, monkeypatch): + original_cache = tmp_path / "original_fast_cache" + original_cache.mkdir() + tokenizer = _ActionTokenizer() + monkeypatch.setattr( + "lerobot.processor.tokenizer_processor.AutoProcessor.from_pretrained", + lambda path, **kwargs: tokenizer, + ) + monkeypatch.setattr( + "lerobot.processor.tokenizer_processor.AutoTokenizer.from_pretrained", + lambda *args, **kwargs: _PaligemmaTokenizer(), + ) + monkeypatch.setattr( + "lerobot.policies.pi052.fit_fast_tokenizer.fit_fast_tokenizer", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("FAST fitting must not run")), + ) + + pipeline = _make_pipeline(original_cache) + expected_tokens = pipeline.steps[-1]._tokenize_action(torch.tensor([[[0.2, 0.8]]]))[0] + expected_recipe = asdict(pipeline.steps[1].recipe) + expected_state = pipeline.steps[0].state_dict() + checkpoint = tmp_path / "checkpoint" + pipeline.save_pretrained(checkpoint) + DataProcessorPipeline( + [], + name="policy_postprocessor", + to_transition=identity_transition, + to_output=identity_transition, + ).save_pretrained(checkpoint) + + saved_config = json.loads((checkpoint / "policy_preprocessor.json").read_text()) + tokenizer_step = saved_config["steps"][2] + assert tokenizer_step["config"]["action_tokenizer_name"] == "action_tokenizer" + assert tokenizer_step["artifacts"] == {"action_tokenizer_name": "action_tokenizer"} + assert (checkpoint / "action_tokenizer" / "processor_config.json").is_file() + + shutil.rmtree(original_cache) + loaded, _ = make_pre_post_processors( + SimpleNamespace(type="pi052", auto_fit_fast_tokenizer=True), + pretrained_path=str(checkpoint), + dataset_repo_id="org/dataset-that-must-not-be-read", + ) + + assert asdict(loaded.steps[1].recipe) == expected_recipe + for key, tensor in expected_state.items(): + torch.testing.assert_close(loaded.steps[0].state_dict()[key], tensor) + torch.testing.assert_close( + loaded.steps[-1]._tokenize_action(torch.tensor([[[0.2, 0.8]]]))[0], + expected_tokens, + ) + + +def test_pi052_pipeline_rejects_missing_fitted_action_tokenizer(tmp_path, monkeypatch): + tokenizer = _ActionTokenizer() + monkeypatch.setattr( + "lerobot.processor.tokenizer_processor.AutoProcessor.from_pretrained", + lambda path, **kwargs: tokenizer, + ) + monkeypatch.setattr( + "lerobot.processor.tokenizer_processor.AutoTokenizer.from_pretrained", + lambda *args, **kwargs: _PaligemmaTokenizer(), + ) + + pipeline = _make_pipeline(tmp_path / "original_fast_cache") + checkpoint = tmp_path / "checkpoint" + pipeline.save_pretrained(checkpoint) + shutil.rmtree(checkpoint / "action_tokenizer") + + with pytest.raises(FileNotFoundError, match="Checkpoint artifacts are incomplete"): + DataProcessorPipeline.from_pretrained( + checkpoint, + config_filename="policy_preprocessor.json", + to_transition=identity_transition, + to_output=identity_transition, + ) diff --git a/tests/policies/pi052/test_pi052_fit_fast_tokenizer.py b/tests/policies/pi052/test_pi052_fit_fast_tokenizer.py index 0f4949172..b4c86d180 100644 --- a/tests/policies/pi052/test_pi052_fit_fast_tokenizer.py +++ b/tests/policies/pi052/test_pi052_fit_fast_tokenizer.py @@ -15,6 +15,7 @@ # limitations under the License. import numpy as np +import pytest from lerobot.policies.pi052.fit_fast_tokenizer import ( _apply_relative_actions, @@ -22,6 +23,7 @@ from lerobot.policies.pi052.fit_fast_tokenizer import ( _is_global_leader, _normalize_actions, _select_episode_indices, + _validate_fast_reconstruction, ) @@ -90,3 +92,31 @@ def test_fast_tokenizer_relative_actions_match_training_transform(): relative = _apply_relative_actions(actions, states, [True, False]) np.testing.assert_allclose(relative, [[[1.0, 10.0], [2.0, 11.0]]]) + + +class _RoundTripTokenizer: + def __init__(self, offset: float = 0.0): + self.offset = offset + + def __call__(self, actions): + return actions + + def decode(self, tokens): + return tokens + self.offset + + +def test_fast_tokenizer_reconstruction_validation_reports_error(): + actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24 + + report, decoded = _validate_fast_reconstruction(_RoundTripTokenizer(0.05), actions, 0.1, 0.1) + + np.testing.assert_allclose(decoded, actions + 0.05) + assert report["reconstruction_rmse"] == pytest.approx(0.05) + assert report["max_dim_rmse"] == pytest.approx(0.05) + + +def test_fast_tokenizer_reconstruction_validation_rejects_large_error(): + actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24 + + with pytest.raises(RuntimeError, match="exceeds the configured limit"): + _validate_fast_reconstruction(_RoundTripTokenizer(0.25), actions, 0.1, 0.2) diff --git a/tests/scripts/test_backfill_pi052_action_tokenizer.py b/tests/scripts/test_backfill_pi052_action_tokenizer.py new file mode 100644 index 000000000..d3d664785 --- /dev/null +++ b/tests/scripts/test_backfill_pi052_action_tokenizer.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from scripts.backfill_pi052_action_tokenizer import ( + CHECKPOINT_DIRECTORIES, + DEFAULT_REPOSITORIES, + artifact_fingerprint, + make_portable_preprocessor, +) + + +def test_atomic4_backfill_covers_every_repository_and_checkpoint(): + assert len(DEFAULT_REPOSITORIES) == 6 + assert CHECKPOINT_DIRECTORIES == ( + "", + "checkpoints/003000/pretrained_model", + "checkpoints/006000/pretrained_model", + "checkpoints/009000/pretrained_model", + "checkpoints/012000/pretrained_model", + ) + + +def test_backfill_embeds_recipe_and_declares_relative_tokenizer(): + recipe = {"messages": [{"role": "user", "content": "${task}", "stream": "low_level"}]} + preprocessor = { + "name": "policy_preprocessor", + "steps": [ + { + "registry_name": "normalizer_processor", + "config": {}, + "state_file": "normalizer.safetensors", + }, + {"registry_name": "render_messages_processor", "config": {"recipe": recipe}}, + { + "registry_name": "action_tokenizer_processor", + "config": {"action_tokenizer_name": "/fsx/original/tokenizer"}, + }, + ], + } + + portable = make_portable_preprocessor(preprocessor) + + assert portable["steps"][1]["config"]["recipe"] == recipe + assert portable["steps"][2]["config"]["action_tokenizer_name"] == "action_tokenizer" + assert portable["steps"][2]["artifacts"] == {"action_tokenizer_name": "action_tokenizer"} + assert preprocessor["steps"][2]["config"]["action_tokenizer_name"].startswith("/fsx/") + + +def test_artifact_fingerprint_includes_paths_and_contents(): + first = artifact_fingerprint([("a/file", b"same"), ("b/file", b"content")]) + + assert first == artifact_fingerprint([("b/file", b"content"), ("a/file", b"same")]) + assert first != artifact_fingerprint([("a/renamed", b"same"), ("b/file", b"content")]) + assert first != artifact_fingerprint([("a/file", b"changed"), ("b/file", b"content")])