feat(pi0-fast): support automatic tokenizer fitting

docs/source/pi0fast.mdx

@@ -109,15 +109,21 @@ lerobot-train \
 
 ### Key Training Parameters
 
-| Parameter                              | Description                                        | Default                         |
-| -------------------------------------- | -------------------------------------------------- | ------------------------------- |
-| `--policy.gradient_checkpointing=true` | Reduces memory usage significantly during training | `false`                         |
-| `--policy.dtype=bfloat16`              | Use mixed precision training for efficiency        | `float32`                       |
-| `--policy.chunk_size`                  | Number of action steps to predict (action horizon) | `50`                            |
-| `--policy.n_action_steps`              | Number of action steps to execute                  | `50`                            |
-| `--policy.max_action_tokens`           | Maximum number of FAST tokens per action chunk     | `256`                           |
-| `--policy.action_tokenizer_name`       | FAST tokenizer to use                              | `lerobot/fast-action-tokenizer` |
-| `--policy.compile_model=true`          | Enable torch.compile for faster training           | `false`                         |
+| Parameter                               | Description                                        | Default                         |
+| --------------------------------------- | -------------------------------------------------- | ------------------------------- |
+| `--policy.gradient_checkpointing=true`  | Reduces memory usage significantly during training | `false`                         |
+| `--policy.dtype=bfloat16`               | Use mixed precision training for efficiency        | `float32`                       |
+| `--policy.chunk_size`                   | Number of action steps to predict (action horizon) | `50`                            |
+| `--policy.n_action_steps`               | Number of action steps to execute                  | `50`                            |
+| `--policy.max_action_tokens`            | Maximum number of FAST tokens per action chunk     | `256`                           |
+| `--policy.action_tokenizer_name`        | FAST tokenizer to use                              | `lerobot/fast-action-tokenizer` |
+| `--policy.auto_fit_fast_tokenizer=true` | Fit and cache a tokenizer for the training dataset | `false`                         |
+| `--policy.compile_model=true`           | Enable torch.compile for faster training           | `false`                         |
+
+Set `--policy.auto_fit_fast_tokenizer=true` to sample action chunks from the
+training dataset and cache a fitted tokenizer under
+`~/.cache/lerobot/fast_tokenizers`. This also works when fine-tuning with
+`--policy.path`; leave it disabled to retain the checkpoint's tokenizer.
 
 ## Inference
 

src/lerobot/policies/factory.py

@@ -408,6 +408,20 @@ def make_pre_post_processors(
         _reconnect_relative_absolute_steps(preprocessor, postprocessor)
         return preprocessor, postprocessor
 
+    if (
+        pretrained_path
+        and getattr(policy_cfg, "type", None) == "pi0_fast"
+        and getattr(policy_cfg, "auto_fit_fast_tokenizer", False)
+        and kwargs.get("dataset_repo_id") is not None
+    ):
+        from .pi0_fast.processor_pi0_fast import make_pi0_fast_pre_post_processors
+
+        return make_pi0_fast_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+            dataset_repo_id=kwargs.get("dataset_repo_id"),
+        )
+
     if pretrained_path:
         if isinstance(policy_cfg, GrootConfig):
             from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained
@@ -509,6 +523,15 @@ def make_pre_post_processors(
             dataset_stats=kwargs.get("dataset_stats"),
         )
 
+    elif policy_cfg.type == "pi0_fast":
+        from .pi0_fast.processor_pi0_fast import make_pi0_fast_pre_post_processors
+
+        processors = make_pi0_fast_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+            dataset_repo_id=kwargs.get("dataset_repo_id"),
+        )
+
     elif policy_cfg.type == "pi052":
         # PI052 must precede PI05 because its config subclasses PI05Config.
         from .pi052.processor_pi052 import make_pi052_pre_post_processors

src/lerobot/policies/pi052/fit_fast_tokenizer.py

@@ -24,6 +24,7 @@ import logging
 import os
 import time
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 
@@ -235,3 +236,21 @@ def fit_fast_tokenizer(
     fitted.save_pretrained(str(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:
+    """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
+
+    try:
+        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,
+        )
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("FAST tokenizer fit failed (%s); using %r instead.", exc, config.action_tokenizer_name)
+        return config.action_tokenizer_name

src/lerobot/policies/pi052/processor_pi052.py

@@ -93,34 +93,11 @@ def make_pi052_pre_post_processors(
 
     # Add FAST action-token supervision only when explicitly enabled.
     if getattr(config, "enable_fast_action_loss", False):
-        # Fit once on this dataset and cache by dataset, base tokenizer, and sample count.
-        action_tokenizer_path = config.action_tokenizer_name
-        if getattr(config, "auto_fit_fast_tokenizer", False) and dataset_repo_id is not None:
-            from .fit_fast_tokenizer import fit_fast_tokenizer  # noqa: PLC0415
-
-            cache_dir = Path(config.fast_tokenizer_cache_dir).expanduser()
-            try:
-                action_tokenizer_path = fit_fast_tokenizer(
-                    dataset_repo_id=dataset_repo_id,
-                    cache_dir=cache_dir,
-                    base_tokenizer_name=config.action_tokenizer_name,
-                    n_samples=config.fast_tokenizer_fit_samples,
-                    chunk_size=config.chunk_size,
-                )
-            except Exception as exc:  # noqa: BLE001
-                import logging  # noqa: PLC0415
-
-                logging.getLogger(__name__).warning(
-                    "FAST tokenizer fit failed (%s) — falling back to "
-                    "the universal base tokenizer %r. Train will still "
-                    "work but compression will be suboptimal.",
-                    exc,
-                    config.action_tokenizer_name,
-                )
+        from .fit_fast_tokenizer import resolve_fast_tokenizer  # noqa: PLC0415
 
         input_steps.append(
             ActionTokenizerProcessorStep(
-                action_tokenizer_name=action_tokenizer_path,
+                action_tokenizer_name=resolve_fast_tokenizer(config, dataset_repo_id),
                 max_action_tokens=config.max_action_tokens,
                 fast_skip_tokens=config.fast_skip_tokens,
                 paligemma_tokenizer_name="google/paligemma-3b-pt-224",

src/lerobot/policies/pi0_fast/configuration_pi0_fast.py

@@ -61,6 +61,9 @@ class PI0FastConfig(PreTrainedConfig):
     tokenizer_max_length: int = 200  # see openpi `__post_init__`
     text_tokenizer_name: str = "google/paligemma-3b-pt-224"
     action_tokenizer_name: str = "lerobot/fast-action-tokenizer"
+    auto_fit_fast_tokenizer: bool = False
+    fast_tokenizer_cache_dir: str = "~/.cache/lerobot/fast_tokenizers"
+    fast_tokenizer_fit_samples: int = 1024
     temperature: float = 0.0
     max_decoding_steps: int = 256
     fast_skip_tokens: int = 128

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -101,6 +101,7 @@ class Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(ProcessorStep):
 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,
 ) -> tuple[
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
@@ -143,6 +144,10 @@ def make_pi0_fast_pre_post_processors(
     # state from the observation but does not change it. NormalizerProcessorStep still runs
     # before Pi0FastPrepareStateAndLanguageTokenizerProcessorStep, so the state tokenizer
     # 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)
+
     input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
         AddBatchDimensionProcessorStep(),
@@ -160,7 +165,7 @@ def make_pi0_fast_pre_post_processors(
             padding="max_length",
         ),
         ActionTokenizerProcessorStep(
-            action_tokenizer_name=config.action_tokenizer_name,
+            action_tokenizer_name=action_tokenizer_path,
             max_action_tokens=config.max_action_tokens,
             fast_skip_tokens=config.fast_skip_tokens,
             paligemma_tokenizer_name=config.text_tokenizer_name,

src/lerobot/scripts/lerobot_train.py

@@ -382,12 +382,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     if cfg.is_reward_model_training:
         processor_kwargs["dataset_meta"] = dataset.meta
 
-    # For pi052 (and any future policy that auto-fits part of its
-    # preprocessing per-dataset), pass the dataset repo id so the
-    # processor factory can locate/refresh dataset-specific artifacts
-    # (e.g. fitted FAST tokenizers per Pertsch et al. 2025 [64],
-    # π0.5 §III.C).
-    if cfg.policy.type == "pi052":
+    # Policies that optionally fit dataset-specific processor artifacts need the repo id.
+    if cfg.policy.type in {"pi0_fast", "pi052"}:
         processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id
 
     if not cfg.is_reward_model_training and processor_pretrained_path is not None:

tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py

@@ -0,0 +1,63 @@
+#!/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 lerobot.policies import factory
+from lerobot.policies.pi0_fast import processor_pi0_fast
+from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
+from lerobot.policies.pi052 import fit_fast_tokenizer as fit_module
+
+
+def test_pi0_fast_resolves_dataset_specific_tokenizer(monkeypatch, tmp_path):
+    config = PI0FastConfig(
+        auto_fit_fast_tokenizer=True,
+        action_tokenizer_name="base-tokenizer",
+        fast_tokenizer_cache_dir=str(tmp_path),
+        fast_tokenizer_fit_samples=17,
+        chunk_size=12,
+        n_action_steps=12,
+    )
+    received = {}
+
+    def fake_fit(**kwargs):
+        received.update(kwargs)
+        return "/cache/fitted-tokenizer"
+
+    monkeypatch.setattr(fit_module, "fit_fast_tokenizer", fake_fit)
+
+    assert fit_module.resolve_fast_tokenizer(config, "user/dataset") == "/cache/fitted-tokenizer"
+    assert received == {
+        "dataset_repo_id": "user/dataset",
+        "cache_dir": tmp_path,
+        "base_tokenizer_name": "base-tokenizer",
+        "n_samples": 17,
+        "chunk_size": 12,
+    }
+
+
+def test_pretrained_pi0_fast_rebuilds_processor_only_during_dataset_fit(monkeypatch):
+    config = PI0FastConfig(auto_fit_fast_tokenizer=True)
+    expected = (object(), object())
+
+    monkeypatch.setattr(processor_pi0_fast, "make_pi0_fast_pre_post_processors", lambda **_: expected)
+
+    assert (
+        factory.make_pre_post_processors(
+            config,
+            pretrained_path="checkpoint",
+            dataset_repo_id="user/dataset",
+        )
+        == expected
+    )