Initial pi0.6 minimal validation
This commit is contained in:
seeker
@@ -0,0 +1,330 @@
|
||||
r"""Tiny discrete state, subtask, advantage, and FAST-lite tokenization.
|
||||
|
||||
The layout mirrors the ordering used by :math:`\pi^*_{0.6}`: the model first
|
||||
predicts a subtask, then receives an optional advantage indicator, and only
|
||||
then predicts the discrete action representation. ``omitted`` is represented
|
||||
by the physical absence of an indicator token, as required for unconditional
|
||||
classifier-free guidance training.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, Literal
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
Advantage = Literal["positive", "negative", "omitted"]
|
||||
|
||||
|
||||
@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 literal byte tokens 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: Tensor | Iterable[Iterable[float]]) -> 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)).long()
|
||||
|
||||
def decode(self, tokens: Tensor | Iterable[Iterable[int]]) -> 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])
|
||||
high = torch.tensor([axis.maximum for axis in self.ranges])
|
||||
return low + encoded.float() / (self.bins - 1) * (high - low)
|
||||
|
||||
|
||||
class SubtaskVocabulary:
|
||||
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 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()
|
||||
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 | Tensor) -> str:
|
||||
try:
|
||||
return self._token_to_label[int(token_id)]
|
||||
except KeyError as error:
|
||||
raise ValueError(f"unknown subtask token: {int(token_id)}") from error
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PolicySequence:
|
||||
"""One policy token sequence plus masks needed by ``Pi06Mini``.
|
||||
|
||||
``labels`` is already shifted for next-token prediction. The advantage
|
||||
token itself is never a prediction target because it is externally
|
||||
supplied conditioning. ``action_prefix_mask`` allows the continuous
|
||||
expert to read observation, predicted subtask, and advantage (if present),
|
||||
while insulating it from teacher-forced FAST-lite tokens.
|
||||
"""
|
||||
|
||||
input_ids: Tensor
|
||||
labels: Tensor
|
||||
generated_token_start: int
|
||||
action_prefix_mask: Tensor
|
||||
subtask_index: int
|
||||
advantage_index: int | None
|
||||
action_start_index: int
|
||||
|
||||
|
||||
class MiniTokenizer:
|
||||
r"""Compose the minimal :math:`\pi^*_{0.6}` policy vocabulary."""
|
||||
|
||||
BOS = 256
|
||||
IMAGE = 257
|
||||
STATE = 258
|
||||
SUBTASK_QUERY = 259
|
||||
ADVANTAGE_POSITIVE = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
|
||||
ADVANTAGE_NEGATIVE = ADVANTAGE_POSITIVE + 1
|
||||
ACTION_START = ADVANTAGE_NEGATIVE + 1
|
||||
ACTION_END = ACTION_START + 1
|
||||
ACTION_TOKEN_START = ACTION_END + 1
|
||||
ACTION_BINS = 64
|
||||
ACTION_TOKEN_END = ACTION_TOKEN_START + ACTION_BINS - 1
|
||||
METADATA_DEMO = ACTION_TOKEN_END + 1
|
||||
METADATA_UNDERSTEER = METADATA_DEMO + 1
|
||||
METADATA_STOP = METADATA_DEMO + 2
|
||||
METADATA_INTERVENTION = METADATA_DEMO + 3
|
||||
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 self.METADATA_INTERVENTION + 1
|
||||
|
||||
def encode_metadata(self, source_type: str | None) -> Tensor | None:
|
||||
if source_type is None:
|
||||
return None
|
||||
if source_type == "demo_success": return torch.tensor(self.METADATA_DEMO)
|
||||
if "understeer" in source_type: return torch.tensor(self.METADATA_UNDERSTEER)
|
||||
if "stop" in source_type: return torch.tensor(self.METADATA_STOP)
|
||||
if "intervention" in source_type: return torch.tensor(self.METADATA_INTERVENTION)
|
||||
raise ValueError(f"unknown metadata source_type: {source_type}")
|
||||
|
||||
@property
|
||||
def action_token_ids(self) -> range:
|
||||
return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)
|
||||
|
||||
def build_observation(self, history_xy: Tensor, metadata: str | None = None) -> Tensor:
|
||||
coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1).cpu()
|
||||
specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
|
||||
query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
|
||||
parts = [specials, coordinate_tokens]
|
||||
metadata_token = self.encode_metadata(metadata)
|
||||
if metadata_token is not None: parts.append(metadata_token.reshape(1))
|
||||
parts.append(query)
|
||||
return torch.cat(parts)
|
||||
|
||||
def build_value_input(self, history_xy: Tensor, metadata: str | None = None) -> Tensor:
|
||||
"""Return image marker plus discretized state, without policy query tokens."""
|
||||
|
||||
coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1).cpu()
|
||||
specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
|
||||
parts = [specials, coordinate_tokens]
|
||||
metadata_token = self.encode_metadata(metadata)
|
||||
if metadata_token is not None: parts.append(metadata_token.reshape(1))
|
||||
return torch.cat(parts)
|
||||
|
||||
# Backward-compatible name used by the pi0.5 data utilities.
|
||||
build_prefix = build_observation
|
||||
|
||||
def encode_subtask(self, reasoning: str) -> Tensor:
|
||||
return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long)
|
||||
|
||||
def encode_advantage(self, advantage: Advantage | bool | None) -> Tensor | None:
|
||||
if advantage is True or advantage == "positive":
|
||||
return torch.tensor(self.ADVANTAGE_POSITIVE, dtype=torch.long)
|
||||
if advantage is False or advantage == "negative":
|
||||
return torch.tensor(self.ADVANTAGE_NEGATIVE, dtype=torch.long)
|
||||
if advantage is None or advantage == "omitted":
|
||||
return None
|
||||
raise ValueError("advantage must be positive, negative, omitted, bool, or None")
|
||||
|
||||
@staticmethod
|
||||
def _orthonormal_dct_matrix(
|
||||
horizon: int, *, device: torch.device, dtype: torch.dtype
|
||||
) -> Tensor:
|
||||
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: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||
values = torch.as_tensor(actions, dtype=torch.float32)
|
||||
if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
|
||||
raise ValueError("actions must have non-empty shape [horizon, 2]")
|
||||
transform = self._orthonormal_dct_matrix(
|
||||
values.shape[0], device=values.device, dtype=values.dtype
|
||||
)
|
||||
return transform @ values
|
||||
|
||||
def idct_actions(self, coefficients: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||
values = torch.as_tensor(coefficients, dtype=torch.float32)
|
||||
if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
|
||||
raise ValueError("DCT coefficients must have non-empty shape [horizon, 2]")
|
||||
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: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||
values = torch.as_tensor(actions, dtype=torch.float32)
|
||||
if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
|
||||
raise ValueError("actions must have non-empty shape [horizon, 2]")
|
||||
values = values.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
|
||||
coefficients = self.dct_actions(values) / math.sqrt(values.shape[0])
|
||||
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)
|
||||
payload = torch.round(unit * (self.ACTION_BINS - 1)).long().reshape(-1)
|
||||
payload = payload + self.ACTION_TOKEN_START
|
||||
boundaries = torch.tensor(
|
||||
[self.ACTION_START, self.ACTION_END], dtype=torch.long, device=payload.device
|
||||
)
|
||||
return torch.cat((boundaries[:1], payload, boundaries[1:]))
|
||||
|
||||
def decode_actions(self, tokens: Tensor | Iterable[int]) -> Tensor:
|
||||
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
|
||||
coefficients = (
|
||||
quantized.float() / (self.ACTION_BINS - 1)
|
||||
* (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
|
||||
+ self.ACTION_MINIMUM
|
||||
)
|
||||
horizon = payload.numel() // 2
|
||||
return self.idct_actions(coefficients.reshape(horizon, 2) * math.sqrt(horizon))
|
||||
|
||||
def build_policy_sequence(
|
||||
self,
|
||||
history_xy: Tensor,
|
||||
subtask: str,
|
||||
actions: Tensor,
|
||||
advantage: Advantage | bool | None,
|
||||
metadata: str | None = None,
|
||||
) -> PolicySequence:
|
||||
observation = self.build_observation(history_xy, metadata)
|
||||
subtask_token = self.encode_subtask(subtask).reshape(1)
|
||||
advantage_token = self.encode_advantage(advantage)
|
||||
action_tokens = self.encode_actions(actions).cpu()
|
||||
pieces = [observation, subtask_token]
|
||||
if advantage_token is not None:
|
||||
pieces.append(advantage_token.reshape(1))
|
||||
pieces.append(action_tokens)
|
||||
input_ids = torch.cat(pieces)
|
||||
|
||||
subtask_index = observation.numel()
|
||||
advantage_index = subtask_index + 1 if advantage_token is not None else None
|
||||
action_start_index = subtask_index + 1 + int(advantage_token is not None)
|
||||
|
||||
labels = torch.full_like(input_ids, -100)
|
||||
# SUBTASK_QUERY predicts the subtask. The external advantage indicator
|
||||
# is deliberately skipped as a target. Advantage/subtask then predicts
|
||||
# ACTION_START and the remaining FAST-lite stream autoregressively.
|
||||
labels[subtask_index - 1] = input_ids[subtask_index]
|
||||
labels[action_start_index - 1 : -1] = input_ids[action_start_index:]
|
||||
|
||||
action_prefix_mask = torch.zeros_like(input_ids, dtype=torch.bool)
|
||||
action_prefix_mask[:action_start_index] = True
|
||||
return PolicySequence(
|
||||
input_ids=input_ids,
|
||||
labels=labels,
|
||||
generated_token_start=subtask_index,
|
||||
action_prefix_mask=action_prefix_mask,
|
||||
subtask_index=subtask_index,
|
||||
advantage_index=advantage_index,
|
||||
action_start_index=action_start_index,
|
||||
)
|
||||
Reference in New Issue
Block a user