diff --git a/docs/source/pi0fast.mdx b/docs/source/pi0fast.mdx index 15dff8071..3c4fc86fd 100644 --- a/docs/source/pi0fast.mdx +++ b/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 diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 351f2d098..40d839ab1 100644 --- a/src/lerobot/policies/factory.py +++ b/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 diff --git a/src/lerobot/policies/pi052/fit_fast_tokenizer.py b/src/lerobot/policies/pi052/fit_fast_tokenizer.py index e3d103519..f694288c1 100644 --- a/src/lerobot/policies/pi052/fit_fast_tokenizer.py +++ b/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 diff --git a/src/lerobot/policies/pi052/processor_pi052.py b/src/lerobot/policies/pi052/processor_pi052.py index 6a33ff7ce..32bdac27b 100644 --- a/src/lerobot/policies/pi052/processor_pi052.py +++ b/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", diff --git a/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py b/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py index e5c6851f4..492f9b7b9 100644 --- a/src/lerobot/policies/pi0_fast/configuration_pi0_fast.py +++ b/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 diff --git a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py b/src/lerobot/policies/pi0_fast/processor_pi0_fast.py index 60a519786..a29e45dd7 100644 --- a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py +++ b/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, diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index a74a545a5..55846c821 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/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: diff --git a/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py b/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py new file mode 100644 index 000000000..d435c9f09 --- /dev/null +++ b/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 + )