from __future__ import annotations import math from dataclasses import dataclass from typing import Iterable import torch @dataclass(frozen=True) class CoordinateRange: minimum: float maximum: float def __post_init__(self) -> None: if self.maximum <= self.minimum: raise ValueError("coordinate maximum must be greater than minimum") class CoordinateTokenizer: """Uniformly discretize XY state values into the literal token range 0..255.""" bins = 256 def __init__( self, x_range: tuple[float, float] = (-10.0, 10.0), y_range: tuple[float, float] = (-5.0, 5.0), ) -> None: self.ranges = ( CoordinateRange(*x_range), CoordinateRange(*y_range), ) def encode(self, coordinates: torch.Tensor | Iterable[Iterable[float]]) -> torch.Tensor: values = torch.as_tensor(coordinates, dtype=torch.float32) if values.shape[-1] != 2: raise ValueError("coordinates must end in an XY dimension of size 2") low = values.new_tensor([axis.minimum for axis in self.ranges]) high = values.new_tensor([axis.maximum for axis in self.ranges]) unit = ((values - low) / (high - low)).clamp(0.0, 1.0) return torch.round(unit * (self.bins - 1)).to(torch.long) def decode(self, tokens: torch.Tensor | Iterable[Iterable[int]]) -> torch.Tensor: encoded = torch.as_tensor(tokens, dtype=torch.long) if encoded.shape[-1] != 2: raise ValueError("coordinate tokens must end in an XY dimension of size 2") if encoded.numel() and (encoded.min() < 0 or encoded.max() >= self.bins): raise ValueError("coordinate tokens must be in the inclusive range 0..255") low = torch.tensor([axis.minimum for axis in self.ranges], dtype=torch.float32) high = torch.tensor([axis.maximum for axis in self.ranges], dtype=torch.float32) return low + encoded.float() / (self.bins - 1) * (high - low) class TaskVocabulary: """Small overall-task vocabulary for the high-level pi0.5 inference pass.""" labels = ( "follow_road", "clean_kitchen", "clean_bedroom", "pick_place", "fold_laundry", "open_close", "unknown", ) def __init__(self, first_token_id: int = 260) -> None: self.first_token_id = int(first_token_id) self._label_to_token = { label: self.first_token_id + index for index, label in enumerate(self.labels) } self._token_to_label = {token: label for label, token in self._label_to_token.items()} @property def token_ids(self) -> tuple[int, ...]: return tuple(self._label_to_token[label] for label in self.labels) def classify(self, instruction: str) -> str: text = instruction.casefold() if any(term in text for term in ("road", "lane", "drive", "waypoint", "steer")): return "follow_road" if any(term in text for term in ("fold", "folding")): return "fold_laundry" if any(term in text for term in ("kitchen", "dish", "sink", "counter")): return "clean_kitchen" if any(term in text for term in ("bedroom", "bed", "pillow", "blanket")): return "clean_bedroom" if any(term in text for term in ("open", "close", "drawer", "cabinet", "fridge")): return "open_close" if any(term in text for term in ("pick", "place", "put", "move")): return "pick_place" return instruction if instruction in self._label_to_token else "unknown" def encode(self, label_or_instruction: str) -> int: label = ( label_or_instruction if label_or_instruction in self._label_to_token else self.classify(label_or_instruction) ) return self._label_to_token[label] def decode(self, token_id: int | torch.Tensor) -> str: try: return self._token_to_label[int(token_id)] except KeyError as error: raise ValueError(f"unknown task token: {int(token_id)}") from error class SubtaskVocabulary: """Small semantic vocabulary standing in for pi0.5's generated language subtask.""" labels = ( "keep_lane", "steer_left", "steer_right", "slow_down", "stop", "accelerate", "unknown", ) def __init__(self, first_token_id: int = 260) -> None: self.first_token_id = int(first_token_id) self._label_to_token = { label: self.first_token_id + index for index, label in enumerate(self.labels) } self._token_to_label = {token: label for label, token in self._label_to_token.items()} @property def size(self) -> int: return len(self.labels) @property def token_ids(self) -> tuple[int, ...]: return tuple(self._label_to_token[label] for label in self.labels) def classify(self, reasoning: str) -> str: text = reasoning.casefold() # More specific commands precede generic lane/follow language. if any(term in text for term in ("stop", "halt", "stationary")): return "stop" if any(term in text for term in ("slow", "deceler", "reduce speed", "brake")): return "slow_down" if any(term in text for term in ("accelerat", "speed up")): return "accelerate" if any(term in text for term in ("left", "counterclockwise")): return "steer_left" if any(term in text for term in ("right", "clockwise")): return "steer_right" if any(term in text for term in ("lane", "straight", "follow", "continue", "proceed")): return "keep_lane" return "unknown" def encode(self, label_or_reasoning: str) -> int: label = ( label_or_reasoning if label_or_reasoning in self._label_to_token else self.classify(label_or_reasoning) ) return self._label_to_token[label] def decode(self, token_id: int | torch.Tensor) -> str: token = int(token_id) try: return self._token_to_label[token] except KeyError as error: raise ValueError(f"unknown subtask token: {token}") from error class MiniTokenizer: """Compose hierarchical task/state, subtask, and FAST-lite action tokens. pi0.5 performs two semantically different calls with the same model. The high-level call reads the overall task and predicts a semantic subtask. The low-level call deliberately replaces the overall task with that predicted subtask, matching ``pi(a | observation, predicted_subtask)`` rather than allowing continuous actions to depend directly on the overall task. Actions are z-score-normalized by the dataset before reaching this tokenizer. A complete ``[horizon, 2]`` chunk is transformed along time with an orthonormal DCT-II. All coefficients are retained and serialized in frequency-major XY order, so low-frequency motion appears first in the token stream. The coefficients are divided by ``sqrt(horizon)`` before being uniformly quantized over ``[-4, 4]`` into 64 shared tokens. This keeps the fixed quantization interval valid for any horizon when the input actions have first been clipped to the same interval. This intentionally implements only the central transform/order idea from FAST. It does not learn dataset-specific coefficient bins or drop trailing high-frequency coefficients. """ DEFAULT_TASK_INSTRUCTION = "Follow the visible road." BOS = 256 IMAGE = 257 STATE = 258 TASK = 259 TASK_TOKEN_START = TASK + 1 TASK_TOKEN_END = TASK_TOKEN_START + len(TaskVocabulary.labels) - 1 SUBTASK_QUERY = TASK_TOKEN_END + 1 SUBTASK_COMMAND = SUBTASK_QUERY + 1 SUBTASK_TOKEN_START = SUBTASK_COMMAND + 1 SUBTASK_TOKEN_END = SUBTASK_TOKEN_START + len(SubtaskVocabulary.labels) - 1 ACTION_START = SUBTASK_TOKEN_END + 1 ACTION_END = ACTION_START + 1 ACTION_TOKEN_START = ACTION_END + 1 ACTION_BINS = 64 ACTION_TOKEN_END = ACTION_TOKEN_START + ACTION_BINS - 1 ACTION_MINIMUM = -4.0 ACTION_MAXIMUM = 4.0 def __init__( self, coordinate_tokenizer: CoordinateTokenizer | None = None, tasks: TaskVocabulary | None = None, subtasks: SubtaskVocabulary | None = None, ) -> None: self.coordinates = coordinate_tokenizer or CoordinateTokenizer() self.tasks = tasks or TaskVocabulary(first_token_id=self.TASK_TOKEN_START) self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_TOKEN_START) @property def vocab_size(self) -> int: return max(max(self.tasks.token_ids), max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1 @property def action_token_ids(self) -> range: """Contiguous vocabulary interval reserved for quantized DCT coefficients.""" return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1) def build_observation(self, history_xy: torch.Tensor) -> torch.Tensor: coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1) specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long) return torch.cat((specials, coordinate_tokens.cpu()), dim=0) def build_prefix( self, history_xy: torch.Tensor, task_instruction: str = DEFAULT_TASK_INSTRUCTION, ) -> torch.Tensor: """Build the high-level ``observation + overall task -> subtask`` prefix.""" observation = self.build_observation(history_xy) task = torch.tensor( [self.TASK, self.tasks.encode(task_instruction), self.SUBTASK_QUERY], dtype=torch.long, ) return torch.cat((observation, task), dim=0) def build_low_level_prefix( self, history_xy: torch.Tensor, subtask: str | int | torch.Tensor, ) -> torch.Tensor: """Build ``observation + predicted subtask`` for FAST/flow action inference.""" token = self.subtasks.encode(subtask) if isinstance(subtask, str) else int(subtask) if token not in self.subtasks.token_ids: raise ValueError(f"low-level command is not a subtask token: {token}") command = torch.tensor([self.SUBTASK_COMMAND, token], dtype=torch.long) return torch.cat((self.build_observation(history_xy), command), dim=0) def encode_subtask(self, reasoning: str) -> torch.Tensor: return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long) @staticmethod def _orthonormal_dct_matrix( horizon: int, *, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: """Return the orthonormal DCT-II matrix ``[frequency, time]``.""" if horizon <= 0: raise ValueError("DCT horizon must be positive") time = torch.arange(horizon, device=device, dtype=dtype) frequency = torch.arange(horizon, device=device, dtype=dtype).unsqueeze(1) matrix = torch.cos(math.pi / horizon * (time + 0.5) * frequency) matrix[0] *= math.sqrt(1.0 / horizon) if horizon > 1: matrix[1:] *= math.sqrt(2.0 / horizon) return matrix def dct_actions( self, actions: torch.Tensor | Iterable[Iterable[float]] ) -> torch.Tensor: """Apply an orthonormal DCT-II along an action chunk's time axis.""" values = torch.as_tensor(actions, dtype=torch.float32) if values.ndim != 2 or values.shape[-1] != 2: raise ValueError("actions must have shape [horizon, 2]") if values.shape[0] == 0: raise ValueError("actions must contain at least one timestep") transform = self._orthonormal_dct_matrix( values.shape[0], device=values.device, dtype=values.dtype ) return transform @ values def idct_actions( self, coefficients: torch.Tensor | Iterable[Iterable[float]] ) -> torch.Tensor: """Invert :meth:`dct_actions` using the transpose of its basis.""" values = torch.as_tensor(coefficients, dtype=torch.float32) if values.ndim != 2 or values.shape[-1] != 2: raise ValueError("DCT coefficients must have shape [horizon, 2]") if values.shape[0] == 0: raise ValueError("DCT coefficients must contain at least one frequency") transform = self._orthonormal_dct_matrix( values.shape[0], device=values.device, dtype=values.dtype ) return transform.transpose(0, 1) @ values def encode_actions( self, actions: torch.Tensor | Iterable[Iterable[float]] ) -> torch.Tensor: """DCT-transform and encode normalized actions as a marked token stream.""" values = torch.as_tensor(actions, dtype=torch.float32) if values.ndim != 2 or values.shape[-1] != 2: raise ValueError("actions must have shape [horizon, 2]") if values.shape[0] == 0: raise ValueError("actions must contain at least one timestep") values = values.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM) coefficients = self.dct_actions(values) / math.sqrt(values.shape[0]) # Exact zero is useful for a stable shared center bin. Tiny non-zero # values here are only floating-point residue from the DCT basis. coefficients = torch.where( coefficients.abs() < 1e-6, torch.zeros_like(coefficients), coefficients, ) unit = ( ( coefficients.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM) - self.ACTION_MINIMUM ) / (self.ACTION_MAXIMUM - self.ACTION_MINIMUM) ) quantized = torch.round(unit * (self.ACTION_BINS - 1)).to(torch.long) # Row-major flattening is frequency-major: DC XY, frequency-1 XY, ... action_tokens = quantized.reshape(-1) + self.ACTION_TOKEN_START boundaries = torch.tensor( [self.ACTION_START, self.ACTION_END], dtype=torch.long, device=action_tokens.device, ) return torch.cat((boundaries[:1], action_tokens, boundaries[1:])) def decode_actions(self, tokens: torch.Tensor | Iterable[int]) -> torch.Tensor: """Decode FAST-lite tokens back to normalized ``[horizon, 2]`` actions.""" encoded = torch.as_tensor(tokens, dtype=torch.long) if encoded.ndim != 1: raise ValueError("action tokens must be a one-dimensional sequence") if ( encoded.numel() < 4 or encoded[0] != self.ACTION_START or encoded[-1] != self.ACTION_END ): raise ValueError("action tokens must be enclosed by ACTION_START and ACTION_END") payload = encoded[1:-1] if payload.numel() % 2: raise ValueError("action token payload must contain complete XY pairs") if payload.min() < self.ACTION_TOKEN_START or payload.max() > self.ACTION_TOKEN_END: raise ValueError( "action token payload is outside the reserved action token interval" ) quantized = payload - self.ACTION_TOKEN_START normalized_coefficients = ( quantized.float() / (self.ACTION_BINS - 1) * (self.ACTION_MAXIMUM - self.ACTION_MINIMUM) + self.ACTION_MINIMUM ) horizon = payload.numel() // 2 coefficients = normalized_coefficients.reshape(horizon, 2) * math.sqrt(horizon) return self.idct_actions(coefficients)