Initial pi0.7 minimal validation
This commit is contained in:
seeker
@@ -0,0 +1,393 @@
|
||||
"""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,
|
||||
control mode, and data provenance. 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"
|
||||
source_type: str | None = None
|
||||
|
||||
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
|
||||
SOURCE_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
|
||||
|
||||
SOURCE_LABELS = (
|
||||
"demonstration",
|
||||
"autonomous",
|
||||
"intervention",
|
||||
"human_video",
|
||||
"web",
|
||||
"unknown",
|
||||
)
|
||||
SOURCE_TOKEN_START = CONTROL_EE + 1
|
||||
SOURCE_TOKEN_END = SOURCE_TOKEN_START + len(SOURCE_LABELS) - 1
|
||||
|
||||
ACTION_START = SOURCE_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) -> None:
|
||||
self.tasks = _LabelVocabulary(self.TASK_LABELS, self.TASK_TOKEN_START)
|
||||
self.subtasks = _LabelVocabulary(self.SUBTASK_LABELS, self.SUBTASK_TOKEN_START)
|
||||
self.sources = _LabelVocabulary(self.SOURCE_LABELS, self.SOURCE_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"
|
||||
|
||||
@staticmethod
|
||||
def _classify_source(source: str) -> str:
|
||||
if source in MiniTokenizer.SOURCE_LABELS:
|
||||
return source
|
||||
text = source.casefold()
|
||||
if any(word in text for word in ("demo", "teleop")):
|
||||
return "demonstration"
|
||||
if any(word in text for word in ("autonomous", "rollout", "policy")):
|
||||
return "autonomous"
|
||||
if any(word in text for word in ("intervention", "correction")):
|
||||
return "intervention"
|
||||
if any(word in text for word in ("human", "ego", "video")):
|
||||
return "human_video"
|
||||
if any(word in text for word in ("web", "internet")):
|
||||
return "web"
|
||||
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,
|
||||
)
|
||||
if metadata.source_type is not None:
|
||||
source = self._classify_source(metadata.source_type)
|
||||
add_metadata(self.SOURCE_MARKER, self.sources.encode(source))
|
||||
|
||||
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)
|
||||
|
||||
@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
|
||||
|
||||
Reference in New Issue
Block a user