Initial pi0.5 minimal validation
This commit is contained in:
seeker
@@ -0,0 +1,285 @@
|
||||
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 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 image/state, semantic, and FAST-lite action tokens.
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
BOS = 256
|
||||
IMAGE = 257
|
||||
STATE = 258
|
||||
SUBTASK_QUERY = 259
|
||||
ACTION_START = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
|
||||
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,
|
||||
subtasks: SubtaskVocabulary | None = None,
|
||||
) -> None:
|
||||
self.coordinates = coordinate_tokenizer or CoordinateTokenizer()
|
||||
self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1)
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return max(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_prefix(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)
|
||||
query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
|
||||
return torch.cat((specials, coordinate_tokens.cpu(), query), 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)
|
||||
Reference in New Issue
Block a user