seeker
389 lines
14 KiB
Python
389 lines
14 KiB
Python
"""Prompt and FAST-lite tokenization for the minimal pi0.7 policy.
|
|
|
|
pi0.7 conditions the low-level policy on a richer prompt than earlier pi
|
|
models: task and subtask language, episode speed/quality/mistake metadata,
|
|
and control mode. This tokenizer represents those concepts
|
|
with explicit marker/value pairs so the architecture can be validated without
|
|
shipping a full language tokenizer.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import Iterable, Literal
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
ControlMode = Literal["joint", "ee"]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EpisodeMetadata:
|
|
speed_steps: int | None = None
|
|
quality: int | None = None
|
|
mistake: bool | None = None
|
|
control_mode: ControlMode = "joint"
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.speed_steps is not None and self.speed_steps <= 0:
|
|
raise ValueError("speed_steps must be positive")
|
|
if self.quality is not None and not 1 <= self.quality <= 5:
|
|
raise ValueError("quality must be in the inclusive range 1..5")
|
|
if self.control_mode not in ("joint", "ee"):
|
|
raise ValueError("control_mode must be 'joint' or 'ee'")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PolicySequence:
|
|
input_ids: Tensor
|
|
labels: Tensor
|
|
action_prefix_mask: Tensor
|
|
generated_token_start: int
|
|
action_start_index: int
|
|
prompt_length: int
|
|
metadata_indices: tuple[int, ...]
|
|
|
|
|
|
class _LabelVocabulary:
|
|
def __init__(self, labels: tuple[str, ...], first_token_id: int) -> None:
|
|
self.labels = labels
|
|
self.first_token_id = first_token_id
|
|
self._to_token = {
|
|
label: first_token_id + index for index, label in enumerate(labels)
|
|
}
|
|
self._to_label = {token: label for label, token in self._to_token.items()}
|
|
|
|
@property
|
|
def token_ids(self) -> tuple[int, ...]:
|
|
return tuple(self._to_token[label] for label in self.labels)
|
|
|
|
def encode(self, label: str) -> int:
|
|
try:
|
|
return self._to_token[label]
|
|
except KeyError as error:
|
|
raise ValueError(f"unknown categorical label: {label}") from error
|
|
|
|
def decode(self, token: int | Tensor) -> str:
|
|
try:
|
|
return self._to_label[int(token)]
|
|
except KeyError as error:
|
|
raise ValueError(f"unknown categorical token: {int(token)}") from error
|
|
|
|
|
|
class MiniTokenizer:
|
|
"""Explicit rich-prompt tokens plus a DCT-based FAST-lite action stream."""
|
|
|
|
BOS = 256
|
|
TASK_MARKER = 257
|
|
SUBTASK_MARKER = 258
|
|
SPEED_MARKER = 259
|
|
QUALITY_MARKER = 260
|
|
MISTAKE_MARKER = 261
|
|
CONTROL_MARKER = 262
|
|
NEXT_SUBTASK_MARKER = 263
|
|
|
|
TASK_LABELS = (
|
|
"drive",
|
|
"pick_place",
|
|
"fold",
|
|
"open_close",
|
|
"clean",
|
|
"cook",
|
|
"unknown",
|
|
)
|
|
TASK_TOKEN_START = 264
|
|
TASK_TOKEN_END = TASK_TOKEN_START + len(TASK_LABELS) - 1
|
|
|
|
SUBTASK_LABELS = (
|
|
"keep_lane",
|
|
"steer_left",
|
|
"steer_right",
|
|
"slow_down",
|
|
"stop",
|
|
"accelerate",
|
|
"pick_place",
|
|
"fold",
|
|
"open_close",
|
|
"unknown",
|
|
)
|
|
SUBTASK_TOKEN_START = TASK_TOKEN_END + 1
|
|
SUBTASK_TOKEN_END = SUBTASK_TOKEN_START + len(SUBTASK_LABELS) - 1
|
|
|
|
SPEED_BINS = (500, 1000, 2000, 4000, 8000, 16000)
|
|
SPEED_TOKEN_START = SUBTASK_TOKEN_END + 1
|
|
SPEED_TOKEN_END = SPEED_TOKEN_START + len(SPEED_BINS) - 1
|
|
QUALITY_TOKEN_START = SPEED_TOKEN_END + 1
|
|
QUALITY_TOKEN_END = QUALITY_TOKEN_START + 4
|
|
MISTAKE_FALSE = QUALITY_TOKEN_END + 1
|
|
MISTAKE_TRUE = MISTAKE_FALSE + 1
|
|
CONTROL_JOINT = MISTAKE_TRUE + 1
|
|
CONTROL_EE = CONTROL_JOINT + 1
|
|
|
|
ACTION_START = CONTROL_EE + 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) -> None:
|
|
self.tasks = _LabelVocabulary(self.TASK_LABELS, self.TASK_TOKEN_START)
|
|
self.subtasks = _LabelVocabulary(self.SUBTASK_LABELS, self.SUBTASK_TOKEN_START)
|
|
|
|
@property
|
|
def vocab_size(self) -> int:
|
|
return self.ACTION_TOKEN_END + 1
|
|
|
|
@property
|
|
def action_token_ids(self) -> range:
|
|
return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)
|
|
|
|
@staticmethod
|
|
def _classify_task(description: str) -> str:
|
|
text = description.casefold()
|
|
if any(word in text for word in ("drive", "lane", "road", "steer", "vehicle")):
|
|
return "drive"
|
|
if any(word in text for word in ("fold", "shirt", "towel", "laundry")):
|
|
return "fold"
|
|
if any(word in text for word in ("open", "close", "door", "drawer", "fridge")):
|
|
return "open_close"
|
|
if any(word in text for word in ("clean", "wipe", "wash", "tidy")):
|
|
return "clean"
|
|
if any(word in text for word in ("cook", "food", "recipe", "microwave")):
|
|
return "cook"
|
|
if any(word in text for word in ("pick", "place", "move", "put")):
|
|
return "pick_place"
|
|
return description if description in MiniTokenizer.TASK_LABELS else "unknown"
|
|
|
|
@staticmethod
|
|
def _classify_subtask(description: str) -> str:
|
|
if description in MiniTokenizer.SUBTASK_LABELS:
|
|
return description
|
|
text = description.casefold()
|
|
if any(word in text for word in ("stop", "halt", "stationary")):
|
|
return "stop"
|
|
if any(word in text for word in ("slow", "brake", "deceler")):
|
|
return "slow_down"
|
|
if any(word in text for word in ("accelerat", "speed up")):
|
|
return "accelerate"
|
|
if "left" in text:
|
|
return "steer_left"
|
|
if "right" in text:
|
|
return "steer_right"
|
|
if any(word in text for word in ("lane", "straight", "continue", "follow")):
|
|
return "keep_lane"
|
|
if "fold" in text:
|
|
return "fold"
|
|
if any(word in text for word in ("open", "close")):
|
|
return "open_close"
|
|
if any(word in text for word in ("pick", "place", "put")):
|
|
return "pick_place"
|
|
return "unknown"
|
|
|
|
def encode_prompt(
|
|
self,
|
|
task: str,
|
|
subtask: str | None,
|
|
metadata: EpisodeMetadata,
|
|
*,
|
|
include_episode_metadata: bool = True,
|
|
include_subtask: bool = True,
|
|
) -> tuple[Tensor, tuple[int, ...]]:
|
|
"""Encode prompt and return positions controlled by metadata CFG.
|
|
|
|
Control mode is deliberately retained even in the unconditional branch,
|
|
matching pi0.7 training where control-mode dropout is not applied.
|
|
"""
|
|
|
|
ids = [self.BOS, self.TASK_MARKER, self.tasks.encode(self._classify_task(task))]
|
|
if include_subtask and subtask is not None:
|
|
ids.extend(
|
|
(self.SUBTASK_MARKER, self.subtasks.encode(self._classify_subtask(subtask)))
|
|
)
|
|
|
|
metadata_indices: list[int] = []
|
|
|
|
def add_metadata(marker: int, value: int) -> None:
|
|
metadata_indices.extend((len(ids), len(ids) + 1))
|
|
ids.extend((marker, value))
|
|
|
|
if include_episode_metadata:
|
|
if metadata.speed_steps is not None:
|
|
bin_index = min(
|
|
range(len(self.SPEED_BINS)),
|
|
key=lambda index: abs(self.SPEED_BINS[index] - metadata.speed_steps),
|
|
)
|
|
add_metadata(self.SPEED_MARKER, self.SPEED_TOKEN_START + bin_index)
|
|
if metadata.quality is not None:
|
|
add_metadata(
|
|
self.QUALITY_MARKER, self.QUALITY_TOKEN_START + metadata.quality - 1
|
|
)
|
|
if metadata.mistake is not None:
|
|
add_metadata(
|
|
self.MISTAKE_MARKER,
|
|
self.MISTAKE_TRUE if metadata.mistake else self.MISTAKE_FALSE,
|
|
)
|
|
control = self.CONTROL_JOINT if metadata.control_mode == "joint" else self.CONTROL_EE
|
|
ids.extend((self.CONTROL_MARKER, control))
|
|
return torch.tensor(ids, dtype=torch.long), tuple(metadata_indices)
|
|
|
|
def encode_high_level_prompt(
|
|
self,
|
|
task: str,
|
|
past_subtasks: Iterable[str],
|
|
metadata: EpisodeMetadata,
|
|
) -> Tensor:
|
|
"""Encode task, past semantic decisions, metadata, and a next-subtask query."""
|
|
|
|
prompt, _ = self.encode_prompt(
|
|
task,
|
|
subtask=None,
|
|
metadata=metadata,
|
|
include_subtask=False,
|
|
)
|
|
control_pair = prompt[-2:]
|
|
body = prompt[:-2]
|
|
history: list[int] = []
|
|
for subtask in past_subtasks:
|
|
history.extend(
|
|
(self.SUBTASK_MARKER, self.subtasks.encode(self._classify_subtask(subtask)))
|
|
)
|
|
return torch.cat(
|
|
(
|
|
body,
|
|
torch.tensor(history, dtype=torch.long),
|
|
control_pair,
|
|
torch.tensor([self.NEXT_SUBTASK_MARKER], dtype=torch.long),
|
|
)
|
|
)
|
|
|
|
@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 FAST-lite 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,
|
|
task: str,
|
|
subtask: str | None,
|
|
actions: Tensor,
|
|
metadata: EpisodeMetadata,
|
|
*,
|
|
include_episode_metadata: bool = True,
|
|
include_subtask: bool = True,
|
|
) -> PolicySequence:
|
|
prompt, metadata_indices = self.encode_prompt(
|
|
task,
|
|
subtask,
|
|
metadata,
|
|
include_episode_metadata=include_episode_metadata,
|
|
include_subtask=include_subtask,
|
|
)
|
|
action_tokens = self.encode_actions(actions).cpu()
|
|
input_ids = torch.cat((prompt, action_tokens))
|
|
action_start_index = prompt.numel()
|
|
labels = torch.full_like(input_ids, -100)
|
|
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,
|
|
action_prefix_mask=action_prefix_mask,
|
|
generated_token_start=0,
|
|
action_start_index=action_start_index,
|
|
prompt_length=prompt.numel(),
|
|
metadata_indices=metadata_indices,
|
|
)
|
|
|
|
def build_metadata_cfg_pair(
|
|
self,
|
|
task: str,
|
|
subtask: str | None,
|
|
actions: Tensor,
|
|
metadata: EpisodeMetadata,
|
|
) -> tuple[PolicySequence, PolicySequence]:
|
|
"""Return conditioned and episode-metadata-dropped policy sequences."""
|
|
|
|
conditioned = self.build_policy_sequence(task, subtask, actions, metadata)
|
|
unconditional = self.build_policy_sequence(
|
|
task,
|
|
subtask,
|
|
actions,
|
|
metadata,
|
|
include_episode_metadata=False,
|
|
)
|
|
return conditioned, unconditional
|