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>

docs/source/language_and_recipes.mdx

@@ -146,6 +146,12 @@ The renderer does not apply a tokenizer chat template. Policy processors decide
 Blend recipes select one weighted sub-recipe deterministically from the sample index.
 `recipes/subtask_mem.yaml` trains the compact core blend — high-level subtask prediction, low-level execution, and memory. `recipes/subtask_mem_vqa_speech.yaml` is the fuller variant that also adds VQA and spoken interjection responses.
 
+A message recipe with a supervised assistant turn on the `low_level` stream trains
+the π0.5 paper's joint sequence instead of a blend: the target span gets text CE
+while also conditioning the action losses in the same forward.
+`recipes/subtask_joint.yaml` is the provided example; pair it with
+`--policy.joint_subtask_conditioning=true` at inference.
+
 ## Graceful absence
 
 If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op.

docs/source/pi052.mdx

@@ -55,9 +55,20 @@ The provided recipes are:
 | Recipe                                | Required annotations                                                    | Trains                                             |
 | ------------------------------------- | ----------------------------------------------------------------------- | -------------------------------------------------- |
 | `recipes/subtask.yaml`                | `subtask`                                                               | Subtask prediction and subtask-conditioned actions |
+| `recipes/subtask_joint.yaml`          | `subtask`                                                               | Paper-style joint sequence: subtask text and actions in one sample |
 | `recipes/subtask_mem.yaml`            | `subtask`, `memory`                                                     | Subtasks, actions, and memory updates              |
 | `recipes/subtask_mem_vqa_speech.yaml` | `subtask`, `memory`, `vqa`; interjection/speech rows for those branches | Subtasks, actions, memory, VQA, and spoken replies |
 
+The blend recipes factorize training into separate high-level (task → subtask)
+and low-level (subtask → actions) samples, matching how inference decomposes
+π(a|o, subtask)·π(subtask|o, task). `recipes/subtask_joint.yaml` instead uses
+the π0.5 paper's single-sequence layout — the supervised subtask span is
+attended causally and conditions the FAST and flow losses in the same forward.
+Checkpoints trained with the joint recipe must set
+`--policy.joint_subtask_conditioning=true` at inference so the flow prefix
+rebuilds the same layout (task turn with state, then the generated subtask as a
+causal assistant turn); leave it `false` for the blend recipes.
+
 Use `lerobot-annotate` to generate these columns. The repository includes a
 Hugging Face Jobs launcher that you can edit for your source and destination
 datasets. For a local annotation run, first install
@@ -117,6 +128,14 @@ the expected prompt, text target, and action endpoints before scaling up.
 | `policy.knowledge_insulation`    |                     `true` | Blocks action-loss gradients through the VLM K/V path          |
 | `policy.flow_num_repeats`        |                        `5` | Reuses one VLM prefix for independent denoising targets        |
 | `policy.lm_head_lr_scale`        |                      `1.0` | Scales language-head learning rate; `1.0` uses the base rate   |
+| `policy.fast_skip_tokens`        |                     `1152` | FAST id offset; skips `<seg>`+`<loc>` so VQA and FAST never collide |
+| `policy.joint_subtask_conditioning` |                  `false` | Rebuilds the joint-sequence prefix at inference (see recipes)  |
+
+`fast_skip_tokens=1152` places FAST codes below PaliGemma's `<loc>` range.
+openpi's pi0-FAST convention is `128` (FAST occupies the `<loc>` ids); use that
+value only to stay weight-compatible with checkpoints trained that way, and
+avoid combining it with the VQA recipe, whose `<loc>` targets would share
+embedding rows with FAST codes.
 
 The loss weights are starting points, not dataset-independent constants. Track
 flow loss and text/FAST losses separately, and inspect generated subtasks rather

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/configs/default.py

@@ -33,6 +33,8 @@ class DatasetConfig:
     # looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
     root: str | None = None
     episodes: list[int] | None = None
+    # Episode indices to drop (e.g. corrupt or heterogeneous ones). Applied on top of `episodes`.
+    exclude_episodes: list[int] | None = None
     image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
     revision: str | None = None
     use_imagenet_stats: bool = True
@@ -62,6 +64,10 @@ class DatasetConfig:
             if len(self.episodes) != len(set(self.episodes)):
                 duplicates = sorted({ep for ep in self.episodes if self.episodes.count(ep) > 1})
                 raise ValueError(f"Episode indices contain duplicates: {duplicates}")
+        if self.exclude_episodes is not None and any(ep < 0 for ep in self.exclude_episodes):
+            raise ValueError(
+                f"exclude_episodes must be non-negative, got: {[ep for ep in self.exclude_episodes if ep < 0]}"
+            )
 
 
 @dataclass

src/lerobot/configs/recipes/subtask_joint.yaml

@@ -0,0 +1,13 @@
+# Paper-style joint sequence (pi0.5 §IV-B): one sample supervises the subtask
+# text with CE and, because the assistant turn is part of the prefix, conditions
+# the FAST and flow action losses on the same annotated subtask in one forward.
+# The supervised span is attended causally; the action losses see task + subtask.
+#
+# Pair with `--policy.joint_subtask_conditioning=true` at inference so the flow
+# prefix reproduces this layout (task turn with state + causal generated subtask).
+# Samples without a `subtask` annotation fall back to a plain task-prompt
+# low-level sample via `if_present`.
+
+messages:
+  - {role: user, content: "${task}", stream: low_level}
+  - {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask}

src/lerobot/datasets/factory.py

@@ -66,6 +66,17 @@ def resolve_delta_timestamps(
     return delta_timestamps
 
 
+def _resolve_episodes(
+    episodes: list[int] | None, exclude_episodes: list[int] | None, total_episodes: int
+) -> list[int] | None:
+    """Apply an episode exclusion list on top of an optional allowlist."""
+    if not exclude_episodes:
+        return episodes
+    base = episodes if episodes is not None else list(range(total_episodes))
+    excluded = set(exclude_episodes)
+    return [episode for episode in base if episode not in excluded]
+
+
 def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
     """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
 
@@ -87,11 +98,14 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
         )
         delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
+        episodes = _resolve_episodes(
+            cfg.dataset.episodes, cfg.dataset.exclude_episodes, ds_meta.total_episodes
+        )
         if not cfg.dataset.streaming:
             dataset = LeRobotDataset(
                 cfg.dataset.repo_id,
                 root=cfg.dataset.root,
-                episodes=cfg.dataset.episodes,
+                episodes=episodes,
                 delta_timestamps=delta_timestamps,
                 image_transforms=image_transforms,
                 revision=cfg.dataset.revision,
@@ -104,7 +118,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             dataset = StreamingLeRobotDataset(
                 cfg.dataset.repo_id,
                 root=cfg.dataset.root,
-                episodes=cfg.dataset.episodes,
+                episodes=episodes,
                 delta_timestamps=delta_timestamps,
                 image_transforms=image_transforms,
                 revision=cfg.dataset.revision,

src/lerobot/policies/factory.py

@@ -277,6 +277,10 @@ class ProcessorConfigKwargs(TypedDict, total=False):
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None
     # Dataset repo used for optional processor fitting; omit it to use universal tokenizers.
     dataset_repo_id: str | None
+    dataset_root: str | None
+    dataset_revision: str | None
+    dataset_episodes: list[int] | None
+    dataset_exclude_episodes: list[int] | None
     dataset_meta: Any | None
 
 
@@ -311,23 +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")),
-        }
-        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
 
@@ -435,6 +426,10 @@ def make_pre_post_processors(
             config=policy_cfg,
             dataset_stats=kwargs.get("dataset_stats"),
             dataset_repo_id=kwargs.get("dataset_repo_id"),
+            dataset_root=kwargs.get("dataset_root"),
+            dataset_revision=kwargs.get("dataset_revision"),
+            episodes=kwargs.get("dataset_episodes"),
+            exclude_episodes=kwargs.get("dataset_exclude_episodes"),
         )
 
     elif policy_cfg.type == "pi052":
@@ -446,6 +441,10 @@ def make_pre_post_processors(
             dataset_stats=kwargs.get("dataset_stats"),
             # Without a dataset repo, FAST auto-fit falls back to the universal tokenizer.
             dataset_repo_id=kwargs.get("dataset_repo_id"),
+            dataset_root=kwargs.get("dataset_root"),
+            dataset_revision=kwargs.get("dataset_revision"),
+            episodes=kwargs.get("dataset_episodes"),
+            exclude_episodes=kwargs.get("dataset_exclude_episodes"),
         )
 
     elif isinstance(policy_cfg, PI05Config):

src/lerobot/policies/pi052/configuration_pi052.py

@@ -63,9 +63,14 @@ class PI052Config(PI05Config):
     max_action_tokens: int = 256
     """Maximum number of FAST tokens per action chunk."""
 
-    fast_skip_tokens: int = 128
-    """Number of low-vocab tokens the FAST tokenizer skips to avoid
-    collisions with PaliGemma's text vocabulary."""
+    fast_skip_tokens: int = 1152
+    """Number of top-of-vocab tokens the FAST id mapping skips.
+
+    1152 skips PaliGemma's 128 ``<seg>`` and 1024 ``<loc>`` special tokens so
+    FAST codes land in plain-text ids below 256000 and never collide with the
+    ``<loc>`` targets used for VQA. openpi's pi0-FAST convention is 128 (FAST
+    occupies the ``<loc>`` range); use 128 only to stay weight-compatible with
+    checkpoints trained that way."""
 
     fast_action_loss_weight: float = 1.0
     """Weight on FAST action-token CE relative to continuous-flow supervision."""
@@ -76,6 +81,16 @@ class PI052Config(PI05Config):
     Non-positive values regenerate each action chunk while still refreshing the action prompt every chunk.
     """
 
+    joint_subtask_conditioning: bool = False
+    """Condition low-level action inference on the task plus the generated subtask.
+
+    Matches paper-style joint-sequence recipes (``recipes/subtask_joint.yaml``)
+    where one sample supervises the subtask text and conditions the action
+    losses on it: the inference prefix becomes
+    ``User: {task}, State: ...;\\nAssistant: {subtask}<eos>`` with the subtask
+    span attended causally, exactly as trained. Leave ``False`` for the blend
+    recipes, whose low-level samples use ``User: {subtask}, State: ...;``."""
+
     auto_fit_fast_tokenizer: bool = False
     """Fit and cache a dataset-specific FAST tokenizer before training.
 
@@ -88,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."""
@@ -147,10 +171,16 @@ class PI052Config(PI05Config):
 
     def __post_init__(self) -> None:
         super().__post_init__()
+        if self.enable_fast_action_loss and not self.recipe_path:
+            raise ValueError("PI052 FAST action loss requires recipe_path to build action supervision.")
         if self.text_loss_weight > 0 and self.unfreeze_lm_head:
             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

@@ -20,8 +20,10 @@ Training invokes this automatically when FAST loss and automatic fitting are ena
 from __future__ import annotations
 
 import hashlib
+import json
 import logging
 import os
+import shutil
 import time
 from pathlib import Path
 from typing import Any
@@ -34,8 +36,20 @@ logger = logging.getLogger(__name__)
 _CACHE_SENTINEL = "processor_config.json"
 
 
-def _is_local_leader() -> bool:
-    return int(os.environ.get("LOCAL_RANK", "0")) == 0
+def _is_global_leader() -> bool:
+    return int(os.environ.get("RANK", "0")) == 0
+
+
+def _jsonable(value: Any) -> Any:
+    if hasattr(value, "detach"):
+        value = value.detach().cpu().numpy()
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    if isinstance(value, dict):
+        return {key: _jsonable(item) for key, item in sorted(value.items())}
+    if isinstance(value, (list, tuple)):
+        return [_jsonable(item) for item in value]
+    return value
 
 
 def _dataset_signature(
@@ -43,22 +57,159 @@ def _dataset_signature(
     base_tokenizer_name: str,
     n_samples: int,
     chunk_size: int,
+    normalization_mode: str,
+    dataset_revision: str | None = None,
+    episodes: list[int] | None = None,
+    exclude_episodes: list[int] | None = None,
+    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:
-    """Deterministic short hash for naming the cache directory.
+    """Hash every input that changes the fitted action distribution."""
+    payload = {
+        "dataset_repo_id": dataset_repo_id,
+        "dataset_revision": dataset_revision,
+        "base_tokenizer_name": base_tokenizer_name,
+        "n_samples": n_samples,
+        "chunk_size": chunk_size,
+        "normalization_mode": normalization_mode,
+        "episodes": episodes,
+        "exclude_episodes": exclude_episodes,
+        "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]
 
-    Keys on (dataset, base tokenizer, sample count, chunk size) so any
-    of those changing re-runs the fit. ``chunk_size`` matters because
-    the tokenizer is fit on chunks of that length.
-    """
-    h = hashlib.sha256()
-    h.update(dataset_repo_id.encode("utf-8"))
-    h.update(b"\0")
-    h.update(base_tokenizer_name.encode("utf-8"))
-    h.update(b"\0")
-    h.update(str(n_samples).encode("utf-8"))
-    h.update(b"\0")
-    h.update(str(chunk_size).encode("utf-8"))
-    return h.hexdigest()[:16]
+
+def _select_episode_indices(
+    available_episodes: list[int],
+    episodes: list[int] | None,
+    exclude_episodes: list[int] | None,
+) -> list[int]:
+    allowed = set(episodes) if episodes is not None else set(available_episodes)
+    excluded = set(exclude_episodes or [])
+    return [episode for episode in available_episodes if episode in allowed and episode not in excluded]
+
+
+def _apply_relative_actions(
+    actions: np.ndarray,
+    states: np.ndarray,
+    relative_action_mask: list[bool] | None,
+) -> np.ndarray:
+    """Match RelativeActionsProcessorStep before tokenizer fitting."""
+    action_dim = actions.shape[-1]
+    mask = list(relative_action_mask) if relative_action_mask is not None else [True] * action_dim
+    if len(mask) < action_dim:
+        mask.extend([True] * (action_dim - len(mask)))
+    mask_array = np.asarray(mask[:action_dim], dtype=np.float32)
+    relative = actions.copy()
+    relative -= states[:, None, :action_dim] * mask_array
+    return relative
+
+
+def _normalize_actions(
+    actions: np.ndarray,
+    normalization_mode: str,
+    action_stats: dict | None = None,
+) -> np.ndarray:
+    """Match the action normalization applied by the training preprocessor."""
+    mode = getattr(normalization_mode, "value", normalization_mode).upper()
+    flat = actions.reshape(-1, actions.shape[-1])
+    stats = action_stats or {}
+
+    def stat(name: str, fallback) -> np.ndarray:
+        value = stats.get(name)
+        if value is None:
+            value = fallback()
+        if hasattr(value, "detach"):
+            value = value.detach().cpu().numpy()
+        return np.asarray(value, dtype=np.float32)
+
+    if mode == "IDENTITY":
+        return actions
+    if mode == "MEAN_STD":
+        mean = stat("mean", lambda: flat.mean(axis=0))
+        std = stat("std", lambda: flat.std(axis=0))
+        return ((actions - mean) / np.where(std == 0, 1e-8, std)).astype(np.float32)
+    if mode in {"QUANTILES", "QUANTILE10"}:
+        low_name, high_name, low_q, high_q = (
+            ("q01", "q99", 0.01, 0.99) if mode == "QUANTILES" else ("q10", "q90", 0.10, 0.90)
+        )
+        low = stat(low_name, lambda: np.quantile(flat, low_q, axis=0))
+        high = stat(high_name, lambda: np.quantile(flat, high_q, axis=0))
+    elif mode == "MIN_MAX":
+        low = stat("min", lambda: flat.min(axis=0))
+        high = stat("max", lambda: flat.max(axis=0))
+    else:
+        raise ValueError(f"Unsupported FAST tokenizer normalization mode: {mode}")
+
+    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(
@@ -69,6 +220,17 @@ def fit_fast_tokenizer(
     n_samples: int = 1024,
     chunk_size: int = 50,
     seed: int = 42,
+    dataset_root: str | Path | None = None,
+    dataset_revision: str | None = None,
+    episodes: list[int] | None = None,
+    exclude_episodes: list[int] | None = None,
+    normalization_mode: str = "QUANTILES",
+    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.
 
@@ -100,7 +262,23 @@ def fit_fast_tokenizer(
         FileNotFoundError: If the dataset can't be loaded.
     """
     cache_dir = Path(cache_dir)
-    sig = _dataset_signature(dataset_repo_id, base_tokenizer_name, n_samples, chunk_size)
+    normalization_mode = getattr(normalization_mode, "value", normalization_mode).upper()
+    sig = _dataset_signature(
+        dataset_repo_id,
+        base_tokenizer_name,
+        n_samples,
+        chunk_size,
+        normalization_mode,
+        dataset_revision,
+        episodes,
+        exclude_episodes,
+        action_stats,
+        use_relative_actions,
+        relative_action_mask,
+        validation_samples,
+        max_reconstruction_rmse,
+        max_dim_rmse,
+    )
     out_dir = cache_dir / sig
 
     if out_dir.exists() and (out_dir / _CACHE_SENTINEL).exists():
@@ -113,8 +291,8 @@ def fit_fast_tokenizer(
         )
         return str(out_dir)
 
-    # Each node fits its node-local cache once; its other local ranks wait.
-    is_leader = _is_local_leader()
+    # One global rank populates the shared cache; every other rank waits for the atomic publish.
+    is_leader = _is_global_leader()
     if not is_leader:
         timeout_s = 1800.0  # 30 min — covers ~1024-sample fits on cold caches
         start = time.monotonic()
@@ -138,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] = []
@@ -147,15 +323,23 @@ def fit_fast_tokenizer(
     # Read v3 parquet shards directly to avoid split lookup failures and repeated metadata parsing.
     import pyarrow as _pa  # noqa: PLC0415
     import pyarrow.parquet as _pq  # noqa: PLC0415
-    from huggingface_hub import snapshot_download  # noqa: PLC0415
 
-    snap = Path(snapshot_download(repo_id=dataset_repo_id, repo_type="dataset"))
+    if dataset_root is not None:
+        snap = Path(dataset_root)
+    else:
+        from huggingface_hub import snapshot_download  # noqa: PLC0415
+
+        snap = Path(
+            snapshot_download(repo_id=dataset_repo_id, repo_type="dataset", revision=dataset_revision)
+        )
     data_files = sorted((snap / "data").glob("chunk-*/file-*.parquet"))
     if not data_files:
         raise RuntimeError(f"FAST fit: no ``data/chunk-*/file-*.parquet`` shards found under {snap!s}.")
 
-    # Load only episode indices and fixed-width actions across all shards.
-    tables = [_pq.read_table(f, columns=["episode_index", "action"]) for f in data_files]
+    columns = ["episode_index", "action"]
+    if use_relative_actions:
+        columns.append("observation.state")
+    tables = [_pq.read_table(f, columns=columns) for f in data_files]
     table = _pa.concat_tables(tables)
     eps = table["episode_index"].to_numpy()
     acts_col = table["action"]
@@ -167,6 +351,16 @@ def fit_fast_tokenizer(
         acts = np.asarray(acts_col.to_pylist(), dtype=np.float32)
     if acts.ndim != 2:
         raise RuntimeError(f"FAST fit: expected ``action`` rows to be 1-D vectors; got shape {acts.shape}.")
+    states = None
+    if use_relative_actions:
+        try:
+            states = np.stack(table["observation.state"].to_numpy(zero_copy_only=False)).astype(np.float32)
+        except Exception:  # noqa: BLE001
+            states = np.asarray(table["observation.state"].to_pylist(), dtype=np.float32)
+        if states.ndim != 2:
+            raise RuntimeError(
+                f"FAST fit: expected ``observation.state`` rows to be 1-D vectors; got {states.shape}."
+            )
 
     # Sort once because episode order is only guaranteed within each shard.
     order = np.argsort(eps, kind="stable")
@@ -181,14 +375,20 @@ def fit_fast_tokenizer(
     # ``acts`` is in original (un-sorted-by-episode) row order; reorder
     # so per-episode slices are contiguous.
     acts = acts[order]
+    if states is not None:
+        states = states[order]
 
-    samples_per_episode = max(1, n_samples // max(num_episodes, 1))
+    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.")
+    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
-    ep_indices = list(ep_to_slice.keys())
+    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]
@@ -198,8 +398,10 @@ def fit_fast_tokenizer(
         starts = rng.integers(0, ep_actions.shape[0] - chunk_size + 1, size=samples_per_episode)
         for s in starts:
             actions_buf.append(ep_actions[int(s) : int(s) + chunk_size])
+            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
 
@@ -213,6 +415,8 @@ def fit_fast_tokenizer(
         )
 
     actions = np.stack(actions_buf, axis=0).astype(np.float32)  # (N, H, D)
+    if states is not None:
+        actions = _apply_relative_actions(actions, np.stack(states_buf), relative_action_mask)
     logger.info(
         "FAST fit: collected %d chunks of shape %s from %d episodes",
         actions.shape[0],
@@ -220,14 +424,9 @@ def fit_fast_tokenizer(
         eps_visited,
     )
 
-    # Match training-time quantile normalization so FAST sees the same bounded action space.
-    flat = actions.reshape(-1, actions.shape[-1])
-    q01 = np.quantile(flat, 0.01, axis=0)
-    q99 = np.quantile(flat, 0.99, axis=0)
-    span = np.where((q99 - q01) > 1e-6, q99 - q01, 1.0)
-    actions = np.clip((actions - q01) / span * 2.0 - 1.0, -1.0, 1.0).astype(np.float32)
+    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()`` "
@@ -235,22 +434,79 @@ def fit_fast_tokenizer(
             "to the current ``physical-intelligence/fast`` revision."
         )
 
-    fitted = base.fit(actions)
-    out_dir.mkdir(parents=True, exist_ok=True)
-    fitted.save_pretrained(str(out_dir))
+    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)
     logger.info("FAST fit: saved fitted tokenizer to %s", out_dir)
     return str(out_dir)
 
 
-def resolve_fast_tokenizer(config: Any, dataset_repo_id: str | None) -> str:
+def resolve_fast_tokenizer(
+    config: Any,
+    dataset_repo_id: str | None,
+    dataset_root: str | Path | None = None,
+    dataset_stats: dict | None = None,
+    dataset_revision: str | None = None,
+    episodes: list[int] | None = None,
+    exclude_episodes: list[int] | None = None,
+) -> str:
     """Return the configured tokenizer, fitting a cached dataset-specific one when requested."""
     if not getattr(config, "auto_fit_fast_tokenizer", False) or dataset_repo_id is None:
         return config.action_tokenizer_name
 
+    relative_action_mask = None
+    if getattr(config, "use_relative_actions", False):
+        action_names = getattr(config, "action_feature_names", None)
+        exclude_tokens = [
+            str(name).lower() for name in getattr(config, "relative_exclude_joints", []) if name
+        ]
+        if action_names is not None and exclude_tokens:
+            relative_action_mask = [
+                not any(token == str(name).lower() or token in str(name).lower() for token in exclude_tokens)
+                for name in action_names
+            ]
+
     return fit_fast_tokenizer(
         dataset_repo_id=dataset_repo_id,
         cache_dir=Path(config.fast_tokenizer_cache_dir).expanduser(),
         base_tokenizer_name=config.action_tokenizer_name,
         n_samples=config.fast_tokenizer_fit_samples,
         chunk_size=config.chunk_size,
+        dataset_root=dataset_root,
+        dataset_revision=dataset_revision,
+        episodes=episodes,
+        exclude_episodes=exclude_episodes,
+        normalization_mode=config.normalization_mapping.get("ACTION", "QUANTILES"),
+        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/policies/pi052/inference/pi052_adapter.py

@@ -41,22 +41,53 @@ class PI052PolicyAdapter(BaseLanguageAdapter):
 
         subtask = state.language_context.get("subtask") or state.task or ""
         # Match the training prompt by conditioning on both subtask and discretized state.
-        content = subtask
+        state_str = None
         obs_state = observation.get(OBS_STATE)
         if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0:
             from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str  # noqa: PLC0415
 
             state_row = obs_state[0] if obs_state.ndim > 1 else obs_state
-            content = f"{subtask}, State: {discretize_state_str(state_row)};"
+            state_str = discretize_state_str(state_row)
 
-        text_batch = _build_text_batch(
-            self.policy,
-            [{"role": "user", "content": content}],
-            add_generation_prompt=False,
-        )
         batch = dict(observation)
-        batch[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"]
-        batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"]
+        if getattr(self.policy.config, "joint_subtask_conditioning", False):
+            # Joint sequences keep the task turn (with state) and render the
+            # subtask as a causal assistant turn, exactly as trained.
+            from transformers import AutoTokenizer  # noqa: PLC0415
+
+            from lerobot.policies.pi052.text_processor_pi052 import (  # noqa: PLC0415
+                encode_prompt_with_targets,
+                register_paligemma_loc_tokens,
+            )
+            from lerobot.utils.constants import OBS_LANGUAGE_CAUSAL_MARKS  # noqa: PLC0415
+
+            task = state.task or ""
+            task_content = task if state_str is None else f"{task}, State: {state_str};"
+            tok_name = getattr(self.policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
+            tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens)
+            ids, attn, marks = encode_prompt_with_targets(
+                tokenizer,
+                [
+                    {"role": "user", "content": task_content},
+                    {"role": "assistant", "content": subtask},
+                ],
+                target_indices=[1],
+            )
+            device = getattr(self.policy.config, "device", None)
+            if device is not None:
+                ids, attn, marks = ids.to(device), attn.to(device), marks.to(device)
+            batch[OBS_LANGUAGE_TOKENS] = ids
+            batch[OBS_LANGUAGE_ATTENTION_MASK] = attn
+            batch[OBS_LANGUAGE_CAUSAL_MARKS] = marks
+        else:
+            content = subtask if state_str is None else f"{subtask}, State: {state_str};"
+            text_batch = _build_text_batch(
+                self.policy,
+                [{"role": "user", "content": content}],
+                add_generation_prompt=False,
+            )
+            batch[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"]
+            batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"]
         return self.policy.predict_action_chunk(batch)
 
     def generate_text(
@@ -67,6 +98,21 @@ class PI052PolicyAdapter(BaseLanguageAdapter):
         user_text: str | None = None,
     ) -> str:
         messages = self.build_messages(kind, state, user_text=user_text)
+        if kind == "subtask" and getattr(self.policy.config, "joint_subtask_conditioning", False):
+            # Joint samples carry state on the task turn, so the subtask must be
+            # generated from the same state-bearing prompt.
+            import torch  # noqa: PLC0415
+
+            from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str  # noqa: PLC0415
+            from lerobot.utils.constants import OBS_STATE  # noqa: PLC0415
+
+            obs_state = (observation or {}).get(OBS_STATE)
+            if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0:
+                state_row = obs_state[0] if obs_state.ndim > 1 else obs_state
+                for m in reversed(messages):
+                    if m.get("role") == "user":
+                        m["content"] = f"{m.get('content', '')}, State: {discretize_state_str(state_row)};"
+                        break
         return _generate_with_policy(
             self.policy,
             messages,
@@ -141,7 +187,10 @@ def _build_text_batch(
     messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in prompt_messages]
     prompt, _spans = _format_messages(messages)
     if add_generation_prompt:
-        prompt = prompt + "Assistant: "
+        # No trailing space: SentencePiece folds it into the first target token
+        # ("▁move"), so a space-suffixed prefill ends in a lone "▁" the model
+        # never saw at this position during training.
+        prompt = prompt + "Assistant:"
 
     encoded = tokenizer(prompt, return_tensors="pt")
     ids = encoded["input_ids"]

src/lerobot/policies/pi052/modeling_pi052.py

@@ -28,6 +28,7 @@ from torch.nn import functional
 from lerobot.utils.constants import (
     ACTION,
     OBS_LANGUAGE_ATTENTION_MASK,
+    OBS_LANGUAGE_CAUSAL_MARKS,
     OBS_LANGUAGE_TOKENS,
     OBS_STATE,
 )
@@ -143,9 +144,42 @@ class PI05Pytorch(PI05PytorchBase):  # see openpi `PI0Pytorch`
         suffix_out = suffix_out.to(dtype=torch.float32)
         return self.action_out_proj(suffix_out)
 
+    def sample_actions(
+        self,
+        images,
+        img_masks,
+        tokens,
+        masks,
+        noise=None,
+        num_steps=None,
+        lang_causal_marks=None,
+        **kwargs,
+    ) -> Tensor:
+        """Sample actions, optionally marking trailing language positions causal.
 
-# FAST tokens occupy the high vocabulary range and must be masked during text generation.
-_FAST_ACTION_VOCAB_SIZE = 2048
+        ``lang_causal_marks`` (B, L_lang bool) flags generated-subtask tokens so
+        joint-sequence checkpoints see the same causal prefix layout at
+        inference as during training (``_mark_target_span_causal``).
+        """
+        self._lang_causal_marks = lang_causal_marks
+        try:
+            return super().sample_actions(
+                images, img_masks, tokens, masks, noise=noise, num_steps=num_steps, **kwargs
+            )
+        finally:
+            self._lang_causal_marks = None
+
+    def embed_prefix(self, images, img_masks, tokens, masks):
+        prefix_embs, prefix_pad, prefix_att = super().embed_prefix(images, img_masks, tokens, masks)
+        marks = getattr(self, "_lang_causal_marks", None)
+        if marks is not None:
+            prefix_att = _apply_causal_language_marks(prefix_att, marks.to(prefix_att.device))
+        return prefix_embs, prefix_pad, prefix_att
+
+
+# The universal `physical-intelligence/fast` tokenizer (and dataset refits of it)
+# uses 1024 BPE codes; text generation must mask any that map below the <loc> range.
+_FAST_ACTION_VOCAB_SIZE = 1024
 
 
 _HF_KERNELS_ENABLED = False
@@ -331,6 +365,17 @@ def _mark_target_span_causal(
     return att
 
 
+def _apply_causal_language_marks(prefix_att_masks: Tensor, marks: Tensor) -> Tensor:
+    """OR per-token causal marks into the trailing language segment of a prefix."""
+    att = prefix_att_masks.clone()
+    n = min(marks.shape[1], att.shape[1])
+    if n <= 0:
+        return att
+    seg = att[:, -n:].bool()
+    att[:, -n:] = (seg | marks[:, -n:].bool()).to(att.dtype)
+    return att
+
+
 def _fast_lin_ce(
     hidden: Tensor,
     lm_head_weight: Tensor,
@@ -362,10 +407,27 @@ def _fast_lin_ce(
                 for sample_hidden, sample_labels in zip(shift_hidden, shift_targets, strict=True)
             ]
         )
-    batch_size, target_length, hidden_size = shift_hidden.shape
-    flat_hidden = shift_hidden.reshape(batch_size * target_length, hidden_size).to(lm_head_weight.dtype)
-    flat_labels = shift_targets.reshape(batch_size * target_length)
-    return _lin_ce_flat(flat_hidden, lm_head_weight, flat_labels, compiled=compiled)
+
+    valid_counts = shift_valid.sum(dim=1)
+    active_samples = valid_counts > 0
+    if not bool(active_samples.any().item()):
+        return shift_hidden.sum() * 0.0
+
+    weighted_losses = []
+    active_count = active_samples.sum()
+    for token_count in torch.unique(valid_counts[active_samples]).tolist():
+        group = active_samples & valid_counts.eq(token_count)
+        group_size = group.sum()
+        group_hidden = shift_hidden[group].reshape(-1, shift_hidden.shape[-1]).to(lm_head_weight.dtype)
+        group_labels = shift_targets[group].reshape(-1)
+        group_loss = _lin_ce_flat(
+            group_hidden,
+            lm_head_weight,
+            group_labels,
+            compiled=compiled,
+        )
+        weighted_losses.append(group_loss * group_size)
+    return torch.stack(weighted_losses).sum() / active_count
 
 
 # ----------------------------------------------------------------------
@@ -957,7 +1019,19 @@ class PI052Policy(PI05Policy):
             action_mask = batch.get(ACTION_TOKEN_MASK)
             action_code_mask = batch.get(ACTION_CODE_TOKEN_MASK)
             if action_tokens is None or action_mask is None or action_code_mask is None:
-                run_fast = False
+                missing = [
+                    key
+                    for key, value in (
+                        (ACTION_TOKENS, action_tokens),
+                        (ACTION_TOKEN_MASK, action_mask),
+                        (ACTION_CODE_TOKEN_MASK, action_code_mask),
+                    )
+                    if value is None
+                ]
+                raise ValueError(
+                    "PI052 FAST action loss is enabled, but the preprocessor did not produce "
+                    f"required batch keys: {missing}."
+                )
 
         # Flow uses one fused prefix/suffix pass; text-only batches skip the suffix.
         if run_flow:
@@ -1593,8 +1667,12 @@ class PI052Policy(PI05Policy):
         return decoded
 
     def _prepare_action_batch(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
-        from .inference.pi052_adapter import _build_text_batch  # noqa: PLC0415
-        from .text_processor_pi052 import discretize_state_str  # noqa: PLC0415
+        from .inference.pi052_adapter import _build_text_batch, _get_loc_tokenizer  # noqa: PLC0415
+        from .text_processor_pi052 import (  # noqa: PLC0415
+            discretize_state_str,
+            encode_prompt_with_targets,
+            register_paligemma_loc_tokens,
+        )
 
         n = self._batch_size_from_observation(batch)
         self._ensure_subtask_state(n)
@@ -1602,6 +1680,14 @@ class PI052Policy(PI05Policy):
         # Mirror training by appending the already normalized state to low-level prompts.
         state_all = batch.get(OBS_STATE)
 
+        joint = bool(getattr(self.config, "joint_subtask_conditioning", False))
+        joint_tokenizer = None
+        if joint:
+            from transformers import AutoTokenizer  # noqa: PLC0415
+
+            tok_name = getattr(self.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
+            joint_tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens)
+
         # Hold subtasks for the configured interval to match training and avoid rapid replanning.
         replan = int(getattr(self.config, "subtask_replan_steps", 0) or 0)
         hold_chunks = max(1, round(replan / self.config.n_action_steps)) if replan > 0 else 1
@@ -1609,7 +1695,7 @@ class PI052Policy(PI05Policy):
         self._subtask_chunk_counter += 1
 
         # Generate and batch one independently conditioned subtask per environment.
-        rows: list[tuple[Tensor, Tensor | None]] = []
+        rows: list[tuple[Tensor, Tensor | None, Tensor | None]] = []
         tokenizer = None
         for i in range(n):
             if regenerate or not self.last_subtasks[i]:
@@ -1619,32 +1705,62 @@ class PI052Policy(PI05Policy):
                 # Hold the previously generated subtask; only the state in the
                 # prompt below is refreshed to the current observation.
                 subtask = self.last_subtasks[i]
-            content = subtask
-            if torch.is_tensor(state_all):
-                content = f"{subtask}, State: {discretize_state_str(state_all[i])};"
-            text_batch = _build_text_batch(
-                self,
-                [{"role": "user", "content": content}],
-                add_generation_prompt=False,
-            )
-            rows.append((text_batch["lang_tokens"], text_batch["lang_masks"]))
-            tokenizer = text_batch["tokenizer"]
+            state_str = discretize_state_str(state_all[i]) if torch.is_tensor(state_all) else None
+            if joint:
+                # Joint sequences keep the task turn (with state) and render the
+                # subtask as a causal assistant turn, exactly as trained.
+                task_content = tasks[i]
+                if state_str is not None:
+                    task_content = f"{task_content}, State: {state_str};"
+                ids, attn, marks = encode_prompt_with_targets(
+                    joint_tokenizer,
+                    [
+                        {"role": "user", "content": task_content},
+                        {"role": "assistant", "content": subtask},
+                    ],
+                    target_indices=[1],
+                )
+                device = getattr(self.config, "device", None)
+                if device is not None:
+                    ids, attn, marks = ids.to(device), attn.to(device), marks.to(device)
+                rows.append((ids, attn, marks))
+                tokenizer = joint_tokenizer
+            else:
+                content = subtask if state_str is None else f"{subtask}, State: {state_str};"
+                text_batch = _build_text_batch(
+                    self,
+                    [{"role": "user", "content": content}],
+                    add_generation_prompt=False,
+                )
+                rows.append((text_batch["lang_tokens"], text_batch["lang_masks"], None))
+                tokenizer = text_batch["tokenizer"]
 
-        tokens, masks = self._stack_token_rows(rows, tokenizer)
+        tokens, masks, marks = self._stack_token_rows(rows, tokenizer)
 
         out = dict(batch)
         out[OBS_LANGUAGE_TOKENS] = tokens
         out[OBS_LANGUAGE_ATTENTION_MASK] = masks
+        if marks is not None:
+            out[OBS_LANGUAGE_CAUSAL_MARKS] = marks
         return out
 
     def _generate_low_level_subtask(self, obs_i: dict[str, Tensor], task: str, i: int) -> str:
         from .inference.pi052_adapter import _generate_with_policy  # noqa: PLC0415
+        from .text_processor_pi052 import discretize_state_str  # noqa: PLC0415
 
         msg = ""
         if task:
+            content = task
+            if getattr(self.config, "joint_subtask_conditioning", False):
+                # Joint samples carry state on the task turn, so the subtask
+                # must be generated from the same state-bearing prompt.
+                state = obs_i.get(OBS_STATE)
+                if torch.is_tensor(state) and state.numel() > 0:
+                    state_row = state[0] if state.ndim > 1 else state
+                    content = f"{task}, State: {discretize_state_str(state_row)};"
             msg = _generate_with_policy(
                 self,
-                [{"role": "user", "content": task}],
+                [{"role": "user", "content": content}],
                 observation=obs_i,
                 label=f"eval subtask gen[{i}]",
                 suppress_loc_tokens=True,
@@ -1711,21 +1827,27 @@ class PI052Policy(PI05Policy):
         return out
 
     @staticmethod
-    def _stack_token_rows(rows: list[tuple[Tensor, Tensor | None]], tokenizer: Any) -> tuple[Tensor, Tensor]:
-        """Right-pad per-env ``(1, L_i)`` token/mask rows and stack to ``(n, L)``.
+    def _stack_token_rows(
+        rows: list[tuple[Tensor, Tensor | None, Tensor | None]], tokenizer: Any
+    ) -> tuple[Tensor, Tensor, Tensor | None]:
+        """Right-pad per-env ``(1, L_i)`` token/mask/marks rows and stack to ``(n, L)``.
 
         Right-padding with a False attention mask matches the training-time
         tokenizer (``padding_side="right"``), so the action expert treats pad
-        positions as masked.
+        positions as masked. Causal marks (third element, optional) pad False.
         """
-        max_len = max(t.shape[1] for t, _ in rows)
+        max_len = max(t.shape[1] for t, _, _ in rows)
         pad_id = getattr(tokenizer, "pad_token_id", None) or 0
+        has_marks = any(m is not None for _, _, m in rows)
         tok_rows: list[Tensor] = []
         mask_rows: list[Tensor] = []
-        for tokens, masks in rows:
+        marks_rows: list[Tensor] = []
+        for tokens, masks, marks in rows:
             length = tokens.shape[1]
             if masks is None:
                 masks = torch.ones((1, length), dtype=torch.bool, device=tokens.device)
+            if has_marks and marks is None:
+                marks = torch.zeros((1, length), dtype=torch.bool, device=tokens.device)
             if length < max_len:
                 pad = max_len - length
                 tokens = torch.cat(
@@ -1736,9 +1858,17 @@ class PI052Policy(PI05Policy):
                     [masks, torch.zeros((1, pad), dtype=masks.dtype, device=masks.device)],
                     dim=1,
                 )
+                if has_marks:
+                    marks = torch.cat(
+                        [marks, torch.zeros((1, pad), dtype=marks.dtype, device=marks.device)],
+                        dim=1,
+                    )
             tok_rows.append(tokens)
             mask_rows.append(masks)
-        return torch.cat(tok_rows, dim=0), torch.cat(mask_rows, dim=0)
+            if has_marks:
+                marks_rows.append(marks)
+        stacked_marks = torch.cat(marks_rows, dim=0) if has_marks else None
+        return torch.cat(tok_rows, dim=0), torch.cat(mask_rows, dim=0), stacked_marks
 
     @staticmethod
     def _fallback_subtask_from_task(task: str) -> str:
@@ -1886,4 +2016,22 @@ class PI052Policy(PI05Policy):
         if self.config.use_flashrt_fp8_mlp and not getattr(self, "_fp8_applied", False):
             self._fp8_applied = True
             self.apply_flashrt_fp8_mlp(batch)
-        return super().predict_action_chunk(batch, **kwargs)
+        marks = batch.get(OBS_LANGUAGE_CAUSAL_MARKS)
+        if marks is None:
+            return super().predict_action_chunk(batch, **kwargs)
+        return self._predict_action_chunk_with_marks(batch, marks, **kwargs)
+
+    @torch.no_grad()
+    def _predict_action_chunk_with_marks(
+        self, batch: dict[str, Tensor], marks: Tensor, **kwargs: Unpack[ActionSelectKwargs]
+    ) -> Tensor:
+        """Base ``predict_action_chunk`` plus causal marks on the generated-subtask span."""
+        self.eval()
+        images, img_masks = self._preprocess_images(batch)
+        tokens = batch[OBS_LANGUAGE_TOKENS]
+        masks = batch[OBS_LANGUAGE_ATTENTION_MASK]
+        actions = self.model.sample_actions(
+            images, img_masks, tokens, masks, lang_causal_marks=marks, **kwargs
+        )
+        original_action_dim = self.config.output_features[ACTION].shape[0]
+        return actions[:, :, :original_action_dim]

src/lerobot/policies/pi052/processor_pi052.py

@@ -53,6 +53,10 @@ def make_pi052_pre_post_processors(
     config: PI052Config,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
     dataset_repo_id: str | None = None,
+    dataset_root: str | None = None,
+    dataset_revision: str | None = None,
+    episodes: list[int] | None = None,
+    exclude_episodes: list[int] | None = None,
 ) -> tuple[
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
@@ -62,6 +66,8 @@ def make_pi052_pre_post_processors(
     Falls through to π0.5's stock pipeline when ``recipe_path`` is unset.
     """
     if not config.recipe_path:
+        if getattr(config, "enable_fast_action_loss", False):
+            raise ValueError("PI052 FAST action loss requires recipe_path to build action supervision.")
         return make_pi05_pre_post_processors(config, dataset_stats=dataset_stats)
 
     recipe = _load_recipe(config.recipe_path)
@@ -97,10 +103,19 @@ def make_pi052_pre_post_processors(
 
         input_steps.append(
             ActionTokenizerProcessorStep(
-                action_tokenizer_name=resolve_fast_tokenizer(config, dataset_repo_id),
+                action_tokenizer_name=resolve_fast_tokenizer(
+                    config,
+                    dataset_repo_id,
+                    dataset_root,
+                    dataset_stats,
+                    dataset_revision,
+                    episodes,
+                    exclude_episodes,
+                ),
                 max_action_tokens=config.max_action_tokens,
                 fast_skip_tokens=config.fast_skip_tokens,
                 paligemma_tokenizer_name="google/paligemma-3b-pt-224",
+                allow_truncation=False,
             )
         )
 

src/lerobot/policies/pi052/text_processor_pi052.py

@@ -251,6 +251,44 @@ def _format_messages(
     return "".join(parts), spans
 
 
+def encode_prompt_with_targets(
+    tokenizer: Any, messages: list[dict[str, Any]], target_indices: list[int]
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Tokenize a flat prompt and mark the token positions of target spans.
+
+    Inference-side twin of ``PI052TextTokenizerStep._encode_messages``: same
+    serialization (role headers, target EOS) and the same offset-overlap span
+    arithmetic, but unpadded and returning a boolean target mask instead of
+    labels. Used to rebuild joint-sequence prompts whose target spans must be
+    attended causally, matching ``_mark_target_span_causal`` at train time.
+
+    Returns ``(input_ids, attention_mask, target_marks)``, each ``(1, L)``.
+    """
+    prompt, spans = _format_messages(messages, target_indices, getattr(tokenizer, "eos_token", None))
+    encoded = tokenizer(prompt, return_tensors="pt", return_offsets_mapping=True)
+    input_ids = encoded["input_ids"][0]
+    attention_mask = encoded.get("attention_mask")
+    if attention_mask is None:
+        attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+    else:
+        attention_mask = attention_mask[0].bool()
+    offsets = encoded["offset_mapping"][0]
+
+    marks = torch.zeros_like(input_ids, dtype=torch.bool)
+    for idx in target_indices:
+        if idx >= len(spans):
+            continue
+        char_start, char_end = spans[idx]
+        for token_pos in range(input_ids.shape[0]):
+            if not attention_mask[token_pos]:
+                continue
+            tok_start, tok_end = int(offsets[token_pos, 0]), int(offsets[token_pos, 1])
+            if tok_end <= char_start or tok_start >= char_end:
+                continue
+            marks[token_pos] = True
+    return input_ids.unsqueeze(0), attention_mask.unsqueeze(0), marks.unsqueeze(0)
+
+
 @dataclass
 @ProcessorStepRegistry.register(name="pi052_text_tokenizer")
 class PI052TextTokenizerStep(ProcessorStep):

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -102,6 +102,10 @@ def make_pi0_fast_pre_post_processors(
     config: PI0FastConfig,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
     dataset_repo_id: str | None = None,
+    dataset_root: str | None = None,
+    dataset_revision: str | None = None,
+    episodes: list[int] | None = None,
+    exclude_episodes: list[int] | None = None,
 ) -> tuple[
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
@@ -146,7 +150,15 @@ def make_pi0_fast_pre_post_processors(
     # continues to receive normalized state in [-1, 1] as expected.
     from ..pi052.fit_fast_tokenizer import resolve_fast_tokenizer  # noqa: PLC0415
 
-    action_tokenizer_path = resolve_fast_tokenizer(config, dataset_repo_id)
+    action_tokenizer_path = resolve_fast_tokenizer(
+        config,
+        dataset_repo_id,
+        dataset_root,
+        dataset_stats,
+        dataset_revision,
+        episodes,
+        exclude_episodes,
+    )
 
     input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one

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
@@ -350,6 +351,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
     max_action_tokens: int = 256
     fast_skip_tokens: int = 128
     paligemma_tokenizer_name: str = "google/paligemma-3b-pt-224"
+    allow_truncation: bool = True
     # Internal tokenizer instance (not part of the config)
     action_tokenizer: Any = field(default=None, init=False, repr=False)
     _paligemma_tokenizer: Any = field(default=None, init=False, repr=False)
@@ -502,6 +504,11 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
 
             # Truncate or pad to max_action_tokens
             if len(tokens) > self.max_action_tokens:
+                if not self.allow_truncation:
+                    raise ValueError(
+                        f"FAST action sequence has {len(tokens)} tokens, exceeding "
+                        f"max_action_tokens={self.max_action_tokens}."
+                    )
                 logging.warning(
                     f"Token length ({len(tokens)}) exceeds max length ({self.max_action_tokens}), truncating. "
                     "Consider increasing the `max_action_tokens` in your model config if this happens frequently."
@@ -567,6 +574,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
             "max_action_tokens": self.max_action_tokens,
             "fast_skip_tokens": self.fast_skip_tokens,
             "paligemma_tokenizer_name": self.paligemma_tokenizer_name,
+            "allow_truncation": self.allow_truncation,
         }
 
         # Only save tokenizer_name if it was used to create the tokenizer
@@ -575,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]]:

src/lerobot/scripts/lerobot_train.py

@@ -352,6 +352,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     if cfg.policy.type in {"pi0_fast", "pi052"}:
         processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id
+        processor_kwargs["dataset_revision"] = cfg.dataset.revision
+        processor_kwargs["dataset_episodes"] = cfg.dataset.episodes
+        processor_kwargs["dataset_exclude_episodes"] = cfg.dataset.exclude_episodes
+        processor_kwargs["dataset_root"] = cfg.dataset.root
 
     if not cfg.is_reward_model_training and processor_pretrained_path is not None:
         preprocessor_overrides = {
@@ -538,6 +542,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             collate_fn=eval_collate_fn,
             prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
             persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
+            multiprocessing_context=mp_context if cfg.num_workers > 0 else None,
         )
 
     # Prepare everything with accelerator

src/lerobot/utils/constants.py

@@ -26,6 +26,7 @@ OBS_IMAGES = OBS_IMAGE + "s"
 OBS_LANGUAGE = OBS_STR + ".language"
 OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
 OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
+OBS_LANGUAGE_CAUSAL_MARKS = OBS_LANGUAGE + ".causal_marks"
 OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask"
 OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens"
 OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask"

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_fast_action_loss.py

@@ -16,14 +16,18 @@
 
 """Regression tests for PI052 FAST action-code supervision."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
+from torch import nn
 from torch.nn import functional as F  # noqa: N812
 
 pytest.importorskip("transformers")
 pytest.importorskip("liger_kernel")
 
-from lerobot.policies.pi052.modeling_pi052 import _fast_lin_ce  # noqa: E402
+from lerobot.policies.pi052.modeling_pi052 import PI052Policy, _fast_lin_ce  # noqa: E402
+from lerobot.policies.pi052.processor_pi052 import make_pi052_pre_post_processors  # noqa: E402
 
 
 def _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t):
@@ -104,3 +108,55 @@ def test_fast_ce_returns_zero_when_no_action_code_positions_are_valid():
     assert loss.item() == 0
     loss.backward()
     assert logits.grad is not None
+
+
+def test_fast_ce_averages_each_action_sample_equally():
+    torch.manual_seed(0)
+    hidden = torch.randn(2, 5, 8)
+    lm_head_weight = torch.eye(8)
+    action_tokens = torch.tensor([[1, 2, 0, 0, 0], [1, 3, 4, 5, 6]])
+    action_code_mask = torch.tensor([[False, True, False, False, False], [False, True, True, True, True]])
+
+    loss = _fast_lin_ce(
+        hidden,
+        lm_head_weight,
+        action_tokens,
+        action_code_mask,
+        predict_actions_t=None,
+        reduction="mean",
+    )
+    per_sample = _fast_lin_ce(
+        hidden,
+        lm_head_weight,
+        action_tokens,
+        action_code_mask,
+        predict_actions_t=None,
+        reduction="none",
+    )
+
+    assert torch.allclose(loss, per_sample.mean())
+
+
+def test_pi052_rejects_fast_loss_without_recipe():
+    config = SimpleNamespace(recipe_path=None, enable_fast_action_loss=True)
+
+    with pytest.raises(ValueError, match="recipe_path"):
+        make_pi052_pre_post_processors(config)
+
+
+def test_pi052_rejects_missing_fast_batch_keys():
+    policy = PI052Policy.__new__(PI052Policy)
+    nn.Module.__init__(policy)
+    policy.config = SimpleNamespace(
+        enable_fast_action_loss=True,
+        fast_action_loss_weight=1.0,
+        flow_loss_weight=0.0,
+        text_loss_weight=1.0,
+    )
+    batch = {
+        "text_labels": torch.tensor([[1, 2]]),
+        "predict_actions": torch.tensor([True]),
+    }
+
+    with pytest.raises(ValueError, match="FAST action loss is enabled"):
+        policy.forward(batch)

tests/policies/pi052/test_pi052_fit_fast_tokenizer.py

@@ -0,0 +1,122 @@
+#!/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
+import pytest
+
+from lerobot.policies.pi052.fit_fast_tokenizer import (
+    _apply_relative_actions,
+    _dataset_signature,
+    _is_global_leader,
+    _normalize_actions,
+    _select_episode_indices,
+    _validate_fast_reconstruction,
+)
+
+
+def test_fast_tokenizer_fit_uses_training_mean_std_normalization():
+    actions = np.array([[[1.0, 7.0], [3.0, 3.0]]], dtype=np.float32)
+    stats = {"mean": [2.0, 5.0], "std": [0.5, 2.0]}
+
+    normalized = _normalize_actions(actions, "MEAN_STD", stats)
+
+    np.testing.assert_allclose(normalized, [[[-2.0, 1.0], [2.0, -1.0]]])
+
+
+def test_fast_tokenizer_fit_quantiles_match_training_without_clipping():
+    actions = np.array([[[-1.0], [3.0]]], dtype=np.float32)
+    stats = {"q01": [0.0], "q99": [2.0]}
+
+    normalized = _normalize_actions(actions, "QUANTILES", stats)
+
+    np.testing.assert_allclose(normalized, [[[-2.0], [2.0]]])
+
+
+def test_fast_tokenizer_cache_signature_tracks_stats_and_episode_selection():
+    kwargs = {
+        "dataset_repo_id": "org/dataset",
+        "base_tokenizer_name": "physical-intelligence/fast",
+        "n_samples": 100,
+        "chunk_size": 20,
+        "normalization_mode": "QUANTILES",
+        "dataset_revision": "main",
+        "episodes": [1, 2, 3],
+        "exclude_episodes": [2],
+        "use_relative_actions": False,
+        "relative_action_mask": None,
+    }
+
+    first = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [1.0]})
+    changed_stats = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [2.0]})
+    changed_selection = _dataset_signature(
+        **{**kwargs, "exclude_episodes": [2, 3]},
+        action_stats={"q01": [0.0], "q99": [1.0]},
+    )
+
+    assert first != changed_stats
+    assert first != changed_selection
+
+
+def test_fast_tokenizer_uses_only_global_rank_zero(monkeypatch):
+    monkeypatch.setenv("RANK", "8")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    assert not _is_global_leader()
+
+    monkeypatch.setenv("RANK", "0")
+    assert _is_global_leader()
+
+
+def test_fast_tokenizer_episode_selection_applies_allowlist_and_exclusions():
+    selected = _select_episode_indices([0, 1, 2, 3], episodes=[1, 2, 3], exclude_episodes=[2])
+
+    assert selected == [1, 3]
+
+
+def test_fast_tokenizer_relative_actions_match_training_transform():
+    actions = np.array([[[2.0, 10.0], [3.0, 11.0]]], dtype=np.float32)
+    states = np.array([[1.0, 4.0]], dtype=np.float32)
+
+    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/policies/pi052/test_pi052_joint_sequence.py

@@ -0,0 +1,148 @@
+#!/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.
+
+"""Tests for PI052 joint-sequence (paper-style) subtask conditioning.
+
+Joint recipes train the subtask text and the action losses in one sequence,
+with the supervised subtask span attended causally. At inference the same
+layout is rebuilt around the *generated* subtask, so these tests pin:
+
+- the inference-side encoder produces the same token ids and target positions
+  as the training-time tokenizer step for the same messages;
+- OR-ing causal marks into a prefix reproduces the training-time attention
+  pattern (prompt cannot see the subtask; subtask is causal over itself);
+- the joint recipe file stays a valid message recipe;
+- the FAST id mapping with the default ``fast_skip_tokens`` stays clear of
+  PaliGemma's ``<loc>`` range so VQA and FAST supervision never collide.
+"""
+
+from pathlib import Path
+
+import torch
+
+from lerobot.configs.recipe import TrainingRecipe
+from lerobot.policies.pi052.text_processor_pi052 import (
+    PI052TextTokenizerStep,
+    encode_prompt_with_targets,
+)
+
+
+class _CharTokenizer:
+    """Char-level stub: 1 char = 1 token, so offsets are trivially aligned."""
+
+    pad_token_id = 0
+    eos_token = "\x1f"  # unit separator — a 1-char "EOS" for testing
+
+    def __call__(self, text, max_length=None, padding=None, return_tensors=None, **kwargs):
+        limit = max_length if max_length is not None else len(text)
+        ids = [ord(c) % 251 + 1 for c in text[:limit]]
+        offsets = [(i, i + 1) for i in range(len(ids))]
+        attention = [1] * len(ids)
+        if padding == "max_length" and max_length is not None and len(ids) < max_length:
+            pad = max_length - len(ids)
+            ids += [self.pad_token_id] * pad
+            offsets += [(0, 0)] * pad
+            attention += [0] * pad
+        return {
+            "input_ids": torch.tensor([ids], dtype=torch.long),
+            "attention_mask": torch.tensor([attention], dtype=torch.long),
+            "offset_mapping": torch.tensor([offsets], dtype=torch.long),
+        }
+
+
+_MESSAGES = [
+    {"role": "user", "content": "fold the towel"},
+    {"role": "assistant", "content": "grab the near corner"},
+]
+
+
+def test_encode_prompt_with_targets_matches_training_labels():
+    tokenizer = _CharTokenizer()
+
+    step = PI052TextTokenizerStep(max_length=120)
+    step._tokenizer = tokenizer
+    train_ids, train_attn, labels, predict_actions, _prompt = step._encode_messages(
+        tokenizer,
+        [dict(m) for m in _MESSAGES],
+        message_streams=["low_level", "low_level"],
+        target_indices=[1],
+        complementary={},
+    )
+    assert bool(predict_actions)
+
+    ids, attn, marks = encode_prompt_with_targets(tokenizer, [dict(m) for m in _MESSAGES], [1])
+
+    n = int(attn.sum())
+    assert n == int(train_attn.sum())
+    assert torch.equal(ids[0, :n], train_ids[:n])
+    # Causal marks at inference must cover exactly the supervised label span.
+    assert torch.equal(marks[0, :n], labels[:n] != -100)
+    assert marks.any(), "the assistant target span must be marked"
+    # The user turn must stay unmarked (bidirectional prompt).
+    user_len = len("User: fold the towel\n")
+    assert not marks[0, :user_len].any()
+
+
+def test_apply_causal_language_marks_reproduces_training_mask():
+    from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks
+    from lerobot.policies.pi052.modeling_pi052 import _apply_causal_language_marks
+
+    n_img, n_lang = 4, 8
+    prefix_len = n_img + n_lang
+    pad = torch.ones((1, prefix_len), dtype=torch.bool)
+    att = torch.zeros((1, prefix_len), dtype=torch.bool)
+    # Subtask span = language positions 5..7 (prefix positions 9..11).
+    marks = torch.zeros((1, n_lang), dtype=torch.bool)
+    marks[0, 5:8] = True
+
+    att_marked = _apply_causal_language_marks(att, marks)
+    att_2d = make_att_2d_masks(pad, att_marked)[0]
+
+    subtask = [n_img + 5, n_img + 6, n_img + 7]
+    # Prompt and images never see the subtask.
+    for q in range(n_img + 5):
+        for k in subtask:
+            assert not att_2d[q, k], f"prompt position {q} must not attend subtask position {k}"
+    # Subtask tokens see the full prompt and earlier subtask tokens only.
+    for qi, q in enumerate(subtask):
+        for k in range(n_img + 5):
+            assert att_2d[q, k]
+        for ki, k in enumerate(subtask):
+            assert bool(att_2d[q, k]) == (ki <= qi)
+
+
+def test_joint_recipe_is_a_valid_message_recipe():
+    recipe_path = Path(__file__).parents[3] / "src" / "lerobot" / "configs" / "recipes" / "subtask_joint.yaml"
+    recipe = TrainingRecipe.from_yaml(recipe_path)
+    assert recipe.messages is not None and len(recipe.messages) == 2
+    assert all(turn.stream == "low_level" for turn in recipe.messages)
+    assert not recipe.messages[0].target
+    assert recipe.messages[1].target
+    assert recipe.messages[1].if_present == "subtask"
+
+
+def test_default_fast_mapping_clears_loc_and_seg_ranges():
+    from lerobot.policies.pi052.configuration_pi052 import PI052Config
+    from lerobot.policies.pi052.modeling_pi052 import _FAST_ACTION_VOCAB_SIZE
+
+    skip = PI052Config.__dataclass_fields__["fast_skip_tokens"].default
+    assert skip == 1152
+
+    paligemma_vocab = 257152
+    fast_ids = paligemma_vocab - 1 - skip - torch.arange(_FAST_ACTION_VOCAB_SIZE)
+    # Below the <loc> range [256000, 257024) and the <seg> range [257024, 257152).
+    assert int(fast_ids.max()) < 256000
+    assert int(fast_ids.min()) >= 0

tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py

@@ -47,6 +47,14 @@ def test_pi0_fast_resolves_dataset_specific_tokenizer(monkeypatch, tmp_path):
         "base_tokenizer_name": "base-tokenizer",
         "n_samples": 17,
         "chunk_size": 12,
+        "dataset_root": None,
+        "dataset_revision": None,
+        "episodes": None,
+        "exclude_episodes": None,
+        "normalization_mode": config.normalization_mapping["ACTION"],
+        "action_stats": None,
+        "use_relative_actions": False,
+        "relative_action_mask": None,
     }
 
 
@@ -62,13 +70,13 @@ def test_fast_fit_failure_is_not_silently_replaced(monkeypatch, tmp_path):
         fit_module.resolve_fast_tokenizer(config, "user/dataset")
 
 
-def test_each_node_uses_its_local_rank_zero_as_fit_leader(monkeypatch):
+def test_only_global_rank_zero_fits_shared_tokenizer(monkeypatch):
     monkeypatch.setenv("RANK", "8")
     monkeypatch.setenv("LOCAL_RANK", "0")
-    assert fit_module._is_local_leader()
+    assert not fit_module._is_global_leader()
 
-    monkeypatch.setenv("LOCAL_RANK", "1")
-    assert not fit_module._is_local_leader()
+    monkeypatch.setenv("RANK", "0")
+    assert fit_module._is_global_leader()
 
 
 def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
@@ -78,7 +86,7 @@ def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
     monkeypatch.setattr(
         fit_module,
         "resolve_fast_tokenizer",
-        lambda config, dataset_repo_id: "/cache/fitted-tokenizer",
+        lambda config, dataset_repo_id, *args: "/cache/fitted-tokenizer",
     )
 
     def fake_from_pretrained(cls, *args, **kwargs):

tests/processor/test_tokenizer_processor.py

@@ -94,6 +94,7 @@ def test_action_tokenizer_config_preserves_token_mapping():
     processor.max_action_tokens = 384
     processor.fast_skip_tokens = 64
     processor.paligemma_tokenizer_name = "custom/paligemma"
+    processor.allow_truncation = False
     processor.action_tokenizer_name = "custom/fast"
     processor.action_tokenizer_input_object = None
 
@@ -102,10 +103,31 @@ def test_action_tokenizer_config_preserves_token_mapping():
         "max_action_tokens": 384,
         "fast_skip_tokens": 64,
         "paligemma_tokenizer_name": "custom/paligemma",
+        "allow_truncation": False,
         "action_tokenizer_name": "custom/fast",
     }
 
 
+def test_action_tokenizer_can_reject_truncated_sequences():
+    processor = object.__new__(ActionTokenizerProcessorStep)
+    processor.max_action_tokens = 4
+    processor.fast_skip_tokens = 128
+    processor.allow_truncation = False
+    processor.action_tokenizer = lambda _actions: [1, 2, 3]
+    processor._paligemma_tokenizer = type(
+        "Tokenizer",
+        (),
+        {
+            "vocab_size": 1000,
+            "bos_token_id": 2,
+            "encode": lambda _self, text, **_kwargs: [10, 11] if text == "Action: " else [12, 1],
+        },
+    )()
+
+    with pytest.raises(ValueError, match="max_action_tokens=4"):
+        processor._tokenize_action(torch.zeros(1, 2, 1))
+
+
 @pytest.fixture
 def mock_tokenizer():
     """Provide a mock tokenizer for testing."""

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