seeker
383 lines
15 KiB
Python
383 lines
15 KiB
Python
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)
|