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/multi_gpu_training.mdx

@@ -113,19 +113,6 @@ accelerate launch --num_processes=2 $(which lerobot-train) \
   --policy.type=act
 ```
 
-When the desired global batch is larger than the per-GPU batch that fits in memory, use gradient
-accumulation. `steps` continues to count optimizer updates:
-
-```bash
-# 8 samples/GPU × 4 GPUs × 2 microbatches = effective global batch 64.
-accelerate launch --num_processes=4 $(which lerobot-train) \
-  --batch_size=8 \
-  --gradient_accumulation_steps=2 \
-  --steps=80000 \
-  --dataset.repo_id=lerobot/pusht \
-  --policy.type=act
-```
-
 ## Training Large Models with FSDP
 
 DDP replicates the full model on every GPU, so a model that doesn't fit on one GPU won't fit under

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

docs/source/robomme.mdx

@@ -44,18 +44,6 @@ docker build -f docker/Dockerfile.benchmark.robomme -t lerobot-robomme .
 
 The `docker/Dockerfile.benchmark.robomme` image overrides `gymnasium==0.29.1` and `numpy==1.26.4` after lerobot's install. Both versions are runtime-safe for lerobot's actual API usage.
 
-When evaluating a checkpoint saved on the host, run the container with the host UID. Checkpoint weights are intentionally private by default (`0600`), so the image's built-in user cannot otherwise read a bind-mounted `model.safetensors` file. Mount the Hugging Face cache at an accessible path as well:
-
-```bash
-docker run --gpus all --rm --ipc=host \
-  --user "$(id -u):$(id -g)" \
-  -e HF_HOME=/tmp/hf-cache \
-  -v "$HOME/.cache/huggingface:/tmp/hf-cache:ro" \
-  -v "$PWD/outputs:/results" \
-  lerobot-robomme \
-  lerobot-eval --policy.path=/results/<checkpoint>/pretrained_model # ...
-```
-
 ## Running Evaluation
 
 ### Default (single task, single episode)
@@ -73,7 +61,7 @@ lerobot-eval \
 
 ### Multi-task evaluation
 
-Evaluate multiple tasks in one run by comma-separating task names. Use `task_ids` to select the dataset-backed episodes evaluated for every task. For the standard 50-episode test protocol, select IDs 0 through 49 and run each selected episode once.
+Evaluate multiple tasks in one run by comma-separating task names. Use `task_ids` to control which episodes are evaluated per task. Recommended: 50 episodes per task for the test split.
 
 ```bash
 lerobot-eval \
@@ -81,22 +69,20 @@ lerobot-eval \
     --env.type=robomme \
     --env.task=PickXtimes,BinFill,StopCube,MoveCube,InsertPeg \
     --env.dataset_split=test \
-    --env.task_ids=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49] \
+    --env.task_ids=[0,1,2,3,4,5,6,7,8,9] \
     --eval.batch_size=1 \
-    --eval.n_episodes=1
+    --eval.n_episodes=50
 ```
 
 ### Key CLI options for `env.type=robomme`
 
-| Option               | Default       | Description                                      |
-| -------------------- | ------------- | ------------------------------------------------ |
-| `env.task`           | `PickXtimes`  | Any of the 16 task names above (comma-separated) |
-| `env.dataset_split`  | `test`        | `train`, `val`, or `test`                        |
-| `env.action_space`   | `joint_angle` | `joint_angle` (8-D) or `ee_pose` (7-D)           |
-| `env.episode_length` | `300`         | Max steps per episode                            |
-| `env.task_ids`       | `null`        | Dataset episode indices (null = `[0]`)           |
-
-`eval.n_episodes` repeats each selected `task_id`; it does not advance to the next dataset episode. Keep it at `1` when enumerating the 50 distinct test episodes with `env.task_ids`.
+| Option               | Default       | Description                                        |
+| -------------------- | ------------- | -------------------------------------------------- |
+| `env.task`           | `PickXtimes`  | Any of the 16 task names above (comma-separated)   |
+| `env.dataset_split`  | `test`        | `train`, `val`, or `test`                          |
+| `env.action_space`   | `joint_angle` | `joint_angle` (8-D) or `ee_pose` (7-D)             |
+| `env.episode_length` | `300`         | Max steps per episode                              |
+| `env.task_ids`       | `null`        | List of episode indices to evaluate (null = `[0]`) |
 
 ## Dataset
 
@@ -110,68 +96,22 @@ dataset = LeRobotDataset("lerobot/robomme")
 
 ### Dataset features
 
-| Feature            | Shape         | Description                      |
-| ------------------ | ------------- | -------------------------------- |
-| `image`            | (256, 256, 3) | Front camera RGB                 |
-| `wrist_image`      | (256, 256, 3) | Wrist camera RGB                 |
-| `actions`          | (8,)          | Joint angles + gripper           |
-| `state`            | (8,)          | Joint positions + gripper state  |
-| `simple_subgoal`   | str           | High-level language annotation   |
-| `grounded_subgoal` | str           | Grounded language annotation     |
-| `exec_start_idx`   | scalar        | First action-execution frame     |
-| `is_demo`          | bool          | Video-demonstration prefix frame |
-| `episode_index`    | int           | Episode ID                       |
-| `frame_index`      | int           | Frame within episode             |
+| Feature            | Shape         | Description                     |
+| ------------------ | ------------- | ------------------------------- |
+| `image`            | (256, 256, 3) | Front camera RGB                |
+| `wrist_image`      | (256, 256, 3) | Wrist camera RGB                |
+| `actions`          | (8,)          | Joint angles + gripper          |
+| `state`            | (8,)          | Joint positions + gripper state |
+| `simple_subgoal`   | str           | High-level language annotation  |
+| `grounded_subgoal` | str           | Grounded language annotation    |
+| `episode_index`    | int           | Episode ID                      |
+| `frame_index`      | int           | Frame within episode            |
 
 ### Feature key alignment (training)
 
 The env wrapper exposes `pixels/image` and `pixels/wrist_image` as observation keys. The `features_map` in `RoboMMEEnv` maps these to `observation.images.image` and `observation.images.wrist_image` for the policy. State is exposed as `agent_pos` and maps to `observation.state`.
 
-The published dataset uses the raw keys shown above. Policies that expect canonical LeRobot keys should map them at training time. For example, the pretrained SmolVLA checkpoint expects three cameras, while RoboMME provides two:
-
-```bash
-uv run lerobot-train \
-  --policy.path=lerobot/smolvla_base \
-  --policy.empty_cameras=1 \
-  --dataset.repo_id=lerobot/robomme \
-  --dataset.training_target_start_feature=exec_start_idx \
-  '--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
-```
-
-RoboMME episodes begin with video-demonstration/context frames that should remain available to a
-memory policy but should not be sampled as action-training targets. Setting
-`dataset.training_target_start_feature=exec_start_idx` starts target sampling at each episode's
-execution boundary while preserving earlier frames for temporal observation deltas. In the published
-training split this keeps 476,857 execution targets out of 768,897 total frames. One execution-target
-epoch therefore takes `ceil(476857 / effective_batch_size)` optimizer steps.
-
-For a sample-matched SmolVLA visual-memory ablation, use
-`examples/robomme/smolvla_visual_memory_ablation.sh`. `TARGET_SAMPLES` counts examples across all
-GPUs, and `NUM_PROCESSES` is included when the script converts that target into optimizer steps. For
-example, the following reproduces 5.12 million example exposures (the exposure of RoboMME's
-80,000-step, global-batch-64 memory-policy recipe) with four GPUs, two accumulated microbatches,
-and an effective global batch of 64:
-
-```bash
-TARGET_SAMPLES=5120000 \
-NUM_PROCESSES=4 \
-BATCH_SIZE=8 \
-GRADIENT_ACCUMULATION_STEPS=2 \
-VARIANT=visual-memory \
-RUN_EVAL=false \
-bash examples/robomme/smolvla_visual_memory_ablation.sh
-```
-
-This becomes 80,000 optimizer steps, or about 10.74 execution-target epochs. Run the baseline with
-the same `TARGET_SAMPLES`, effective batch size, seed, and scheduler settings to isolate the visual
-memory change. RoboMME's released baseline uses a different global batch (128), so its nominal
-80,000-step recipe is not sample-matched to its global-batch-64 memory recipe.
-
-At evaluation time the environment wrapper has already converted observations to canonical keys. Only the two camera suffixes need to be aligned with the checkpoint:
-
-```bash
-'--rename_map={"observation.images.image":"observation.images.camera1","observation.images.wrist_image":"observation.images.camera2"}'
-```
+The dataset's `image` and `wrist_image` columns already align with the policy input keys, so no renaming is needed when fine-tuning.
 
 ## Action Spaces
 

docs/source/smolvla.mdx

@@ -76,37 +76,6 @@ Fine-tuning is an art. For a complete overview of the options for finetuning, ru
 lerobot-train --help
 ```
 
-### Experimental MEM visual memory
-
-SmolVLA can optionally fuse a short history of camera frames using the
-space-time separable vision encoder from [MEM](https://arxiv.org/abs/2603.03596).
-Every fourth SigLIP layer adds causal attention across time for matching image
-patches. Historical tokens are discarded inside the vision tower, so the VLM
-receives the same number of image tokens as the baseline policy.
-
-The option is disabled by default. The following example uses six observations
-spaced one second apart for a 10 fps dataset:
-
-```bash
-lerobot-train \
-  --policy.path=lerobot/smolvla_base \
-  --policy.use_visual_memory=true \
-  --policy.visual_memory_frames=6 \
-  --policy.visual_memory_stride=10 \
-  --policy.visual_memory_temporal_attention_every=4 \
-  --policy.freeze_vision_encoder=false \
-  --policy.train_expert_only=false \
-  --dataset.repo_id=${HF_USER}/mydataset \
-  --output_dir=outputs/train/my_smolvla_mem
-```
-
-`visual_memory_stride` is measured in dataset or environment steps. Keep all
-other settings and the random seed fixed when comparing against a baseline.
-Because the public SmolVLA checkpoint was not pretrained with this temporal
-attention pattern, this is a post-training-only MEM ablation; the MEM paper
-reports that memory-aware pretraining performs better than introducing visual
-memory only during task-specific fine-tuning.
-
 <p align="center">
   <img
     src="https://cdn-uploads.huggingface.co/production/uploads/640e21ef3c82bd463ee5a76d/S-3vvVCulChREwHDkquoc.gif"

examples/robomme/smolvla_visual_memory_ablation.sh

@@ -1,131 +0,0 @@
-#!/usr/bin/env bash
-# Matched SmolVLA ablation for the 10 fps lerobot/robomme dataset.
-# Run from the LeRobot repository root on a CUDA machine.
-
-set -euo pipefail
-
-BATCH_SIZE="${BATCH_SIZE:-4}"
-NUM_PROCESSES="${NUM_PROCESSES:-1}"
-GRADIENT_ACCUMULATION_STEPS="${GRADIENT_ACCUMULATION_STEPS:-1}"
-# The published training split has 476,857 execution frames. By default, train on one
-# execution-frame epoch; set STEPS explicitly to use a different optimizer-step budget.
-TARGET_SAMPLES="${TARGET_SAMPLES:-476857}"
-EFFECTIVE_BATCH_SIZE=$((BATCH_SIZE * NUM_PROCESSES * GRADIENT_ACCUMULATION_STEPS))
-STEPS="${STEPS:-$(((TARGET_SAMPLES + EFFECTIVE_BATCH_SIZE - 1) / EFFECTIVE_BATCH_SIZE))}"
-SCHEDULER_WARMUP_STEPS="${SCHEDULER_WARMUP_STEPS:-$(((STEPS + 29) / 30))}"
-SCHEDULER_DECAY_STEPS="${SCHEDULER_DECAY_STEPS:-${STEPS}}"
-SEED="${SEED:-1000}"
-OUTPUT_ROOT="${OUTPUT_ROOT:-outputs/robomme-smolvla-mem-ablation}"
-WANDB_ENABLE="${WANDB_ENABLE:-false}"
-RUN_TRAIN="${RUN_TRAIN:-true}"
-RUN_EVAL="${RUN_EVAL:-true}"
-VARIANT="${VARIANT:-both}"
-TASKS="BinFill,PickXtimes,SwingXtimes,StopCube,VideoUnmask,VideoUnmaskSwap,ButtonUnmask,ButtonUnmaskSwap,PickHighlight,VideoRepick,VideoPlaceButton,VideoPlaceOrder,MoveCube,InsertPeg,PatternLock,RouteStick"
-TASK_IDS="[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49]"
-
-case "${VARIANT}" in
-  baseline) VARIANTS=(baseline) ;;
-  visual-memory) VARIANTS=(visual-memory) ;;
-  both) VARIANTS=(baseline visual-memory) ;;
-  *) echo "VARIANT must be baseline, visual-memory, or both" >&2; exit 2 ;;
-esac
-
-TRAIN_COMMAND=()
-if [[ "${RUN_TRAIN}" == "true" ]]; then
-  if ((NUM_PROCESSES > 1)); then
-    TRAIN_ENTRYPOINT="$(uv run which lerobot-train)"
-    TRAIN_COMMAND=(
-      uv run accelerate launch
-      --multi_gpu
-      --num_processes="${NUM_PROCESSES}"
-      --mixed_precision=bf16
-      "${TRAIN_ENTRYPOINT}"
-    )
-  else
-    TRAIN_COMMAND=(uv run lerobot-train)
-  fi
-fi
-
-echo "Training ${VARIANT}: ${TARGET_SAMPLES} target samples, global batch ${EFFECTIVE_BATCH_SIZE}, ${STEPS} optimizer steps"
-
-COMMON_TRAIN_ARGS=(
-  --policy.path=lerobot/smolvla_base
-  --policy.device=cuda
-  --policy.push_to_hub=false
-  --policy.empty_cameras=1
-  --policy.freeze_vision_encoder=false
-  --policy.train_expert_only=false
-  --dataset.repo_id=lerobot/robomme
-  --dataset.training_target_start_feature=exec_start_idx
-  '--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
-  --batch_size="${BATCH_SIZE}"
-  --gradient_accumulation_steps="${GRADIENT_ACCUMULATION_STEPS}"
-  --steps="${STEPS}"
-  --policy.scheduler_warmup_steps="${SCHEDULER_WARMUP_STEPS}"
-  --policy.scheduler_decay_steps="${SCHEDULER_DECAY_STEPS}"
-  --seed="${SEED}"
-  --env_eval_freq=0
-  --save_freq=5000
-  --wandb.enable="${WANDB_ENABLE}"
-)
-
-if [[ "${RUN_TRAIN}" == "true" ]]; then
-  for variant in "${VARIANTS[@]}"; do
-    MEMORY_ARGS=(--policy.use_visual_memory=false)
-    if [[ "${variant}" == "visual-memory" ]]; then
-      MEMORY_ARGS=(
-        --policy.use_visual_memory=true
-        --policy.visual_memory_frames=6
-        --policy.visual_memory_stride=10
-        --policy.visual_memory_temporal_attention_every=4
-      )
-    fi
-    "${TRAIN_COMMAND[@]}" \
-      "${COMMON_TRAIN_ARGS[@]}" \
-      "${MEMORY_ARGS[@]}" \
-      --output_dir="${OUTPUT_ROOT}/${variant}" \
-      --job_name="robomme-smolvla-${variant}"
-  done
-fi
-
-if [[ "${RUN_EVAL}" == "true" ]]; then
-  for variant in "${VARIANTS[@]}"; do
-    uv run lerobot-eval \
-      --policy.path="${OUTPUT_ROOT}/${variant}/checkpoints/last/pretrained_model" \
-      --env.type=robomme \
-      --env.task="${TASKS}" \
-      --env.dataset_split=test \
-      --env.task_ids="${TASK_IDS}" \
-      '--rename_map={"observation.images.image":"observation.images.camera1","observation.images.wrist_image":"observation.images.camera2"}' \
-      --eval.batch_size=1 \
-      --eval.n_episodes=1 \
-      --seed="${SEED}" \
-      --output_dir="${OUTPUT_ROOT}/eval-${variant}"
-  done
-fi
-
-if [[ "${RUN_EVAL}" == "true" ]]; then
-  uv run python - "${OUTPUT_ROOT}" <<'PY'
-import json
-import pathlib
-import sys
-
-root = pathlib.Path(sys.argv[1])
-for variant in ("baseline", "visual-memory"):
-    result_path = root / f"eval-{variant}" / "eval_info.json"
-    if not result_path.exists():
-        continue
-    with result_path.open() as handle:
-        info = json.load(handle)
-    overall = info["overall"]
-    print(
-        f"{variant}: success={overall['pc_success']:.2f}% "
-        f"avg_reward={overall['avg_sum_reward']:.4f}"
-    )
-    for task, metrics in info["per_group"].items():
-        print(
-            f"  {task}: success={metrics['pc_success']:.2f}% "
-            f"avg_reward={metrics['avg_sum_reward']:.4f}"
-        )
-PY
-fi

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,10 +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
-    # Optional per-frame feature containing the episode-relative index of the first frame that may be
-    # sampled as a training target. Earlier frames remain in the dataset and can still be loaded through
-    # temporal observation deltas. This is useful for datasets with demonstration/context prefixes.
-    training_target_start_feature: str | 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
@@ -66,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/configs/train.py

@@ -99,9 +99,6 @@ class TrainPipelineConfig(HubMixin):
     # Number of workers for the dataloader.
     num_workers: int = 4
     batch_size: int = 8
-    # Number of microbatches accumulated before each optimizer update. The effective global batch is
-    # batch_size * accelerator.num_processes * gradient_accumulation_steps.
-    gradient_accumulation_steps: int = 1
     prefetch_factor: int = 4
     persistent_workers: bool = True
     steps: int = 100_000
@@ -224,8 +221,6 @@ class TrainPipelineConfig(HubMixin):
             )
 
         active_cfg = self.trainable_config
-        if self.gradient_accumulation_steps < 1:
-            raise ValueError("gradient_accumulation_steps must be at least 1.")
         if self.rename_map and active_cfg.pretrained_path is None:
             raise ValueError(
                 "`rename_map` requires a pretrained policy checkpoint. "

src/lerobot/datasets/factory.py

@@ -32,9 +32,7 @@ from .streaming_dataset import StreamingLeRobotDataset
 
 
 def resolve_delta_timestamps(
-    cfg: PreTrainedConfig | RewardModelConfig,
-    ds_meta: LeRobotDatasetMetadata,
-    rename_map: dict[str, str] | None = None,
+    cfg: PreTrainedConfig | RewardModelConfig, ds_meta: LeRobotDatasetMetadata
 ) -> dict[str, list] | None:
     """Resolves delta_timestamps by reading from the 'delta_indices' properties of the config.
 
@@ -55,12 +53,11 @@ def resolve_delta_timestamps(
     """
     delta_timestamps = {}
     for key in ds_meta.features:
-        policy_key = (rename_map or {}).get(key, key)
-        if policy_key == REWARD and cfg.reward_delta_indices is not None:
+        if key == REWARD and cfg.reward_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.reward_delta_indices]
-        if policy_key == ACTION and cfg.action_delta_indices is not None:
+        if key == ACTION and cfg.action_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.action_delta_indices]
-        if policy_key.startswith(OBS_PREFIX) and cfg.observation_delta_indices is not None:
+        if key.startswith(OBS_PREFIX) and cfg.observation_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.observation_delta_indices]
 
     if len(delta_timestamps) == 0:
@@ -69,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.
 
@@ -89,12 +97,15 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         ds_meta = LeRobotDatasetMetadata(
             cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
         )
-        delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta, cfg.rename_map)
+        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,
@@ -107,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,
@@ -133,9 +144,6 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         for key in dataset.meta.camera_keys:
             if key in dataset.meta.depth_keys:
                 continue  # Exclude depth keys from ImageNet stats
-            # Some video-only datasets omit camera statistics entirely. Visual
-            # normalization can still use the requested ImageNet defaults.
-            dataset.meta.stats.setdefault(key, {})
             for stats_type, stats in IMAGENET_STATS.items():
                 dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
 
@@ -181,7 +189,7 @@ def make_train_eval_datasets(
         f"(eval_split={cfg.dataset.eval_split}, {len(task_to_episodes)} tasks)"
     )
 
-    delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, full_dataset.meta, cfg.rename_map)
+    delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, full_dataset.meta)
 
     train_image_transforms = (
         ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None

src/lerobot/datasets/sampler.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 import logging
 import math
-from collections.abc import Iterator, Sequence
+from collections.abc import Iterator
 
 import numpy as np
 import torch
@@ -49,7 +49,7 @@ class EpisodeAwareSampler:
         dataset_from_indices: list[int],
         dataset_to_indices: list[int],
         episode_indices_to_use: list | None = None,
-        drop_n_first_frames: int | Sequence[int] = 0,
+        drop_n_first_frames: int = 0,
         drop_n_last_frames: int = 0,
         shuffle: bool = False,
         seed: int = 0,
@@ -60,12 +60,13 @@ class EpisodeAwareSampler:
             dataset_from_indices: Start index of each episode in the dataset.
             dataset_to_indices: End index of each episode in the dataset.
             episode_indices_to_use: Episode indices to use; None means all.
-            drop_n_first_frames: Frames to drop from the start of each episode. An integer applies the
-                same offset to every episode; a sequence supplies one offset per episode.
+            drop_n_first_frames: Frames to drop from the start of each episode.
             drop_n_last_frames: Frames to drop from the end of each episode.
             shuffle: Whether to shuffle the indices.
             seed: Seed the permutation is derived from (together with the epoch).
         """
+        if drop_n_first_frames < 0:
+            raise ValueError(f"drop_n_first_frames must be >= 0, got {drop_n_first_frames}")
         if drop_n_last_frames < 0:
             raise ValueError(f"drop_n_last_frames must be >= 0, got {drop_n_last_frames}")
 
@@ -77,27 +78,12 @@ class EpisodeAwareSampler:
                 f"got {len(from_indices)} and {len(to_indices)}"
             )
 
-        if isinstance(drop_n_first_frames, int):
-            first_frame_offsets = np.full(len(from_indices), drop_n_first_frames, dtype=np.int64)
-        else:
-            first_frame_offsets = np.asarray(drop_n_first_frames, dtype=np.int64)
-            if first_frame_offsets.shape != from_indices.shape:
-                raise ValueError(
-                    "drop_n_first_frames must be an integer or have one value per episode; "
-                    f"got {len(first_frame_offsets)} values for {len(from_indices)} episodes"
-                )
-        if np.any(first_frame_offsets < 0):
-            raise ValueError(
-                "drop_n_first_frames must be >= 0, got "
-                f"{first_frame_offsets[first_frame_offsets < 0].tolist()}"
-            )
-
         used = np.ones(len(from_indices), dtype=bool)
         if episode_indices_to_use is not None:
             used = np.zeros(len(from_indices), dtype=bool)
             used[np.asarray(episode_indices_to_use, dtype=np.int64)] = True
 
-        starts = from_indices + first_frame_offsets
+        starts = from_indices + drop_n_first_frames
         lengths = to_indices - drop_n_last_frames - starts
         for episode_idx in np.flatnonzero(used & (lengths <= 0)):
             logger.warning(
@@ -105,7 +91,7 @@ class EpisodeAwareSampler:
                 "drop_n_last_frames=%d removes all frames. Skipping.",
                 episode_idx,
                 to_indices[episode_idx] - from_indices[episode_idx],
-                first_frame_offsets[episode_idx],
+                drop_n_first_frames,
                 drop_n_last_frames,
             )
         used &= lengths > 0

src/lerobot/envs/robomme.py

@@ -63,8 +63,6 @@ class RoboMMEGymEnv(gym.Env):
         from robomme.env_record_wrapper import BenchmarkEnvBuilder
 
         self._task = task
-        self.task = task
-        self.task_description = task
         self._action_space_type = action_space_type
         self._dataset = dataset
         self._episode_idx = episode_idx
@@ -107,12 +105,6 @@ class RoboMMEGymEnv(gym.Env):
         )
         obs, info = self._env.reset()
         self._last_raw_obs = obs
-        task_goal = info.get("task_goal")
-        # RoboMME returns [simple_subgoal, grounded_subgoal]. The published
-        # LeRobot training dataset uses the simple subgoal as its `task` text.
-        if isinstance(task_goal, (list, tuple)):
-            task_goal = task_goal[0] if task_goal else ""
-        self.task_description = str(task_goal or self._task)
         return self._convert_obs(obs), self._convert_info(info)
 
     def step(self, action):

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):
@@ -623,9 +622,6 @@ def make_policy(
             raise ValueError("env_cfg cannot be None when ds_meta is not provided")
         features = env_to_policy_features(env_cfg)
 
-    if rename_map:
-        features = {rename_map.get(key, key): feature for key, feature in features.items()}
-
     cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
     if not cfg.input_features:
         cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}

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/policies/smolvla/configuration_smolvla.py

@@ -44,14 +44,6 @@ class SmolVLAConfig(PreTrainedConfig):
     # Image preprocessing
     resize_imgs_with_padding: tuple[int, int] = (512, 512)
 
-    # MEM short-horizon visual memory (https://arxiv.org/abs/2603.03596).
-    # Past frames are fused inside the vision tower and dropped before the VLM,
-    # keeping the number of prefix tokens identical to single-frame SmolVLA.
-    use_visual_memory: bool = False
-    visual_memory_frames: int = 6
-    visual_memory_stride: int = 10
-    visual_memory_temporal_attention_every: int = 4
-
     # Add empty images. Used by smolvla_aloha_sim which adds the empty
     # left and right wrist cameras in addition to the top camera.
     empty_cameras: int = 0
@@ -127,12 +119,6 @@ class SmolVLAConfig(PreTrainedConfig):
             raise NotImplementedError(
                 "`use_delta_joint_actions_aloha` is used by smolvla for aloha real models. It is not ported yet in LeRobot."
             )
-        if self.visual_memory_frames < 1:
-            raise ValueError("visual_memory_frames must be at least 1")
-        if self.visual_memory_stride < 1:
-            raise ValueError("visual_memory_stride must be at least 1")
-        if self.visual_memory_temporal_attention_every < 1:
-            raise ValueError("visual_memory_temporal_attention_every must be at least 1")
 
     def validate_features(self) -> None:
         for i in range(self.empty_cameras):
@@ -162,10 +148,7 @@ class SmolVLAConfig(PreTrainedConfig):
 
     @property
     def observation_delta_indices(self) -> list:
-        if not self.use_visual_memory:
-            return [0]
-        horizon = (self.visual_memory_frames - 1) * self.visual_memory_stride
-        return list(range(-horizon, 1, self.visual_memory_stride))
+        return [0]
 
     @property
     def action_delta_indices(self) -> list:

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -71,7 +71,6 @@ from ..utils import (
 )
 from .configuration_smolvla import SmolVLAConfig
 from .smolvlm_with_expert import SmolVLMWithExpertModel
-from .visual_memory import sample_visual_history
 
 
 class ActionSelectKwargs(TypedDict, total=False):
@@ -254,10 +253,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
         self._queues = {
             ACTION: deque(maxlen=self.config.n_action_steps),
         }
-        if self.config.use_visual_memory:
-            history_length = (self.config.visual_memory_frames - 1) * self.config.visual_memory_stride + 1
-            self._queues.update({key: deque(maxlen=history_length) for key in self.config.image_features})
-            self._visual_memory_steps_seen = 0
 
     def init_rtc_processor(self):
         """Initialize RTC processor if RTC is enabled in config."""
@@ -286,15 +281,9 @@ class SmolVLAPolicy(PreTrainedPolicy):
         # In the case of offline inference, we have the action in the batch
         # that why without the k != ACTION check, it will raise an error because we are trying to stack
         # on an empty container.
-        for k in list(batch):
+        for k in batch:
             if k in self._queues and k != ACTION:
-                history = list(self._queues[k])
-                batch[k], batch[f"{k}_is_pad"] = sample_visual_history(
-                    history,
-                    num_frames=self.config.visual_memory_frames,
-                    stride=self.config.visual_memory_stride,
-                    steps_seen=self._visual_memory_steps_seen,
-                )
+                batch[k] = torch.stack(list(self._queues[k]), dim=1)
 
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
@@ -320,11 +309,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
 
         return batch
 
-    def _populate_observation_queues(self, batch: dict[str, Tensor]) -> None:
-        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
-        if self.config.use_visual_memory:
-            self._visual_memory_steps_seen += 1
-
     @torch.no_grad()
     def predict_action_chunk(
         self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
@@ -332,7 +316,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
         self.eval()
 
         batch = self._prepare_batch(batch)
-        self._populate_observation_queues(batch)
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
         actions = self._get_action_chunk(batch, noise, **kwargs)
         return actions
@@ -354,7 +338,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
 
         self.eval()
         batch = self._prepare_batch(batch)
-        self._populate_observation_queues(batch)
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
         if self._check_get_actions_condition():
             actions = self._get_action_chunk(batch, noise)
@@ -443,30 +427,19 @@ class SmolVLAPolicy(PreTrainedPolicy):
             )
         # Preprocess image features present in the batch
         for key in present_img_keys:
-            img = batch[key]
-            if img.ndim == 5 and not self.config.use_visual_memory:
-                img = img[:, -1, :, :, :]
+            img = batch[key][:, -1, :, :, :] if batch[key].ndim == 5 else batch[key]
             if self.config.resize_imgs_with_padding is not None:
-                if img.ndim == 5:
-                    batch_size, num_frames = img.shape[:2]
-                    img = resize_with_pad(
-                        img.flatten(0, 1), *self.config.resize_imgs_with_padding, pad_value=0
-                    ).unflatten(0, (batch_size, num_frames))
-                else:
-                    img = resize_with_pad(img, *self.config.resize_imgs_with_padding, pad_value=0)
+                img = resize_with_pad(img, *self.config.resize_imgs_with_padding, pad_value=0)
 
             # Normalize from range [0,1] to [-1,1] as expacted by siglip
             img = img * 2.0 - 1.0
 
             bsize = img.shape[0]
             device = img.device
-            if f"{key}_is_pad" in batch:
-                mask = ~batch[f"{key}_is_pad"].bool()
-            elif f"{key}_padding_mask" in batch:
+            if f"{key}_padding_mask" in batch:
                 mask = batch[f"{key}_padding_mask"].bool()
             else:
-                mask_shape = img.shape[:2] if img.ndim == 5 else (bsize,)
-                mask = torch.ones(mask_shape, dtype=torch.bool, device=device)
+                mask = torch.ones(bsize, dtype=torch.bool, device=device)
             images.append(img)
             img_masks.append(mask)
 
@@ -689,11 +662,7 @@ class VLAFlowMatching(nn.Module):
                 embs.append(image_start_token)
                 pad_masks.append(image_start_mask)
 
-            img_emb = self.vlm_with_expert.embed_image(
-                img,
-                frame_mask=img_mask if img.ndim == 5 else None,
-                temporal_attention_every=self.config.visual_memory_temporal_attention_every,
-            )
+            img_emb = self.vlm_with_expert.embed_image(img)
             img_emb = img_emb
 
             # Normalize image embeddings
@@ -701,8 +670,6 @@ class VLAFlowMatching(nn.Module):
             img_emb = img_emb * torch.tensor(img_emb_dim**0.5, dtype=img_emb.dtype, device=img_emb.device)
 
             bsize, num_img_embs = img_emb.shape[:2]
-            if img_mask.ndim == 2:
-                img_mask = img_mask[:, -1]
             img_mask = img_mask[:, None].expand(bsize, num_img_embs)
 
             embs.append(img_emb)

src/lerobot/policies/smolvla/smolvlm_with_expert.py

@@ -20,8 +20,6 @@ from torch import nn
 
 from lerobot.utils.import_utils import _transformers_available, require_package
 
-from .visual_memory import encode_video_with_mem
-
 if TYPE_CHECKING or _transformers_available:
     from transformers import (
         AutoConfig,
@@ -190,31 +188,17 @@ class SmolVLMWithExpertModel(nn.Module):
         if self.train_expert_only:
             self.vlm.eval()
 
-    def embed_image(
-        self,
-        image: torch.Tensor,
-        *,
-        frame_mask: torch.Tensor | None = None,
-        temporal_attention_every: int = 4,
-    ):
+    def embed_image(self, image: torch.Tensor):
         patch_attention_mask = None
         # Get sequence from the vision encoder
-        vision_model = self.get_vlm_model().vision_model
-        image = image.to(dtype=vision_model.dtype)
-        if image.ndim == 5:
-            if frame_mask is None:
-                frame_mask = torch.ones(image.shape[:2], dtype=torch.bool, device=image.device)
-            image_hidden_states = encode_video_with_mem(
-                vision_model,
-                image,
-                frame_mask,
-                temporal_attention_every=temporal_attention_every,
-            )
-        else:
-            image_hidden_states = vision_model(
-                pixel_values=image,
+        image_hidden_states = (
+            self.get_vlm_model()
+            .vision_model(
+                pixel_values=image.to(dtype=self.get_vlm_model().vision_model.dtype),
                 patch_attention_mask=patch_attention_mask,
-            ).last_hidden_state
+            )
+            .last_hidden_state
+        )
         # Modality projection & resampling
         image_hidden_states = self.get_vlm_model().connector(image_hidden_states)
         return image_hidden_states

src/lerobot/policies/smolvla/visual_memory.py

@@ -1,161 +0,0 @@
-# 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.
-
-"""Short-horizon visual memory from MEM (arXiv:2603.03596).
-
-The encoder keeps SigLIP's pretrained parameters and alternates its ordinary
-per-frame spatial attention with causal attention across time for matching
-spatial patches. Only the current frame's patch tokens leave the vision tower,
-so enabling memory does not increase the VLM prefix length.
-"""
-
-import math
-
-import torch
-from torch import Tensor
-
-
-def sample_visual_history(
-    history: list[Tensor], *, num_frames: int, stride: int, steps_seen: int
-) -> tuple[Tensor, Tensor]:
-    """Subsample an inference queue and mark pre-episode padding frames."""
-    sampled = history[::stride][-num_frames:]
-    if len(sampled) != num_frames:
-        raise ValueError(f"visual history has {len(sampled)} samples, expected {num_frames}")
-    video = torch.stack(sampled, dim=1)
-    required_ages = range((num_frames - 1) * stride, -1, -stride)
-    valid = torch.tensor([steps_seen > age for age in required_ages], dtype=torch.bool, device=video.device)
-    padding_mask = (~valid)[None, :].expand(video.shape[0], -1)
-    return video, padding_mask
-
-
-def temporal_sinusoidal_embedding(
-    num_frames: int, hidden_size: int, *, device: torch.device, dtype: torch.dtype
-) -> Tensor:
-    """Return fixed temporal embeddings with an exactly-zero current position."""
-    if hidden_size % 2:
-        raise ValueError(f"hidden_size must be even, got {hidden_size}")
-
-    # History is ordered oldest -> current, matching t in [-K, 0] in MEM.
-    positions = torch.arange(1 - num_frames, 1, device=device, dtype=torch.float32)[:, None]
-    frequencies = torch.exp(
-        torch.arange(0, hidden_size, 2, device=device, dtype=torch.float32)
-        * (-math.log(10_000.0) / hidden_size)
-    )[None, :]
-    angles = positions * frequencies
-    embedding = torch.zeros(num_frames, hidden_size, device=device, dtype=torch.float32)
-    embedding[:, 0::2] = torch.sin(angles)
-    embedding[:, 1::2] = torch.cos(angles) - 1.0
-    return embedding.to(dtype=dtype)
-
-
-def causal_temporal_mask(frame_mask: Tensor, *, dtype: torch.dtype, num_patches: int) -> Tensor:
-    """Build an additive causal/key-padding mask for per-patch temporal attention."""
-    if frame_mask.ndim != 2:
-        raise ValueError(f"frame_mask must have shape (batch, frames), got {tuple(frame_mask.shape)}")
-
-    batch_size, num_frames = frame_mask.shape
-    allowed = torch.ones(num_frames, num_frames, dtype=torch.bool, device=frame_mask.device).tril()
-    allowed = allowed[None, :, :] & frame_mask[:, None, :].bool()
-    # The current frame is always present in normal use. Keeping the diagonal
-    # valid also prevents NaNs for fully padded historical query rows.
-    diagonal = torch.eye(num_frames, dtype=torch.bool, device=frame_mask.device)[None, :, :]
-    allowed = allowed | diagonal
-    mask = torch.zeros(batch_size, 1, num_frames, num_frames, dtype=dtype, device=frame_mask.device)
-    mask.masked_fill_(~allowed[:, None, :, :], torch.finfo(dtype).min)
-    return mask.repeat_interleave(num_patches, dim=0)
-
-
-def encode_video_with_mem(
-    vision_model,
-    pixel_values: Tensor,
-    frame_mask: Tensor,
-    *,
-    temporal_attention_every: int,
-) -> Tensor:
-    """Encode ``(B,T,C,H,W)`` video with MEM's space-time separable attention.
-
-    No modules or learnable parameters are added. At temporal layers, the same
-    layer-normalization and Q/K/V/out projections are reused for a second,
-    causal attention operation along time. For a one-frame input the temporal
-    branch is skipped, making this exactly equivalent to the original image
-    encoder.
-    """
-    if pixel_values.ndim != 5:
-        raise ValueError(f"pixel_values must have shape (B,T,C,H,W), got {tuple(pixel_values.shape)}")
-    if temporal_attention_every < 1:
-        raise ValueError("temporal_attention_every must be at least 1")
-
-    batch_size, num_frames, channels, height, width = pixel_values.shape
-    if frame_mask.shape != (batch_size, num_frames):
-        raise ValueError(
-            f"frame_mask must have shape {(batch_size, num_frames)}, got {tuple(frame_mask.shape)}"
-        )
-
-    required = ("embeddings", "encoder", "post_layernorm")
-    if any(not hasattr(vision_model, name) for name in required):
-        raise TypeError("MEM visual memory currently requires a SigLIP-compatible vision tower")
-
-    flat_pixels = pixel_values.reshape(batch_size * num_frames, channels, height, width)
-    if hasattr(vision_model, "patch_size"):
-        patch_size = vision_model.patch_size
-        patch_attention_mask = torch.ones(
-            batch_size * num_frames,
-            height // patch_size,
-            width // patch_size,
-            dtype=torch.bool,
-            device=pixel_values.device,
-        )
-        hidden_states = vision_model.embeddings(
-            pixel_values=flat_pixels,
-            patch_attention_mask=patch_attention_mask,
-        )
-    else:
-        hidden_states = vision_model.embeddings(flat_pixels)
-    num_patches, hidden_size = hidden_states.shape[1:]
-    hidden_states = hidden_states.reshape(batch_size, num_frames, num_patches, hidden_size)
-
-    temporal_positions = temporal_sinusoidal_embedding(
-        num_frames, hidden_size, device=hidden_states.device, dtype=hidden_states.dtype
-    )[None, :, None, :]
-    temporal_mask = causal_temporal_mask(frame_mask, dtype=hidden_states.dtype, num_patches=num_patches)
-
-    for layer_index, layer in enumerate(vision_model.encoder.layers):
-        spatial_input = hidden_states.reshape(batch_size * num_frames, num_patches, hidden_size)
-        if num_frames == 1 or (layer_index + 1) % temporal_attention_every:
-            hidden_states = layer(spatial_input, attention_mask=None).reshape(
-                batch_size, num_frames, num_patches, hidden_size
-            )
-            continue
-
-        residual = hidden_states
-        spatial_norm = layer.layer_norm1(spatial_input)
-        spatial_output, _ = layer.self_attn(hidden_states=spatial_norm, attention_mask=None)
-        spatial_output = spatial_output.reshape(batch_size, num_frames, num_patches, hidden_size)
-
-        temporal_input = (hidden_states + temporal_positions).permute(0, 2, 1, 3)
-        temporal_input = temporal_input.reshape(batch_size * num_patches, num_frames, hidden_size)
-        temporal_norm = layer.layer_norm1(temporal_input)
-        temporal_output, _ = layer.self_attn(
-            hidden_states=temporal_norm,
-            attention_mask=temporal_mask,
-        )
-        temporal_output = temporal_output.reshape(batch_size, num_patches, num_frames, hidden_size)
-        temporal_output = temporal_output.permute(0, 2, 1, 3)
-
-        hidden_states = residual + spatial_output + temporal_output
-        hidden_states = hidden_states + layer.mlp(layer.layer_norm2(hidden_states))
-
-    # MEM drops all historical tokens before the VLA backbone.
-    return vision_model.post_layernorm(hidden_states[:, -1])

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/rename_processor.py

@@ -21,23 +21,6 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
 
 from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
 
-_AUXILIARY_KEY_SUFFIXES = ("_is_pad", "_padding_mask")
-
-
-def rename_transition_key(key: str, rename_map: dict[str, str]) -> str:
-    """Rename a feature key and any sampling metadata derived from it."""
-    if key in rename_map:
-        return rename_map[key]
-    for suffix in _AUXILIARY_KEY_SUFFIXES:
-        if key.endswith(suffix) and key[: -len(suffix)] in rename_map:
-            return f"{rename_map[key[: -len(suffix)]]}{suffix}"
-    return key
-
-
-def rename_transition_keys(data: dict[str, Any], rename_map: dict[str, str]) -> dict[str, Any]:
-    """Rename all transition keys, including delta-sampling padding masks."""
-    return {rename_transition_key(key, rename_map): value for key, value in data.items()}
-
 
 @dataclass
 @ProcessorStepRegistry.register(name="rename_observations_processor")
@@ -58,7 +41,14 @@ class RenameObservationsProcessorStep(ObservationProcessorStep):
     rename_map: dict[str, str] = field(default_factory=dict)
 
     def observation(self, observation):
-        return rename_transition_keys(observation, self.rename_map)
+        processed_obs = {}
+        for key, value in observation.items():
+            if key in self.rename_map:
+                processed_obs[self.rename_map[key]] = value
+            else:
+                processed_obs[key] = value
+
+        return processed_obs
 
     def get_config(self) -> dict[str, Any]:
         return {"rename_map": self.rename_map}

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

@@ -57,7 +57,6 @@ from lerobot.envs import close_envs, make_env, make_env_pre_post_processors
 from lerobot.jobs import submit_to_hf
 from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
-from lerobot.processor.rename_processor import rename_transition_keys
 from lerobot.rewards import make_reward_pre_post_processors
 from lerobot.utils.collate import lerobot_collate_fn
 from lerobot.utils.import_utils import register_third_party_plugins
@@ -85,7 +84,6 @@ def update_policy(
     lock=None,
     sample_weighter=None,
     log_metrics: bool = True,
-    track_update_time: bool = True,
 ) -> tuple[MetricsTracker, dict | None]:
     """
     Performs a single training step to update the policy's weights.
@@ -149,16 +147,13 @@ def update_policy(
     # Use accelerator's backward method
     accelerator.backward(loss)
 
-    # Accelerate suppresses gradient synchronization and optimizer updates on intermediate
-    # microbatches. Clip and report the norm only when the accumulated update is complete.
-    grad_norm = None
-    if accelerator.sync_gradients:
-        if grad_clip_norm > 0:
-            grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
-        else:
-            grad_norm = torch.nn.utils.clip_grad_norm_(
-                policy.parameters(), float("inf"), error_if_nonfinite=False
-            )
+    # Clip gradients if specified
+    if grad_clip_norm > 0:
+        grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
+    else:
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            policy.parameters(), float("inf"), error_if_nonfinite=False
+        )
 
     # Optimizer step
     with lock if lock is not None else nullcontext():
@@ -167,24 +162,19 @@ def update_policy(
     optimizer.zero_grad()
 
     # Step through pytorch scheduler at every batch instead of epoch
-    if lr_scheduler is not None and accelerator.sync_gradients:
+    if lr_scheduler is not None:
         lr_scheduler.step()
 
     # Update internal buffers if policy has update method
-    if accelerator.sync_gradients and has_method(
-        accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"
-    ):
+    if has_method(accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"):
         accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update()
 
-    if accelerator.sync_gradients:
-        train_metrics.lr = optimizer.param_groups[0]["lr"]
-    if torch.cuda.is_available() and accelerator.sync_gradients:
+    train_metrics.lr = optimizer.param_groups[0]["lr"]
+    if torch.cuda.is_available():
         train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
     train_metrics.accumulate_tensor("loss", loss)
-    if grad_norm is not None:
-        train_metrics.accumulate_tensor("grad_norm", grad_norm)
-    if track_update_time:
-        train_metrics.update_s = time.perf_counter() - start_time
+    train_metrics.accumulate_tensor("grad_norm", grad_norm)
+    train_metrics.update_s = time.perf_counter() - start_time
     # Synchronize accumulated GPU metrics only when logging.
     if log_metrics:
         train_metrics.materialize_tensors()
@@ -247,7 +237,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         mixed_precision = {"bfloat16": "bf16", "float16": "fp16", "float32": "no"}.get(policy_dtype)
         accelerator = Accelerator(
             step_scheduler_with_optimizer=False,
-            gradient_accumulation_steps=cfg.gradient_accumulation_steps,
             mixed_precision=mixed_precision,
             kwargs_handlers=[ddp_kwargs, ipg_kwargs],
             cpu=force_cpu,
@@ -354,22 +343,25 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         )
         processor_pretrained_path = None
 
-    dataset_stats = {cfg.rename_map.get(key, key): value for key, value in dataset.meta.stats.items()}
     processor_kwargs = {}
     if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
-        processor_kwargs["dataset_stats"] = dataset_stats
+        processor_kwargs["dataset_stats"] = dataset.meta.stats
 
     if cfg.is_reward_model_training:
         processor_kwargs["dataset_meta"] = dataset.meta
 
     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 = {
             "device_processor": {"device": device.type},
             "normalizer_processor": {
-                "stats": dataset_stats,
+                "stats": dataset.meta.stats,
                 "features": {**policy.config.input_features, **policy.config.output_features},
                 "norm_map": policy.config.normalization_mapping,
             },
@@ -377,7 +369,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         }
         postprocessor_overrides = {
             "unnormalizer_processor": {
-                "stats": dataset_stats,
+                "stats": dataset.meta.stats,
                 "features": policy.config.output_features,
                 "norm_map": policy.config.normalization_mapping,
             },
@@ -449,11 +441,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
         logging.info(f"{dataset.num_episodes=}")
         num_processes = accelerator.num_processes
-        effective_bs = cfg.batch_size * num_processes * cfg.gradient_accumulation_steps
-        logging.info(
-            "Effective batch size: "
-            f"{cfg.batch_size} x {num_processes} x {cfg.gradient_accumulation_steps} = {effective_bs}"
-        )
+        effective_bs = cfg.batch_size * num_processes
+        logging.info(f"Effective batch size: {cfg.batch_size} x {num_processes} = {effective_bs}")
         logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
         logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
 
@@ -468,46 +457,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         to_indices = dataset.meta.episodes["dataset_to_index"]
         seed = cfg.seed if cfg.seed is not None else 0
 
-        drop_n_first_frames: int | list[int] = 0
-        target_start_feature = cfg.dataset.training_target_start_feature
-        if target_start_feature is not None:
-            if not hasattr(dataset, "hf_dataset"):
-                raise ValueError(
-                    "dataset.training_target_start_feature is only supported for a single map-style "
-                    "LeRobotDataset."
-                )
-            if target_start_feature not in dataset.hf_dataset.column_names:
-                raise ValueError(
-                    f"Training target start feature {target_start_feature!r} is not present in the dataset. "
-                    f"Available features: {dataset.hf_dataset.column_names}"
-                )
-            start_column = dataset.hf_dataset.data.column(target_start_feature)
-            absolute_to_relative_idx = dataset.absolute_to_relative_idx
-            episode_indices = dataset.episodes if dataset.episodes is not None else range(len(from_indices))
-            drop_n_first_frames = [0] * len(from_indices)
-            for episode_index in episode_indices:
-                absolute_index = int(from_indices[episode_index])
-                relative_index = (
-                    absolute_to_relative_idx[absolute_index]
-                    if absolute_to_relative_idx is not None
-                    else absolute_index
-                )
-                drop_n_first_frames[episode_index] = int(start_column[relative_index].as_py())
-            logging.info(
-                "Restricting training targets with %s: keeping %d of %d frames",
-                target_start_feature,
-                sum(
-                    int(to_indices[index]) - int(from_indices[index]) - drop_n_first_frames[index]
-                    for index in episode_indices
-                ),
-                sum(int(to_indices[index]) - int(from_indices[index]) for index in episode_indices),
-            )
-
         sampler = EpisodeAwareSampler(
             from_indices,
             to_indices,
             episode_indices_to_use=dataset.episodes,
-            drop_n_first_frames=drop_n_first_frames,
             drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
             shuffle=True,
             seed=seed,
@@ -533,12 +486,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     f"batch_size={saved_batch_size}. The data order resumes at the right epoch/offset, "
                     "but per-rank sample-exactness requires the same batch size."
                 )
-            sampler_state = compute_sampler_state(
-                step,
-                len(sampler),
-                ckpt_batch_size * cfg.gradient_accumulation_steps,
-                ckpt_num_processes,
-            )
+            sampler_state = compute_sampler_state(step, len(sampler), ckpt_batch_size, ckpt_num_processes)
             sampler.load_state_dict(sampler_state)
             if is_main_process:
                 logging.info(
@@ -594,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
@@ -635,9 +584,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         train_metrics["gpu_mem_gb"] = AverageMeter("mem_gb", ":.2f", reduction="max")
 
     # Keep global batch size for logging; MetricsTracker handles world size internally.
-    effective_batch_size = cfg.batch_size * accelerator.num_processes * cfg.gradient_accumulation_steps
+    effective_batch_size = cfg.batch_size * accelerator.num_processes
     train_tracker = MetricsTracker(
-        cfg.batch_size * cfg.gradient_accumulation_steps,
+        cfg.batch_size,
         dataset.num_frames,
         dataset.num_episodes,
         train_metrics,
@@ -659,42 +608,28 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         )
 
     for _ in range(step, cfg.steps):
-        update_start_time = time.perf_counter()
-        dataloading_s = 0.0
-        output_dict = None
-        for microbatch_idx in range(cfg.gradient_accumulation_steps):
-            start_time = time.perf_counter()
-            batch = next(dl_iter)
-            for cam_key in dataset.meta.camera_keys:
-                if cam_key in batch and batch[cam_key].dtype == torch.uint8:
-                    batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0
-            if cfg.rename_map:
-                batch = rename_transition_keys(batch, cfg.rename_map)
-            batch = preprocessor(batch)
-            dataloading_s += time.perf_counter() - start_time
+        start_time = time.perf_counter()
+        batch = next(dl_iter)
+        for cam_key in dataset.meta.camera_keys:
+            if cam_key in batch and batch[cam_key].dtype == torch.uint8:
+                batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0
+        batch = preprocessor(batch)
+        train_tracker.dataloading_s = time.perf_counter() - start_time
 
-            # Synchronize GPU metrics only on the final microbatch of logged optimizer updates.
-            log_metrics = (
-                cfg.log_freq > 0
-                and (step + 1) % cfg.log_freq == 0
-                and microbatch_idx == cfg.gradient_accumulation_steps - 1
-            )
+        # Synchronize GPU metrics only for updates that will be logged.
+        log_metrics = cfg.log_freq > 0 and (step + 1) % cfg.log_freq == 0
 
-            with accelerator.accumulate(policy):
-                train_tracker, output_dict = update_policy(
-                    train_tracker,
-                    policy,
-                    batch,
-                    optimizer,
-                    cfg.optimizer.grad_clip_norm,
-                    accelerator=accelerator,
-                    lr_scheduler=lr_scheduler,
-                    sample_weighter=sample_weighter,
-                    log_metrics=log_metrics,
-                    track_update_time=False,
-                )
-        train_tracker.dataloading_s = dataloading_s
-        train_tracker.update_s = time.perf_counter() - update_start_time - dataloading_s
+        train_tracker, output_dict = update_policy(
+            train_tracker,
+            policy,
+            batch,
+            optimizer,
+            cfg.optimizer.grad_clip_norm,
+            accelerator=accelerator,
+            lr_scheduler=lr_scheduler,
+            sample_weighter=sample_weighter,
+            log_metrics=log_metrics,
+        )
 
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
         # increment `step` here.

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/datasets/test_sampler.py

@@ -60,26 +60,6 @@ def test_drop_n_first_frames():
     assert list(sampler) == [1, 4, 5]
 
 
-def test_drop_different_number_of_first_frames_per_episode():
-    sampler = EpisodeAwareSampler(
-        dataset_from_indices=[0, 3, 5],
-        dataset_to_indices=[3, 5, 9],
-        drop_n_first_frames=[1, 0, 3],
-    )
-    assert sampler.indices == [1, 2, 3, 4, 8]
-    assert len(sampler) == 5
-
-
-def test_drop_first_frames_sequence_must_match_episode_count():
-    with pytest.raises(ValueError, match="one value per episode"):
-        EpisodeAwareSampler([0, 3], [3, 6], drop_n_first_frames=[1])
-
-
-def test_drop_first_frames_sequence_must_be_non_negative():
-    with pytest.raises(ValueError, match="must be >= 0"):
-        EpisodeAwareSampler([0, 3], [3, 6], drop_n_first_frames=[1, -1])
-
-
 def test_drop_n_last_frames():
     dataset = Dataset.from_dict(
         {

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/policies/smolvla/test_visual_memory.py

@@ -1,210 +0,0 @@
-# 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 pytest
-import torch
-
-from lerobot.datasets.factory import resolve_delta_timestamps
-from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
-from lerobot.policies.smolvla.visual_memory import (
-    causal_temporal_mask,
-    encode_video_with_mem,
-    sample_visual_history,
-    temporal_sinusoidal_embedding,
-)
-
-
-def test_visual_memory_observation_delta_indices():
-    baseline = SmolVLAConfig()
-    memory = SmolVLAConfig(use_visual_memory=True, visual_memory_frames=6, visual_memory_stride=10)
-
-    assert baseline.observation_delta_indices == [0]
-    assert memory.observation_delta_indices == [-50, -40, -30, -20, -10, 0]
-
-
-def test_delta_timestamps_respect_raw_dataset_rename_map():
-    class RawMetadata:
-        fps = 10
-        features = {"image": {}, "state": {}, "actions": {}}
-
-    config = SmolVLAConfig(use_visual_memory=True, visual_memory_frames=3, visual_memory_stride=5)
-    delta_timestamps = resolve_delta_timestamps(
-        config,
-        RawMetadata(),
-        {
-            "image": "observation.images.camera1",
-            "state": "observation.state",
-            "actions": "action",
-        },
-    )
-
-    assert delta_timestamps["image"] == [-1.0, -0.5, 0.0]
-    assert delta_timestamps["state"] == [-1.0, -0.5, 0.0]
-    assert delta_timestamps["actions"] == [index / 10 for index in range(50)]
-
-
-@pytest.mark.parametrize(
-    ("field", "value"),
-    [
-        ("visual_memory_frames", 0),
-        ("visual_memory_stride", 0),
-        ("visual_memory_temporal_attention_every", 0),
-    ],
-)
-def test_visual_memory_config_rejects_non_positive_values(field, value):
-    with pytest.raises(ValueError, match=field):
-        SmolVLAConfig(**{field: value})
-
-
-def test_current_temporal_position_is_exactly_zero():
-    embedding = temporal_sinusoidal_embedding(4, 16, device=torch.device("cpu"), dtype=torch.float32)
-
-    torch.testing.assert_close(embedding[-1], torch.zeros(16))
-    assert torch.count_nonzero(embedding[:-1]) > 0
-
-
-def test_causal_temporal_mask_combines_causality_and_padding():
-    frame_mask = torch.tensor([[False, True, True]])
-    mask = causal_temporal_mask(frame_mask, dtype=torch.float32, num_patches=2)
-
-    assert mask.shape == (2, 1, 3, 3)
-    assert mask[0, 0, 1, 0] < -1e30
-    assert mask[0, 0, 1, 1] == 0
-    assert mask[0, 0, 1, 2] < -1e30
-    assert mask[0, 0, 2, 1] == 0
-
-
-def test_inference_history_matches_training_order_and_padding():
-    history = [torch.full((2, 1), value) for value in range(11)]
-
-    initial_video, initial_padding = sample_visual_history(history, num_frames=3, stride=5, steps_seen=1)
-    full_video, full_padding = sample_visual_history(history, num_frames=3, stride=5, steps_seen=11)
-
-    assert initial_video[:, :, 0].tolist() == [[0, 5, 10], [0, 5, 10]]
-    assert initial_padding.tolist() == [[True, True, False], [True, True, False]]
-    torch.testing.assert_close(full_video[:, :, 0], torch.tensor([[0, 5, 10], [0, 5, 10]]))
-    assert not full_padding.any()
-
-
-def test_single_frame_mem_matches_original_siglip_encoder():
-    transformers = pytest.importorskip("transformers")
-    config = transformers.SiglipVisionConfig(
-        image_size=16,
-        patch_size=8,
-        hidden_size=16,
-        intermediate_size=32,
-        num_hidden_layers=4,
-        num_attention_heads=4,
-        vision_use_head=False,
-    )
-    from transformers.models.siglip.modeling_siglip import SiglipVisionTransformer
-
-    model = SiglipVisionTransformer(config).eval()
-    image = torch.randn(2, 3, 16, 16)
-
-    expected = model(image).last_hidden_state
-    actual = encode_video_with_mem(
-        model,
-        image[:, None],
-        torch.ones(2, 1, dtype=torch.bool),
-        temporal_attention_every=4,
-    )
-
-    torch.testing.assert_close(actual, expected)
-
-
-def test_mem_video_encoder_compresses_time_without_new_parameters():
-    transformers = pytest.importorskip("transformers")
-    config = transformers.SiglipVisionConfig(
-        image_size=16,
-        patch_size=8,
-        hidden_size=16,
-        intermediate_size=32,
-        num_hidden_layers=4,
-        num_attention_heads=4,
-        vision_use_head=False,
-    )
-    from transformers.models.siglip.modeling_siglip import SiglipVisionTransformer
-
-    model = SiglipVisionTransformer(config)
-    parameter_ids = {id(parameter) for parameter in model.parameters()}
-    video = torch.randn(2, 3, 3, 16, 16, requires_grad=True)
-
-    output = encode_video_with_mem(
-        model,
-        video,
-        torch.ones(2, 3, dtype=torch.bool),
-        temporal_attention_every=4,
-    )
-    output.sum().backward()
-
-    assert output.shape == (2, 4, 16)
-    assert {id(parameter) for parameter in model.parameters()} == parameter_ids
-    assert video.grad is not None
-
-
-def test_mem_video_encoder_supports_smolvlm_vision_tower():
-    transformers = pytest.importorskip("transformers")
-    config = transformers.SmolVLMVisionConfig(
-        image_size=16,
-        patch_size=8,
-        hidden_size=16,
-        intermediate_size=32,
-        num_hidden_layers=4,
-        num_attention_heads=4,
-    )
-    from transformers.models.smolvlm.modeling_smolvlm import SmolVLMVisionTransformer
-
-    model = SmolVLMVisionTransformer(config).eval()
-    video = torch.randn(2, 3, 3, 16, 16)
-
-    output = encode_video_with_mem(
-        model,
-        video,
-        torch.ones(2, 3, dtype=torch.bool),
-        temporal_attention_every=4,
-    )
-    single_frame = encode_video_with_mem(
-        model,
-        video[:, -1:],
-        torch.ones(2, 1, dtype=torch.bool),
-        temporal_attention_every=4,
-    )
-
-    assert output.shape == (2, 4, 16)
-    torch.testing.assert_close(single_frame, model(video[:, -1]).last_hidden_state)
-
-
-def test_masked_history_cannot_change_current_embedding():
-    transformers = pytest.importorskip("transformers")
-    config = transformers.SmolVLMVisionConfig(
-        image_size=16,
-        patch_size=8,
-        hidden_size=16,
-        intermediate_size=32,
-        num_hidden_layers=4,
-        num_attention_heads=4,
-    )
-    from transformers.models.smolvlm.modeling_smolvlm import SmolVLMVisionTransformer
-
-    model = SmolVLMVisionTransformer(config).eval()
-    first_video = torch.randn(1, 3, 3, 16, 16)
-    second_video = first_video.clone()
-    second_video[:, :2] = torch.randn_like(second_video[:, :2]) * 100
-    frame_mask = torch.tensor([[False, False, True]])
-
-    first_output = encode_video_with_mem(model, first_video, frame_mask, temporal_attention_every=4)
-    second_output = encode_video_with_mem(model, second_video, frame_mask, temporal_attention_every=4)
-
-    torch.testing.assert_close(first_output, second_output)

tests/processor/test_rename_processor.py

@@ -27,7 +27,7 @@ from lerobot.processor import (
     TransitionKey,
 )
 from lerobot.processor.converters import create_transition, identity_transition
-from lerobot.processor.rename_processor import rename_stats, rename_transition_keys
+from lerobot.processor.rename_processor import rename_stats
 from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from tests.conftest import assert_contract_is_typed
 
@@ -64,22 +64,6 @@ def test_basic_renaming():
     assert processed_obs["unchanged_key"] == "keep_me"
 
 
-def test_renaming_preserves_feature_suffixes_for_sampling_metadata():
-    data = {
-        "image": torch.zeros(1),
-        "image_is_pad": torch.ones(1, dtype=torch.bool),
-        "image_padding_mask": torch.ones(1, dtype=torch.bool),
-    }
-
-    result = rename_transition_keys(data, {"image": "observation.images.camera1"})
-
-    assert set(result) == {
-        "observation.images.camera1",
-        "observation.images.camera1_is_pad",
-        "observation.images.camera1_padding_mask",
-    }
-
-
 def test_empty_rename_map():
     """Test processor with empty rename map (should pass through unchanged)."""
     processor = RenameObservationsProcessorStep(rename_map={})

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

tests/test_robomme_env.py

@@ -44,10 +44,7 @@ def _install_robomme_stub():
                 "joint_state_list": [np.zeros(7, dtype=np.float32)],
                 "gripper_state_list": [np.zeros(2, dtype=np.float32)],
             }
-            env.reset.return_value = (
-                obs,
-                {"status": "ongoing", "task_goal": ["pick the cube", "pick the blue cube"]},
-            )
+            env.reset.return_value = (obs, {"status": "ongoing", "task_goal": "pick the cube"})
             env.step.return_value = (obs, 0.0, False, False, {"status": "ongoing", "task_goal": ""})
             return env
 
@@ -119,21 +116,6 @@ def test_robomme_features_action_dim_ee_pose():
 # ---------------------------------------------------------------------------
 
 
-def test_reset_exposes_episode_task_description():
-    """VLA evaluation receives the episode-specific language instruction."""
-    _install_robomme_stub()
-    try:
-        from lerobot.envs.robomme import RoboMMEGymEnv
-
-        env = RoboMMEGymEnv(task="PickXtimes")
-        env.reset()
-
-        assert env.task == "PickXtimes"
-        assert env.task_description == "pick the cube"
-    finally:
-        _uninstall_robomme_stub()
-
-
 def test_convert_obs_list_format():
     """_convert_obs takes the last element from list-format obs fields and
     emits a nested ``pixels`` dict (image, wrist_image) plus ``agent_pos``.

tests/training/test_gradient_accumulation.py

@@ -1,84 +0,0 @@
-#!/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 pytest
-import torch
-from accelerate import Accelerator
-from torch import nn
-
-from lerobot.scripts.lerobot_train import update_policy
-from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
-
-
-class TinyPolicy(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.projection = nn.Linear(2, 1, bias=False)
-
-    def forward(self, batch):
-        loss = self.projection(batch["x"]).square().mean()
-        return loss, {}
-
-
-def test_gradient_accumulation_steps_optimizer_and_scheduler_once():
-    accelerator = Accelerator(
-        cpu=True,
-        gradient_accumulation_steps=2,
-        step_scheduler_with_optimizer=False,
-    )
-    policy = TinyPolicy()
-    optimizer = torch.optim.SGD(policy.parameters(), lr=0.1)
-    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.5)
-    policy, optimizer, scheduler = accelerator.prepare(policy, optimizer, scheduler)
-    metrics = {
-        "loss": AverageMeter("loss"),
-        "grad_norm": AverageMeter("grad_norm"),
-        "lr": AverageMeter("lr"),
-        "update_s": AverageMeter("update_s"),
-    }
-    tracker = MetricsTracker(1, 2, 1, metrics, accelerator=accelerator)
-    batch = {"x": torch.ones(1, 2)}
-    before = policy.projection.weight.detach().clone()
-
-    with accelerator.accumulate(policy):
-        update_policy(
-            tracker,
-            policy,
-            batch,
-            optimizer,
-            grad_clip_norm=0,
-            accelerator=accelerator,
-            lr_scheduler=scheduler,
-            log_metrics=False,
-        )
-    after_first_microbatch = policy.projection.weight.detach().clone()
-
-    with accelerator.accumulate(policy):
-        update_policy(
-            tracker,
-            policy,
-            batch,
-            optimizer,
-            grad_clip_norm=0,
-            accelerator=accelerator,
-            lr_scheduler=scheduler,
-            log_metrics=False,
-        )
-    after_optimizer_step = policy.projection.weight.detach().clone()
-
-    torch.testing.assert_close(after_first_microbatch, before)
-    assert not torch.equal(after_optimizer_step, after_first_microbatch)
-    assert optimizer.param_groups[0]["lr"] == pytest.approx(0.05)