From 727f98021bd37730c889ec13fc668c63a6fd7341 Mon Sep 17 00:00:00 2001 From: pepijn Date: Fri, 17 Jul 2026 12:23:17 +0000 Subject: [PATCH] fix pi052 FAST training consistency Align tokenizer fitting and loss reduction with the effective training dataset, and fail early when FAST supervision cannot be produced safely. Co-authored-by: Cursor --- src/lerobot/configs/default.py | 6 + src/lerobot/datasets/factory.py | 18 +- src/lerobot/policies/factory.py | 22 +- .../policies/pi052/configuration_pi052.py | 2 + .../policies/pi052/fit_fast_tokenizer.py | 231 +++++++++++++++--- src/lerobot/policies/pi052/modeling_pi052.py | 39 ++- src/lerobot/policies/pi052/processor_pi052.py | 17 +- .../policies/pi0_fast/processor_pi0_fast.py | 14 +- src/lerobot/processor/tokenizer_processor.py | 7 + src/lerobot/scripts/lerobot_train.py | 4 + .../pi052/test_pi052_fast_action_loss.py | 58 ++++- .../pi052/test_pi052_fit_fast_tokenizer.py | 92 +++++++ .../pi0_fast/test_pi0_fast_tokenizer_fit.py | 18 +- tests/processor/test_tokenizer_processor.py | 22 ++ 14 files changed, 500 insertions(+), 50 deletions(-) create mode 100644 tests/policies/pi052/test_pi052_fit_fast_tokenizer.py diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 38991a665..32134ccc9 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -33,6 +33,8 @@ class DatasetConfig: # looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub. root: str | None = None episodes: list[int] | None = None + # Episode indices to drop (e.g. corrupt or heterogeneous ones). Applied on top of `episodes`. + exclude_episodes: list[int] | None = None image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig) revision: str | None = None use_imagenet_stats: bool = True @@ -62,6 +64,10 @@ class DatasetConfig: if len(self.episodes) != len(set(self.episodes)): duplicates = sorted({ep for ep in self.episodes if self.episodes.count(ep) > 1}) raise ValueError(f"Episode indices contain duplicates: {duplicates}") + if self.exclude_episodes is not None and any(ep < 0 for ep in self.exclude_episodes): + raise ValueError( + f"exclude_episodes must be non-negative, got: {[ep for ep in self.exclude_episodes if ep < 0]}" + ) @dataclass diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py index da7b4365a..ad7f207ec 100644 --- a/src/lerobot/datasets/factory.py +++ b/src/lerobot/datasets/factory.py @@ -66,6 +66,17 @@ def resolve_delta_timestamps( return delta_timestamps +def _resolve_episodes( + episodes: list[int] | None, exclude_episodes: list[int] | None, total_episodes: int +) -> list[int] | None: + """Apply an episode exclusion list on top of an optional allowlist.""" + if not exclude_episodes: + return episodes + base = episodes if episodes is not None else list(range(total_episodes)) + excluded = set(exclude_episodes) + return [episode for episode in base if episode not in excluded] + + def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset: """Handles the logic of setting up delta timestamps and image transforms before creating a dataset. @@ -87,11 +98,14 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision ) delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta) + episodes = _resolve_episodes( + cfg.dataset.episodes, cfg.dataset.exclude_episodes, ds_meta.total_episodes + ) if not cfg.dataset.streaming: dataset = LeRobotDataset( cfg.dataset.repo_id, root=cfg.dataset.root, - episodes=cfg.dataset.episodes, + episodes=episodes, delta_timestamps=delta_timestamps, image_transforms=image_transforms, revision=cfg.dataset.revision, @@ -104,7 +118,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas dataset = StreamingLeRobotDataset( cfg.dataset.repo_id, root=cfg.dataset.root, - episodes=cfg.dataset.episodes, + episodes=episodes, delta_timestamps=delta_timestamps, image_transforms=image_transforms, revision=cfg.dataset.revision, diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 6235faad8..315812530 100644 --- a/src/lerobot/policies/factory.py +++ b/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 @@ -322,7 +326,15 @@ def make_pre_post_processors( 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")), + "action_tokenizer_name": resolve_fast_tokenizer( + policy_cfg, + kwargs.get("dataset_repo_id"), + kwargs.get("dataset_root"), + kwargs.get("dataset_stats"), + kwargs.get("dataset_revision"), + kwargs.get("dataset_episodes"), + kwargs.get("dataset_exclude_episodes"), + ), } overrides["action_tokenizer_processor"] = action_tokenizer_override kwargs["preprocessor_overrides"] = overrides @@ -435,6 +447,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 +462,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): diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index 63f0f2808..a0a88e560 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -147,6 +147,8 @@ 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: diff --git a/src/lerobot/policies/pi052/fit_fast_tokenizer.py b/src/lerobot/policies/pi052/fit_fast_tokenizer.py index 4e4501b9e..2098deaed 100644 --- a/src/lerobot/policies/pi052/fit_fast_tokenizer.py +++ b/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,95 @@ 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, ) -> 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, + } + 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 fit_fast_tokenizer( @@ -69,6 +156,14 @@ 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, ) -> str: """Fit a FAST tokenizer on a LeRobot dataset's action distribution. @@ -100,7 +195,20 @@ 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, + ) out_dir = cache_dir / sig if out_dir.exists() and (out_dir / _CACHE_SENTINEL).exists(): @@ -113,8 +221,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() @@ -147,15 +255,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 +283,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,12 +307,17 @@ 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.") + samples_per_episode = max(1, n_samples // 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: break @@ -198,6 +329,8 @@ 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: break @@ -213,6 +346,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,12 +355,7 @@ 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) if not hasattr(base, "fit"): @@ -236,21 +366,54 @@ def fit_fast_tokenizer( ) fitted = base.fit(actions) - out_dir.mkdir(parents=True, exist_ok=True) - fitted.save_pretrained(str(out_dir)) + 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)) + 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, ) diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index bf344f348..9b838a79f 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/src/lerobot/policies/pi052/modeling_pi052.py @@ -362,10 +362,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 +974,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: diff --git a/src/lerobot/policies/pi052/processor_pi052.py b/src/lerobot/policies/pi052/processor_pi052.py index 32bdac27b..8202346e5 100644 --- a/src/lerobot/policies/pi052/processor_pi052.py +++ b/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, ) ) diff --git a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py b/src/lerobot/policies/pi0_fast/processor_pi0_fast.py index a29e45dd7..e4f68f222 100644 --- a/src/lerobot/policies/pi0_fast/processor_pi0_fast.py +++ b/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 diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py index f45cd2cb4..ff4c2d77c 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -350,6 +350,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 +503,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 +573,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 diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index f7a4b033a..4db9ce105 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -352,6 +352,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if cfg.policy.type in {"pi0_fast", "pi052"}: processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id + processor_kwargs["dataset_revision"] = cfg.dataset.revision + processor_kwargs["dataset_episodes"] = cfg.dataset.episodes + processor_kwargs["dataset_exclude_episodes"] = cfg.dataset.exclude_episodes + processor_kwargs["dataset_root"] = cfg.dataset.root if not cfg.is_reward_model_training and processor_pretrained_path is not None: preprocessor_overrides = { diff --git a/tests/policies/pi052/test_pi052_fast_action_loss.py b/tests/policies/pi052/test_pi052_fast_action_loss.py index a32a66597..5250aeb90 100644 --- a/tests/policies/pi052/test_pi052_fast_action_loss.py +++ b/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) diff --git a/tests/policies/pi052/test_pi052_fit_fast_tokenizer.py b/tests/policies/pi052/test_pi052_fit_fast_tokenizer.py new file mode 100644 index 000000000..0f4949172 --- /dev/null +++ b/tests/policies/pi052/test_pi052_fit_fast_tokenizer.py @@ -0,0 +1,92 @@ +#!/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.policies.pi052.fit_fast_tokenizer import ( + _apply_relative_actions, + _dataset_signature, + _is_global_leader, + _normalize_actions, + _select_episode_indices, +) + + +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]]]) diff --git a/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py b/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py index 790bbfe4c..5ec59fa69 100644 --- a/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py +++ b/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): diff --git a/tests/processor/test_tokenizer_processor.py b/tests/processor/test_tokenizer_processor.py index 0290ae78a..aab0c14ec 100644 --- a/tests/processor/test_tokenizer_processor.py +++ b/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."""