Spaces:
Sleeping
Sleeping
File size: 6,951 Bytes
34f3bc9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """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
|