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 <cursoragent@cursor.com>

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()

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
 

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}"

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,
     )

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,
             )

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]]:

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))

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,
+        )

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)

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")])