"""Thin wrapper around HuggingFace-format FAST action tokenizer. FAST (Pertsch et al., Physical Intelligence) tokenizes continuous action chunks via DCT + quantization + BPE. Input must be in [-1, 1] range (post-normalization). The default asset (``physical-intelligence/fast``) is configured via --fast_tokenizer_path / LABVLA_FAST_TOKENIZER_PATH; the wrapper itself takes the path explicitly. The underlying processor is a UniversalActionProcessor (HuggingFace ProcessorMixin). """ from __future__ import annotations from pathlib import Path import numpy as np from transformers import AutoProcessor class FastTokenizerWrapper: """Load once, encode many. Stateless after load; safe for DataLoader workers. Encode path: actions (K, D) in [-1, 1] -> list[int] tokens -> zero-padded np.int64 (max_length,) -> mask (max_length,) bool """ # Tolerance for the out-of-range diagnostic in encode(). See encode()'s # docstring for why FAST's required [-1, 1] clip is asymmetric with the # continuous MSE branch and why we only warn rather than clip the MSE target. _OOR_TOLERANCE: float = 1e-3 def __init__(self, path: str, vocab_size: int, max_length: int): self.path = Path(path) self.vocab_size = int(vocab_size) self.max_length = int(max_length) # trust_remote_code needed because the processor class is defined locally # in processing_action_tokenizer.py (not in transformers upstream). self._processor = AutoProcessor.from_pretrained( str(self.path), trust_remote_code=True, ) # Fail loud at LOAD time when the processor's own vocab differs from the # configured one. The KI head's nn.Embedding is sized by # --discrete_action_vocab_size; a processor emitting ids >= that size # would crash later as an unreadable device-side CUDA assert. _proc_vocab = getattr(self._processor, "vocab_size", None) if _proc_vocab is not None and int(_proc_vocab) != self.vocab_size: raise ValueError( f"FastTokenizerWrapper: processor at {self.path} declares " f"vocab_size={int(_proc_vocab)} but the configured " f"--discrete_action_vocab_size is {self.vocab_size}. The KI " f"head embedding/classifier are sized by the configured value; " f"mismatched ids would index out of bounds (CUDA assert). " f"Align the config with the tokenizer asset." ) self._truncated = 0 # count how many encode() calls hit max_length self._out_of_range = 0 # count encode() calls with |action| > 1+tol def encode(self, actions: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """Encode a single (K, D) action chunk. Args: actions: (K, D) normalized to [-1, 1], float32/float64 ok Returns: tokens: (max_length,) int64, zero-padded mask: (max_length,) bool, True for valid positions Note: FAST's DCT+BPE pipeline requires inputs in [-1, 1], so this method clips `actions` before tokenizing (intentional, part of the FAST contract). This is ASYMMETRIC with the continuous flow-matching MSE branch, which trains on the UNCLIPPED normalized actions — so any element with |action| > 1 gets a clipped discrete CE target but an unclipped MSE target. We warn-once (no behavioural change) past a small tolerance, since it usually means stats drift or outliers, but deliberately do NOT clip the continuous MSE target (that would discard precision). """ assert actions.ndim == 2, f"expected (K,D), got {actions.shape}" a_f32 = actions.astype(np.float32) # Diagnostic: detect normalized actions outside the FAST-required # [-1, 1] range (signals stats drift / outliers / norm mismatch). max_abs = float(np.abs(a_f32).max()) if a_f32.size else 0.0 if max_abs > 1.0 + self._OOR_TOLERANCE: self._out_of_range += 1 if self._out_of_range == 1 or self._out_of_range % 5000 == 0: import logging logging.warning( "[FastTokenizerWrapper] %d action chunks had normalized " "values outside [-1, 1] (most recent max|a|=%.4f) and were " "clipped for FAST tokenization. The continuous flow-matching " "MSE branch trains on the UNCLIPPED values, so the discrete " "(FAST CE) and continuous (MSE) targets diverge for those " "elements. This usually indicates normalization-stats drift " "or outliers — re-check the dataset stats / canonicalization.", self._out_of_range, max_abs, ) a = np.clip(a_f32, -1.0, 1.0) # UniversalActionProcessor.__call__ handles (K,D) -> wraps to (1,K,D) internally # and returns list[list[int]] of length 1 batched = self._processor(a) token_ids = batched[0] if isinstance(batched, list) and len(batched) > 0 else [] tokens = np.zeros(self.max_length, dtype=np.int64) mask = np.zeros(self.max_length, dtype=bool) raw_len = len(token_ids) L = min(raw_len, self.max_length) if raw_len > self.max_length: self._truncated += 1 # Sample-log on 1st and every 5000th truncation to avoid flooding but # still surface systematic truncation without needing a separate counter # report. if self._truncated == 1 or self._truncated % 5000 == 0: import logging logging.warning( "[FastTokenizerWrapper] %d samples truncated to max_length=%d " "(most recent raw_len=%d). If this count grows, consider " "raising max_length.", self._truncated, self.max_length, raw_len, ) if L > 0: tokens[:L] = np.asarray(token_ids[:L], dtype=np.int64) mask[:L] = True # Per-encode id-range guard (cheap max over <=max_length ints). # Catches any id escaping the load-time vocab check (e.g. a processor # without a vocab_size attribute) BEFORE it reaches the KI embedding # as an opaque device-side assert. _mx = int(tokens[:L].max()) _mn = int(tokens[:L].min()) if _mx >= self.vocab_size or _mn < 0: raise ValueError( f"FastTokenizerWrapper.encode: token id range [{_mn}, {_mx}] " f"outside [0, {self.vocab_size}) — the FAST processor at " f"{self.path} emits ids incompatible with " f"--discrete_action_vocab_size={self.vocab_size}." ) return tokens, mask