mirror of
https://github.com/huggingface/lerobot.git
synced 2026-06-04 12:51:27 +00:00
0f4faddc01
Per Pertsch et al. 2025 (FAST paper, [64] in π0.5) and π0.5 §III.C,
the recommended practice is to *fit* the FAST action tokenizer on
the specific dataset's action distribution rather than using the
published universal codebook off the shelf. The universal tokenizer
works on any 6-DoF action sequence but produces suboptimal
compression, which slows CE convergence and wastes vocab capacity.
New utility ``lerobot.policies.pi052.fit_fast_tokenizer``:
* samples N action chunks from the LeRobotDataset (default 1024)
* loads ``physical-intelligence/fast`` as the base
* calls ``.fit(actions)`` (the AutoProcessor API the HF model card
documents) — produces a per-dataset codebook
* saves to ``{cache_dir}/{sha256(dataset, base, n_samples)[:16]}/``
* returns the local path, ready to feed
``ActionTokenizerProcessorStep(action_tokenizer_name=...)``.
Cache is keyed on (dataset, base tokenizer, sample count) so changing
any of them re-runs the fit. Re-running training on the same dataset
re-uses the cache (one fit per dataset per machine).
Auto-fit wiring:
* PI052Config gets ``auto_fit_fast_tokenizer`` (default True),
``fast_tokenizer_cache_dir`` (default ~/.cache/lerobot/...),
``fast_tokenizer_fit_samples`` (default 1024).
* make_pi052_pre_post_processors now takes ``dataset_repo_id``;
when ``enable_fast_action_loss`` and ``auto_fit_fast_tokenizer``
are both True and a repo_id is provided, the factory calls
``fit_fast_tokenizer`` before constructing the processor step
and points it at the fitted path.
* ProcessorConfigKwargs gains ``dataset_repo_id``; the global
factory dispatch threads it through for ``pi052`` policies.
* lerobot_train.py populates ``processor_kwargs['dataset_repo_id']``
from ``--dataset.repo_id`` for pi052 runs.
Failure mode: if ``.fit()`` fails (e.g. older transformers without
the method, or no usable action chunks in the dataset), the factory
logs a warning and falls back to the universal base tokenizer. Train
still works; you just lose the compression improvement.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>