Initial pi0.6 minimal validation
This commit is contained in:
seeker
@@ -0,0 +1,32 @@
|
||||
"""Minimal pi0.6 / Recap architecture validation package."""
|
||||
|
||||
from .model import Pi06Config, Pi06Mini, Pi06Output, flow_matching_sample
|
||||
from .tokenizer import (
|
||||
Advantage,
|
||||
CoordinateTokenizer,
|
||||
MiniTokenizer,
|
||||
PolicySequence,
|
||||
SubtaskVocabulary,
|
||||
)
|
||||
from .value import (
|
||||
DistributionalValueModel,
|
||||
ValueConfig,
|
||||
ValueModel,
|
||||
ValueOutput,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Advantage",
|
||||
"CoordinateTokenizer",
|
||||
"DistributionalValueModel",
|
||||
"MiniTokenizer",
|
||||
"Pi06Config",
|
||||
"Pi06Mini",
|
||||
"Pi06Output",
|
||||
"PolicySequence",
|
||||
"SubtaskVocabulary",
|
||||
"ValueConfig",
|
||||
"ValueModel",
|
||||
"ValueOutput",
|
||||
"flow_matching_sample",
|
||||
]
|
||||
@@ -0,0 +1,354 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Sequence
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image, ImageEnhance
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from .experience import (
|
||||
ExperienceEpisode,
|
||||
build_synthetic_experience,
|
||||
discounted_returns,
|
||||
n_step_advantages,
|
||||
normalize_returns_to_unit_interval,
|
||||
quantile_advantage_labels,
|
||||
stack_episode_field,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_SOURCE_MANIFEST = (
|
||||
Path(__file__).resolve().parents[4]
|
||||
/ "alpamayo"
|
||||
/ "alpamayo-qwen05b-validation"
|
||||
/ "data"
|
||||
/ "processed"
|
||||
/ "overfit_manifest.jsonl"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SampleRecord:
|
||||
clip_id: str
|
||||
t0_us: int
|
||||
image_path: Path
|
||||
trajectory_path: Path
|
||||
reasoning: str = ""
|
||||
|
||||
@property
|
||||
def sample_id(self) -> str:
|
||||
return f"{self.clip_id}_{self.t0_us}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ActionStatistics:
|
||||
mean: torch.Tensor
|
||||
scale: torch.Tensor
|
||||
|
||||
def as_dict(self) -> dict[str, list[float]]:
|
||||
return {
|
||||
"mean": [float(value) for value in self.mean],
|
||||
"scale": [float(value) for value in self.scale],
|
||||
}
|
||||
|
||||
|
||||
class ActionNormalizer:
|
||||
"""Invertibly normalize XY waypoint deltas for flow matching."""
|
||||
|
||||
def __init__(self, mean: torch.Tensor, scale: torch.Tensor) -> None:
|
||||
self.mean = torch.as_tensor(mean, dtype=torch.float32).reshape(2)
|
||||
self.scale = torch.as_tensor(scale, dtype=torch.float32).reshape(2).clamp_min(0.05)
|
||||
|
||||
@staticmethod
|
||||
def to_deltas(history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor:
|
||||
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||
future = torch.as_tensor(future_xy, dtype=torch.float32)
|
||||
if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 1:
|
||||
raise ValueError("history_xy must have shape [history_steps, 2]")
|
||||
if future.ndim != 2 or future.shape[-1] != 2 or len(future) < 1:
|
||||
raise ValueError("future_xy must have shape [action_horizon, 2]")
|
||||
previous = torch.cat((history[-1:], future[:-1]), dim=0)
|
||||
return future - previous
|
||||
|
||||
@classmethod
|
||||
def fit(
|
||||
cls, histories: Sequence[torch.Tensor], futures: Sequence[torch.Tensor]
|
||||
) -> ActionNormalizer:
|
||||
if not histories or len(histories) != len(futures):
|
||||
raise ValueError("histories and futures must be non-empty and equally sized")
|
||||
flat = torch.cat(
|
||||
[
|
||||
cls.to_deltas(history, future).reshape(-1, 2)
|
||||
for history, future in zip(histories, futures, strict=True)
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
return cls(flat.mean(dim=0), flat.std(dim=0, unbiased=False))
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, values: dict[str, Sequence[float]]) -> ActionNormalizer:
|
||||
return cls(torch.tensor(values["mean"]), torch.tensor(values["scale"]))
|
||||
|
||||
@property
|
||||
def statistics(self) -> ActionStatistics:
|
||||
return ActionStatistics(self.mean.clone(), self.scale.clone())
|
||||
|
||||
def encode(self, history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor:
|
||||
return (self.to_deltas(history_xy, future_xy) - self.mean) / self.scale
|
||||
|
||||
def decode(self, actions: torch.Tensor, history_xy: torch.Tensor) -> torch.Tensor:
|
||||
actions = torch.as_tensor(actions, dtype=torch.float32)
|
||||
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||
deltas = actions * self.scale.to(actions.device) + self.mean.to(actions.device)
|
||||
return history[-1:].to(actions.device) + torch.cumsum(deltas, dim=0)
|
||||
|
||||
def as_dict(self) -> dict[str, list[float]]:
|
||||
return self.statistics.as_dict()
|
||||
|
||||
def save(self, path: str | Path) -> Path:
|
||||
destination = Path(path).expanduser().resolve()
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
destination.write_text(json.dumps(self.as_dict(), indent=2) + "\n", encoding="utf-8")
|
||||
return destination
|
||||
|
||||
|
||||
def _resolve_data_path(manifest_path: Path, raw_path: str | Path) -> Path:
|
||||
path = Path(raw_path).expanduser()
|
||||
if path.is_absolute():
|
||||
return path.resolve()
|
||||
candidates = [manifest_path.parent / path]
|
||||
if len(manifest_path.parents) >= 3:
|
||||
candidates.insert(0, manifest_path.parents[2] / path)
|
||||
for candidate in candidates:
|
||||
if candidate.exists():
|
||||
return candidate.resolve()
|
||||
return candidates[0].resolve()
|
||||
|
||||
|
||||
def load_manifest(path: str | Path) -> list[SampleRecord]:
|
||||
manifest_path = Path(path).expanduser().resolve()
|
||||
records: list[SampleRecord] = []
|
||||
for line in manifest_path.read_text(encoding="utf-8").splitlines():
|
||||
if not line.strip():
|
||||
continue
|
||||
row = json.loads(line)
|
||||
records.append(
|
||||
SampleRecord(
|
||||
clip_id=str(row["clip_id"]),
|
||||
t0_us=int(row["t0_us"]),
|
||||
image_path=_resolve_data_path(manifest_path, row["image_path"]),
|
||||
trajectory_path=_resolve_data_path(manifest_path, row["trajectory_path"]),
|
||||
reasoning=str(row.get("reasoning", "")),
|
||||
)
|
||||
)
|
||||
if not records:
|
||||
raise ValueError(f"no samples found in {manifest_path}")
|
||||
return records
|
||||
|
||||
|
||||
def materialize_subset(
|
||||
source_manifest: str | Path = DEFAULT_SOURCE_MANIFEST,
|
||||
output_dir: str | Path = "data/overfit",
|
||||
*,
|
||||
max_samples: int = 1,
|
||||
) -> Path:
|
||||
"""Copy a self-contained PAI seed demonstration for synthetic experience."""
|
||||
|
||||
if max_samples < 1:
|
||||
raise ValueError("max_samples must be positive")
|
||||
destination = Path(output_dir).expanduser().resolve()
|
||||
image_dir = destination / "images"
|
||||
trajectory_dir = destination / "trajectories"
|
||||
image_dir.mkdir(parents=True, exist_ok=True)
|
||||
trajectory_dir.mkdir(parents=True, exist_ok=True)
|
||||
rows: list[dict[str, Any]] = []
|
||||
for index, record in enumerate(load_manifest(source_manifest)[:max_samples]):
|
||||
stem = f"sample_{index:02d}_{record.sample_id}"
|
||||
image_path = image_dir / f"{stem}{record.image_path.suffix.lower()}"
|
||||
trajectory_path = trajectory_dir / f"{stem}.npz"
|
||||
shutil.copy2(record.image_path, image_path)
|
||||
shutil.copy2(record.trajectory_path, trajectory_path)
|
||||
rows.append(
|
||||
{
|
||||
"clip_id": record.clip_id,
|
||||
"t0_us": record.t0_us,
|
||||
"reasoning": record.reasoning,
|
||||
"image_path": str(image_path.relative_to(destination)),
|
||||
"trajectory_path": str(trajectory_path.relative_to(destination)),
|
||||
}
|
||||
)
|
||||
manifest = destination / "overfit_manifest.jsonl"
|
||||
manifest.write_text(
|
||||
"".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return manifest
|
||||
|
||||
|
||||
def _resample_xy(xy: torch.Tensor, steps: int) -> torch.Tensor:
|
||||
xy = torch.as_tensor(xy, dtype=torch.float32)
|
||||
if xy.ndim != 2 or xy.shape[-1] != 2 or len(xy) < 2:
|
||||
raise ValueError("trajectory must contain at least two XY waypoints")
|
||||
if len(xy) == steps:
|
||||
return xy
|
||||
return F.interpolate(
|
||||
xy.T.unsqueeze(0), size=steps, mode="linear", align_corners=True
|
||||
)[0].T
|
||||
|
||||
|
||||
class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
|
||||
"""Four-source ReCAP overfit buffer synthesized from each PAI demonstration."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path: str | Path,
|
||||
*,
|
||||
image_size: int = 64,
|
||||
augmentation_factor: int = 1,
|
||||
history_steps: int = 8,
|
||||
action_horizon: int = 16,
|
||||
n_step: int = 4,
|
||||
gamma: float = 0.99,
|
||||
positive_quantile: float = 0.7,
|
||||
seed: int = 7,
|
||||
action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None,
|
||||
) -> None:
|
||||
if image_size < 1 or augmentation_factor < 1:
|
||||
raise ValueError("image_size and augmentation_factor must be positive")
|
||||
self.records = load_manifest(manifest_path)
|
||||
self.image_size = int(image_size)
|
||||
self.augmentation_factor = int(augmentation_factor)
|
||||
self.history_steps = int(history_steps)
|
||||
self.action_horizon = int(action_horizon)
|
||||
self.seed = int(seed)
|
||||
self.state_dim = self.history_steps * 2
|
||||
self.action_dim = 2
|
||||
|
||||
episodes: list[ExperienceEpisode] = []
|
||||
episode_record_indices: list[int] = []
|
||||
for record_index, record in enumerate(self.records):
|
||||
with np.load(record.trajectory_path) as arrays:
|
||||
history = _resample_xy(
|
||||
torch.from_numpy(np.asarray(arrays["history_xyz"], dtype=np.float32))[:, :2],
|
||||
self.history_steps,
|
||||
)
|
||||
future = _resample_xy(
|
||||
torch.from_numpy(np.asarray(arrays["future_xyz"], dtype=np.float32))[:, :2],
|
||||
self.action_horizon,
|
||||
)
|
||||
generated = build_synthetic_experience(
|
||||
history, future, episode_prefix=record.sample_id
|
||||
)
|
||||
episodes.extend(generated)
|
||||
episode_record_indices.extend([record_index] * len(generated))
|
||||
self.episodes = episodes
|
||||
self._episode_record_indices = episode_record_indices
|
||||
|
||||
raw_rewards = stack_episode_field(episodes, "rewards")
|
||||
raw_returns = discounted_returns(raw_rewards, gamma=gamma)
|
||||
self.normalized_returns = normalize_returns_to_unit_interval(raw_returns)
|
||||
baseline = raw_returns.mean(dim=0, keepdim=True).expand_as(raw_returns)
|
||||
dones = torch.zeros_like(raw_rewards, dtype=torch.bool)
|
||||
dones[:, -1] = True
|
||||
self.advantages = n_step_advantages(
|
||||
raw_rewards, baseline, n_step=n_step, gamma=gamma, dones=dones
|
||||
)
|
||||
intervention_masks = stack_episode_field(episodes, "intervention_mask").bool()
|
||||
self.positive_steps, self.advantage_threshold = quantile_advantage_labels(
|
||||
self.advantages,
|
||||
quantile=positive_quantile,
|
||||
intervention_mask=intervention_masks,
|
||||
)
|
||||
episode_advantages = self.advantages.mean(dim=-1)
|
||||
episode_interventions = intervention_masks.any(dim=-1)
|
||||
self.positive_episodes, self.episode_advantage_threshold = (
|
||||
quantile_advantage_labels(
|
||||
episode_advantages,
|
||||
quantile=positive_quantile,
|
||||
intervention_mask=episode_interventions,
|
||||
)
|
||||
)
|
||||
|
||||
if action_statistics is None:
|
||||
self.action_normalizer = ActionNormalizer.fit(
|
||||
[episode.history_xy for episode in episodes],
|
||||
[episode.future_xy for episode in episodes],
|
||||
)
|
||||
elif isinstance(action_statistics, ActionStatistics):
|
||||
self.action_normalizer = ActionNormalizer(
|
||||
action_statistics.mean, action_statistics.scale
|
||||
)
|
||||
else:
|
||||
self.action_normalizer = ActionNormalizer.from_dict(action_statistics)
|
||||
|
||||
@property
|
||||
def action_statistics(self) -> ActionStatistics:
|
||||
return self.action_normalizer.statistics
|
||||
|
||||
def save_action_statistics(self, path: str | Path) -> Path:
|
||||
return self.action_normalizer.save(path)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.episodes) * self.augmentation_factor
|
||||
|
||||
def _indices(self, index: int) -> tuple[int, int]:
|
||||
if index < 0:
|
||||
index += len(self)
|
||||
if index < 0 or index >= len(self):
|
||||
raise IndexError(index)
|
||||
return divmod(index, self.augmentation_factor)
|
||||
|
||||
def _load_image(self, record_index: int, variant: int) -> torch.Tensor:
|
||||
with Image.open(self.records[record_index].image_path) as source:
|
||||
image = source.convert("RGB")
|
||||
if variant:
|
||||
rng = np.random.default_rng(
|
||||
np.random.SeedSequence([self.seed, record_index, variant])
|
||||
)
|
||||
image = ImageEnhance.Brightness(image).enhance(float(rng.uniform(0.96, 1.04)))
|
||||
image = ImageEnhance.Contrast(image).enhance(float(rng.uniform(0.97, 1.03)))
|
||||
image = image.resize(
|
||||
(self.image_size, self.image_size), resample=Image.Resampling.BILINEAR
|
||||
)
|
||||
array = np.asarray(image, dtype=np.float32).copy() / 255.0
|
||||
return torch.from_numpy(array).permute(2, 0, 1).contiguous()
|
||||
|
||||
def __getitem__(self, index: int) -> dict[str, Any]:
|
||||
episode_index, variant = self._indices(index)
|
||||
episode = self.episodes[episode_index]
|
||||
record_index = self._episode_record_indices[episode_index]
|
||||
actions = self.action_normalizer.encode(
|
||||
episode.history_xy, episode.future_xy
|
||||
)
|
||||
return {
|
||||
"image": self._load_image(record_index, variant),
|
||||
"state": episode.history_xy.reshape(-1),
|
||||
"actions": actions,
|
||||
"history_xy": episode.history_xy.clone(),
|
||||
"future_xy": episode.future_xy.clone(),
|
||||
"rewards": episode.rewards.clone(),
|
||||
"returns": self.normalized_returns[episode_index].clone(),
|
||||
# The critic predicts the beginning-of-episode return for this
|
||||
# minimal trajectory-level buffer; the full vector remains
|
||||
# available for diagnostics and advantage recomputation.
|
||||
"return_target": self.normalized_returns[episode_index, 0].clone(),
|
||||
"value_target": self.normalized_returns[episode_index, 0].clone(),
|
||||
"advantages": self.advantages[episode_index].clone(),
|
||||
"positive_steps": self.positive_steps[episode_index].clone(),
|
||||
"advantage": self.advantages[episode_index].mean().clone(),
|
||||
"positive_condition": self.positive_episodes[episode_index].clone(),
|
||||
"advantage_label": (
|
||||
"positive" if bool(self.positive_episodes[episode_index]) else "negative"
|
||||
),
|
||||
"intervention_mask": episode.intervention_mask.clone(),
|
||||
"is_intervention": bool(episode.intervention_mask.any()),
|
||||
"source_type": episode.source_type,
|
||||
"reasoning": self.records[record_index].reasoning,
|
||||
"sample_id": f"{episode.episode_id}:aug{variant}",
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Literal, Sequence
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
ExperienceSource = Literal[
|
||||
"demo_success",
|
||||
"autonomous_understeer",
|
||||
"autonomous_stop",
|
||||
"intervention_correction",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExperienceEpisode:
|
||||
"""A tiny trajectory-level stand-in for the heterogeneous pi0.6 experience buffer."""
|
||||
|
||||
episode_id: str
|
||||
source_type: ExperienceSource
|
||||
history_xy: torch.Tensor
|
||||
future_xy: torch.Tensor
|
||||
rewards: torch.Tensor
|
||||
intervention_mask: torch.Tensor
|
||||
metadata: dict[str, float | str] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
history = torch.as_tensor(self.history_xy, dtype=torch.float32)
|
||||
future = torch.as_tensor(self.future_xy, dtype=torch.float32)
|
||||
rewards = torch.as_tensor(self.rewards, dtype=torch.float32)
|
||||
intervention = torch.as_tensor(self.intervention_mask, dtype=torch.bool)
|
||||
if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 1:
|
||||
raise ValueError("history_xy must have shape [history_steps, 2]")
|
||||
if future.ndim != 2 or future.shape[-1] != 2 or len(future) < 2:
|
||||
raise ValueError("future_xy must have shape [horizon, 2]")
|
||||
if rewards.shape != (len(future),):
|
||||
raise ValueError("rewards must have one value per future waypoint")
|
||||
if intervention.shape != rewards.shape:
|
||||
raise ValueError("intervention_mask must match rewards")
|
||||
object.__setattr__(self, "history_xy", history)
|
||||
object.__setattr__(self, "future_xy", future)
|
||||
object.__setattr__(self, "rewards", rewards)
|
||||
object.__setattr__(self, "intervention_mask", intervention)
|
||||
|
||||
|
||||
def discounted_returns(rewards: torch.Tensor, gamma: float = 0.99) -> torch.Tensor:
|
||||
"""Compute reward-to-go while preserving arbitrary leading batch dimensions."""
|
||||
|
||||
if not 0.0 <= gamma <= 1.0:
|
||||
raise ValueError("gamma must lie in [0, 1]")
|
||||
rewards = torch.as_tensor(rewards, dtype=torch.float32)
|
||||
if rewards.ndim < 1 or rewards.shape[-1] < 1:
|
||||
raise ValueError("rewards must have a non-empty time dimension")
|
||||
returns = torch.empty_like(rewards)
|
||||
running = torch.zeros_like(rewards[..., 0])
|
||||
for index in range(rewards.shape[-1] - 1, -1, -1):
|
||||
running = rewards[..., index] + gamma * running
|
||||
returns[..., index] = running
|
||||
return returns
|
||||
|
||||
|
||||
def normalize_returns_to_unit_interval(returns: torch.Tensor) -> torch.Tensor:
|
||||
"""Map the worst observed return to -1 and the best to 0.
|
||||
|
||||
pi0.6 uses non-positive reward conditioning. A constant collection is
|
||||
treated as equally optimal and therefore maps to zero.
|
||||
"""
|
||||
|
||||
returns = torch.as_tensor(returns, dtype=torch.float32)
|
||||
if returns.numel() == 0:
|
||||
raise ValueError("returns must be non-empty")
|
||||
minimum = returns.amin()
|
||||
maximum = returns.amax()
|
||||
span = maximum - minimum
|
||||
if float(span) <= torch.finfo(returns.dtype).eps:
|
||||
return torch.zeros_like(returns)
|
||||
return ((returns - maximum) / span).clamp(-1.0, 0.0)
|
||||
|
||||
|
||||
def n_step_advantages(
|
||||
rewards: torch.Tensor,
|
||||
values: torch.Tensor,
|
||||
*,
|
||||
n_step: int = 4,
|
||||
gamma: float = 0.99,
|
||||
dones: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Estimate A(s_t,a_t) from an n-step bootstrapped target.
|
||||
|
||||
``values`` may contain T values (the final bootstrap is zero) or T+1
|
||||
values. ``dones`` stops both reward accumulation and bootstrapping.
|
||||
"""
|
||||
|
||||
rewards = torch.as_tensor(rewards, dtype=torch.float32)
|
||||
values = torch.as_tensor(values, dtype=torch.float32, device=rewards.device)
|
||||
if rewards.ndim < 1 or rewards.shape[-1] < 1:
|
||||
raise ValueError("rewards must have a non-empty time dimension")
|
||||
if n_step < 1:
|
||||
raise ValueError("n_step must be positive")
|
||||
if not 0.0 <= gamma <= 1.0:
|
||||
raise ValueError("gamma must lie in [0, 1]")
|
||||
horizon = rewards.shape[-1]
|
||||
if values.shape[:-1] != rewards.shape[:-1] or values.shape[-1] not in {
|
||||
horizon,
|
||||
horizon + 1,
|
||||
}:
|
||||
raise ValueError("values must have shape [..., T] or [..., T+1]")
|
||||
if values.shape[-1] == horizon:
|
||||
values = torch.cat((values, torch.zeros_like(values[..., :1])), dim=-1)
|
||||
if dones is None:
|
||||
dones = torch.zeros_like(rewards, dtype=torch.bool)
|
||||
dones[..., -1] = True
|
||||
else:
|
||||
dones = torch.as_tensor(dones, dtype=torch.bool, device=rewards.device)
|
||||
if dones.shape != rewards.shape:
|
||||
raise ValueError("dones must match rewards")
|
||||
|
||||
advantages = torch.empty_like(rewards)
|
||||
for start in range(horizon):
|
||||
target = torch.zeros_like(rewards[..., 0])
|
||||
alive = torch.ones_like(target)
|
||||
discount = 1.0
|
||||
stop = min(start + n_step, horizon)
|
||||
for index in range(start, stop):
|
||||
target = target + discount * alive * rewards[..., index]
|
||||
alive = alive * (~dones[..., index]).to(rewards.dtype)
|
||||
discount *= gamma
|
||||
# These inputs are complete episodes rather than replay fragments, so
|
||||
# reaching the final recorded step is terminal even if callers omit a
|
||||
# final done flag. Bootstrap only when the n-step window ends earlier.
|
||||
if stop < horizon:
|
||||
target = target + discount * alive * values[..., stop]
|
||||
advantages[..., start] = target - values[..., start]
|
||||
return advantages
|
||||
|
||||
|
||||
def quantile_advantage_labels(
|
||||
advantages: torch.Tensor,
|
||||
*,
|
||||
quantile: float = 0.7,
|
||||
intervention_mask: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Convert advantages into ReCAP positive/negative conditions.
|
||||
|
||||
Intervention actions are always positive, even when their estimated
|
||||
numerical advantage falls below the empirical threshold.
|
||||
"""
|
||||
|
||||
advantages = torch.as_tensor(advantages, dtype=torch.float32)
|
||||
if advantages.numel() == 0:
|
||||
raise ValueError("advantages must be non-empty")
|
||||
if not 0.0 <= quantile <= 1.0:
|
||||
raise ValueError("quantile must lie in [0, 1]")
|
||||
threshold = torch.quantile(advantages.reshape(-1), quantile)
|
||||
labels = advantages >= threshold
|
||||
if intervention_mask is not None:
|
||||
intervention_mask = torch.as_tensor(
|
||||
intervention_mask, dtype=torch.bool, device=advantages.device
|
||||
)
|
||||
if intervention_mask.shape != advantages.shape:
|
||||
raise ValueError("intervention_mask must match advantages")
|
||||
labels = labels | intervention_mask
|
||||
return labels, threshold
|
||||
|
||||
|
||||
def _trajectory_rewards(candidate: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
error = torch.linalg.vector_norm(candidate - target, dim=-1)
|
||||
scale = torch.linalg.vector_norm(target - target[:1], dim=-1).amax().clamp_min(1.0)
|
||||
return -(error / scale).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def build_synthetic_experience(
|
||||
history_xy: torch.Tensor,
|
||||
expert_future_xy: torch.Tensor,
|
||||
*,
|
||||
episode_prefix: str = "sample",
|
||||
) -> list[ExperienceEpisode]:
|
||||
"""Build four deterministic experience sources from one PAI demonstration.
|
||||
|
||||
The small buffer deliberately includes successful demonstrations, two
|
||||
autonomous failure modes, and a human correction. It is only intended to
|
||||
validate the pi0.6/ReCAP data path and overfit behavior.
|
||||
"""
|
||||
|
||||
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||
expert = torch.as_tensor(expert_future_xy, dtype=torch.float32)
|
||||
if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 1:
|
||||
raise ValueError("history_xy must have shape [history_steps, 2]")
|
||||
if expert.ndim != 2 or expert.shape[-1] != 2 or len(expert) < 4:
|
||||
raise ValueError("expert_future_xy must have shape [horizon>=4, 2]")
|
||||
horizon = len(expert)
|
||||
anchor = history[-1]
|
||||
displacement = expert - anchor
|
||||
|
||||
understeer = anchor + displacement * expert.new_tensor([1.0, 0.2])
|
||||
stop_index = max(2, horizon // 3)
|
||||
stopped = expert.clone()
|
||||
stopped[stop_index:] = stopped[stop_index - 1]
|
||||
|
||||
progress = torch.linspace(0.0, 1.0, horizon, dtype=expert.dtype).unsqueeze(-1)
|
||||
correction_weight = ((progress - 0.2) / 0.55).clamp(0.0, 1.0)
|
||||
correction = understeer * (1.0 - correction_weight) + expert * correction_weight
|
||||
intervention = torch.arange(horizon) >= max(1, horizon // 4)
|
||||
|
||||
zero_rewards = torch.zeros(horizon, dtype=expert.dtype)
|
||||
return [
|
||||
ExperienceEpisode(
|
||||
f"{episode_prefix}:demo",
|
||||
"demo_success",
|
||||
history,
|
||||
expert,
|
||||
zero_rewards,
|
||||
torch.zeros(horizon, dtype=torch.bool),
|
||||
),
|
||||
ExperienceEpisode(
|
||||
f"{episode_prefix}:understeer",
|
||||
"autonomous_understeer",
|
||||
history,
|
||||
understeer,
|
||||
_trajectory_rewards(understeer, expert),
|
||||
torch.zeros(horizon, dtype=torch.bool),
|
||||
{"failure": "understeer"},
|
||||
),
|
||||
ExperienceEpisode(
|
||||
f"{episode_prefix}:stop",
|
||||
"autonomous_stop",
|
||||
history,
|
||||
stopped,
|
||||
_trajectory_rewards(stopped, expert),
|
||||
torch.zeros(horizon, dtype=torch.bool),
|
||||
{"failure": "premature_stop"},
|
||||
),
|
||||
ExperienceEpisode(
|
||||
f"{episode_prefix}:intervention",
|
||||
"intervention_correction",
|
||||
history,
|
||||
correction,
|
||||
_trajectory_rewards(correction, expert),
|
||||
intervention,
|
||||
{"failure": "understeer_then_corrected"},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def stack_episode_field(
|
||||
episodes: Sequence[ExperienceEpisode], field_name: str
|
||||
) -> torch.Tensor:
|
||||
if not episodes:
|
||||
raise ValueError("episodes must be non-empty")
|
||||
return torch.stack([torch.as_tensor(getattr(item, field_name)) for item in episodes])
|
||||
@@ -0,0 +1,44 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class FlowMatching:
|
||||
"""Low-timestep-biased rectified flow with pi*0.6 noise weighting."""
|
||||
|
||||
def __init__(self, inference_steps: int = 5, cutoff: float = 0.999) -> None:
|
||||
self.inference_steps = inference_steps
|
||||
self.cutoff = cutoff
|
||||
|
||||
def training_batch(self, target: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
|
||||
batch = target.shape[0]
|
||||
uniform = torch.rand(batch, device=target.device, dtype=target.dtype)
|
||||
timestep = self.cutoff * (1.0 - uniform.pow(1.0 / 1.5))
|
||||
noise = torch.randn_like(target)
|
||||
noisy = timestep[:, None, None] * target + (1 - timestep[:, None, None]) * noise
|
||||
velocity = target - noise
|
||||
weight = torch.exp(-timestep / 2.0)
|
||||
return noisy, timestep, velocity, weight
|
||||
|
||||
@staticmethod
|
||||
def weighted_loss(prediction: Tensor, target: Tensor, weight: Tensor) -> Tensor:
|
||||
per_example = (prediction.float() - target.float()).square().mean(dim=(1, 2))
|
||||
return (per_example * weight.float()).mean()
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(
|
||||
self,
|
||||
model_fn: Callable[[Tensor, Tensor], Tensor],
|
||||
shape: tuple[int, ...],
|
||||
device: torch.device | str,
|
||||
steps: int | None = None,
|
||||
) -> Tensor:
|
||||
count = steps or self.inference_steps
|
||||
x = torch.randn(shape, device=device)
|
||||
grid = torch.linspace(0.0, 1.0, count + 1, device=device)
|
||||
for index in range(count):
|
||||
t = grid[index].expand(shape[0])
|
||||
x = x + (grid[index + 1] - grid[index]) * model_fn(x, t)
|
||||
return x
|
||||
@@ -0,0 +1,98 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .data import Pi06ExperienceDataset
|
||||
from .flow import FlowMatching
|
||||
from .model import Pi06Config, Pi06Mini
|
||||
from .recap import classifier_free_guidance
|
||||
from .runtime import choose_device, write_json
|
||||
from .tokenizer import MiniTokenizer
|
||||
from .value import ValueConfig, ValueModel
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
|
||||
checkpoint=torch.load(checkpoint_path,map_location="cpu",weights_only=False)
|
||||
cfg=checkpoint["config"]; root=Path(__file__).resolve().parents[2]
|
||||
device=choose_device(cfg.get("device","auto")); tokenizer=MiniTokenizer()
|
||||
dataset=Pi06ExperienceDataset(
|
||||
root/"data/overfit/overfit_manifest.jsonl",image_size=cfg["data"]["image_size"],
|
||||
augmentation_factor=cfg["experience"]["augmentation_factor"],
|
||||
action_horizon=cfg["data"]["action_horizon"],n_step=cfg["experience"]["n_step"],
|
||||
gamma=cfg["experience"]["gamma"],positive_quantile=cfg["experience"]["positive_quantile"],
|
||||
action_statistics=checkpoint["action_statistics"],
|
||||
)
|
||||
mixed=Pi06Mini(Pi06Config(**checkpoint["policy_config"]));mixed.load_state_dict(checkpoint["mixed"]);mixed.to(device).eval()
|
||||
recap=Pi06Mini(Pi06Config(**checkpoint["policy_config"]));recap.load_state_dict(checkpoint["recap"]);recap.to(device).eval()
|
||||
value=ValueModel(ValueConfig(**checkpoint["value_config"]));value.load_state_dict(checkpoint["value"]);value.to(device).eval()
|
||||
flow=FlowMatching(cfg["inference"]["flow_steps"])
|
||||
demo=next(dataset[i] for i in range(len(dataset)) if dataset[i]["source_type"]=="demo_success")
|
||||
failure=next(dataset[i] for i in range(len(dataset)) if "stop" in dataset[i]["source_type"])
|
||||
image=demo["image"][None].to(device)
|
||||
|
||||
def condition(model:Pi06Mini, advantage:str|None, x:torch.Tensor,t:torch.Tensor)->torch.Tensor:
|
||||
seq=tokenizer.build_policy_sequence(
|
||||
demo["history_xy"],demo["reasoning"],demo["actions"],advantage,
|
||||
metadata=None,
|
||||
)
|
||||
ids=seq.input_ids[:seq.action_start_index][None].to(device)
|
||||
return model(image,ids,x,t,generated_token_start=seq.generated_token_start).vector_field
|
||||
|
||||
def sample_field(fn,seed:int=2):
|
||||
torch.manual_seed(seed)
|
||||
return flow.sample(fn,(1,dataset.action_horizon,2),device,cfg["inference"]["flow_steps"])[0].cpu()
|
||||
|
||||
mixed_actions=sample_field(lambda x,t:condition(mixed,None,x,t))
|
||||
uncond_actions=sample_field(lambda x,t:condition(recap,None,x,t))
|
||||
positive_actions=sample_field(lambda x,t:condition(recap,"positive",x,t))
|
||||
negative_actions=sample_field(lambda x,t:condition(recap,"negative",x,t))
|
||||
def cfg_field(x,t):
|
||||
return classifier_free_guidance(
|
||||
condition(recap,"positive",x,t),condition(recap,None,x,t),
|
||||
guidance_scale=cfg["inference"]["cfg_beta"],
|
||||
)
|
||||
cfg_actions=sample_field(cfg_field)
|
||||
decoded={
|
||||
"mixed":dataset.action_normalizer.decode(mixed_actions,demo["history_xy"]).numpy(),
|
||||
"unconditional":dataset.action_normalizer.decode(uncond_actions,demo["history_xy"]).numpy(),
|
||||
"positive":dataset.action_normalizer.decode(positive_actions,demo["history_xy"]).numpy(),
|
||||
"negative":dataset.action_normalizer.decode(negative_actions,demo["history_xy"]).numpy(),
|
||||
"cfg":dataset.action_normalizer.decode(cfg_actions,demo["history_xy"]).numpy(),
|
||||
}
|
||||
target=demo["future_xy"].numpy(); failure_xy=failure["future_xy"].numpy()
|
||||
metrics={name:float(np.mean((xy-target)**2)) for name,xy in decoded.items()}
|
||||
metrics["negative_to_failure_mse"]=float(np.mean((decoded["negative"]-failure_xy)**2))
|
||||
fig,ax=plt.subplots(figsize=(9,6));ax.plot(target[:,0],target[:,1],"o-",label="success target",linewidth=2)
|
||||
ax.plot(failure_xy[:,0],failure_xy[:,1],":",label="failure experience")
|
||||
for name,color in [("mixed","gray"),("unconditional","purple"),("positive","green"),("negative","red"),("cfg","orange")]:
|
||||
xy=decoded[name];ax.plot(xy[:,0],xy[:,1],"x--",label=f"{name} ({metrics[name]:.3f})",color=color)
|
||||
ax.scatter([0],[0],c="black",marker="s",label="ego");ax.axis("equal");ax.grid(True,alpha=.3);ax.legend(fontsize=8)
|
||||
ax.set_xlabel("forward x (m)");ax.set_ylabel("left y (m)");ax.set_title("π*0.6 Recap advantage-conditioned trajectories")
|
||||
outdir=Path(checkpoint_path).parent;trajectory_path=outdir/"recap_trajectory_comparison.png";fig.tight_layout();fig.savefig(trajectory_path,dpi=170);plt.close(fig)
|
||||
|
||||
value_rows=[]
|
||||
for i in range(len(dataset)):
|
||||
s=dataset[i];ids=tokenizer.build_value_input(s["history_xy"],s["source_type"])[None].to(device)
|
||||
prediction=float(value(s["image"][None].to(device),ids).expectation.cpu())
|
||||
value_rows.append({"source_type":s["source_type"],"target":float(s["value_target"]),"prediction":prediction})
|
||||
fig,ax=plt.subplots(figsize=(8,4));x=np.arange(len(value_rows));w=.35
|
||||
ax.bar(x-w/2,[r["target"] for r in value_rows],w,label="return target")
|
||||
ax.bar(x+w/2,[r["prediction"] for r in value_rows],w,label="value prediction")
|
||||
ax.set_xticks(x,[r["source_type"] for r in value_rows],rotation=20,ha="right");ax.set_ylim(-1.05,.05);ax.grid(True,axis="y",alpha=.3);ax.legend();ax.set_title("201-bin distributional value overfit")
|
||||
value_path=outdir/"value_comparison.png";fig.tight_layout();fig.savefig(value_path,dpi=170);plt.close(fig)
|
||||
write_json(outdir/"evaluation.json",{"trajectory_metrics":metrics,"values":value_rows})
|
||||
print({"trajectory_metrics":metrics,"values":value_rows})
|
||||
return trajectory_path,value_path
|
||||
|
||||
|
||||
def main():
|
||||
parser=argparse.ArgumentParser();parser.add_argument("--checkpoint",default="outputs/overfit/checkpoint.pt")
|
||||
evaluate(parser.parse_args().checkpoint)
|
||||
|
||||
|
||||
if __name__=="__main__":main()
|
||||
@@ -0,0 +1,444 @@
|
||||
r"""Minimal :math:`\pi^*_{0.6}` policy with Knowledge Insulation.
|
||||
|
||||
This is intentionally a small, pure-PyTorch architecture test rather than a
|
||||
scale reproduction. It preserves the paper's important information paths:
|
||||
|
||||
* a VLM stream predicts subtask and FAST-lite action tokens;
|
||||
* an optional positive/negative advantage token sits after the subtask and
|
||||
before both action representations;
|
||||
* a separately parameterized flow expert reads observation/subtask/advantage;
|
||||
* flow queries never read teacher-forced FAST-lite tokens; and
|
||||
* VLM keys and values are stop-gradient inputs to the flow expert (Knowledge
|
||||
Insulation), so continuous-action loss cannot update the VLM.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterator, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
@dataclass
|
||||
class Pi06Config:
|
||||
vocab_size: int = 512
|
||||
action_dim: int = 2
|
||||
image_channels: int = 3
|
||||
image_size: int = 64
|
||||
patch_size: int = 16
|
||||
d_model: int = 128
|
||||
n_heads: int = 4
|
||||
n_layers: int = 4
|
||||
mlp_ratio: int = 4
|
||||
max_text_tokens: int = 96
|
||||
max_action_tokens: int = 32
|
||||
action_start_token_id: int = 269
|
||||
dropout: float = 0.0
|
||||
rms_eps: float = 1e-6
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.d_model % self.n_heads:
|
||||
raise ValueError("d_model must be divisible by n_heads")
|
||||
if self.image_size % self.patch_size:
|
||||
raise ValueError("image_size must be divisible by patch_size")
|
||||
if not 0 <= self.action_start_token_id < self.vocab_size:
|
||||
raise ValueError("action_start_token_id must be inside the vocabulary")
|
||||
|
||||
@property
|
||||
def max_image_tokens(self) -> int:
|
||||
return (self.image_size // self.patch_size) ** 2
|
||||
|
||||
@property
|
||||
def max_prefix_tokens(self) -> int:
|
||||
return self.max_image_tokens + self.max_text_tokens
|
||||
|
||||
|
||||
@dataclass
|
||||
class Pi06Output:
|
||||
text_logits: Tensor
|
||||
vector_field: Tensor
|
||||
prefix_hidden: Tensor
|
||||
action_hidden: Tensor
|
||||
text_token_offset: int
|
||||
loss: Optional[Tensor] = None
|
||||
text_loss: Optional[Tensor] = None
|
||||
flow_loss: Optional[Tensor] = None
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
rms = x.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
return x * torch.rsqrt(rms.to(x.dtype) + self.eps) * self.weight
|
||||
|
||||
|
||||
class AdaRMSNorm(nn.Module):
|
||||
"""RMSNorm modulated at every flow-expert layer by the flow timestep."""
|
||||
|
||||
def __init__(self, dim: int, cond_dim: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.norm = RMSNorm(dim, eps)
|
||||
self.modulation = nn.Linear(cond_dim, 2 * dim)
|
||||
nn.init.zeros_(self.modulation.weight)
|
||||
nn.init.zeros_(self.modulation.bias)
|
||||
|
||||
def forward(self, x: Tensor, condition: Tensor) -> Tensor:
|
||||
shift, scale = self.modulation(condition).chunk(2, dim=-1)
|
||||
return self.norm(x) * (1 + scale[:, None]) + shift[:, None]
|
||||
|
||||
|
||||
class TimeEmbedding(nn.Module):
|
||||
def __init__(self, dim: int) -> None:
|
||||
super().__init__()
|
||||
fourier_dim = max(8, dim // 2)
|
||||
self.fourier_dim = fourier_dim + fourier_dim % 2
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(self.fourier_dim, dim), nn.SiLU(), nn.Linear(dim, dim)
|
||||
)
|
||||
|
||||
def forward(self, timestep: Tensor) -> Tensor:
|
||||
timestep = timestep.reshape(-1).float()
|
||||
half = self.fourier_dim // 2
|
||||
frequencies = torch.exp(
|
||||
torch.linspace(
|
||||
math.log(1.0), math.log(1000.0), half,
|
||||
device=timestep.device, dtype=timestep.dtype,
|
||||
)
|
||||
)
|
||||
angles = 2 * math.pi * timestep[:, None] * frequencies[None]
|
||||
return self.mlp(torch.cat((angles.sin(), angles.cos()), dim=-1))
|
||||
|
||||
|
||||
class AttentionWeights(nn.Module):
|
||||
"""One expert's independent Q/K/V/output projections."""
|
||||
|
||||
def __init__(self, dim: int) -> None:
|
||||
super().__init__()
|
||||
self.q = nn.Linear(dim, dim, bias=False)
|
||||
self.k = nn.Linear(dim, dim, bias=False)
|
||||
self.v = nn.Linear(dim, dim, bias=False)
|
||||
self.out = nn.Linear(dim, dim, bias=False)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim: int, ratio: int, dropout: float) -> None:
|
||||
super().__init__()
|
||||
hidden = dim * ratio
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(dim, hidden),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden, dim),
|
||||
nn.Dropout(dropout),
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.net(x)
|
||||
|
||||
|
||||
def attention(
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
n_heads: int,
|
||||
allowed: Optional[Tensor],
|
||||
dropout: float,
|
||||
training: bool,
|
||||
) -> Tensor:
|
||||
batch, q_len, width = query.shape
|
||||
head_dim = width // n_heads
|
||||
|
||||
def split(x: Tensor) -> Tensor:
|
||||
return x.reshape(batch, -1, n_heads, head_dim).transpose(1, 2)
|
||||
|
||||
q, k, v = split(query), split(key), split(value)
|
||||
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim)
|
||||
if allowed is not None:
|
||||
scores = scores.masked_fill(
|
||||
~allowed[:, None], torch.finfo(scores.dtype).min
|
||||
)
|
||||
probabilities = scores.float().softmax(dim=-1).to(scores.dtype)
|
||||
probabilities = F.dropout(probabilities, p=dropout, training=training)
|
||||
output = probabilities @ v
|
||||
return output.transpose(1, 2).contiguous().reshape(batch, q_len, width)
|
||||
|
||||
|
||||
class KnowledgeInsulatedBlock(nn.Module):
|
||||
"""A dual-expert block whose VLM-to-flow connection is stop-gradient."""
|
||||
|
||||
def __init__(self, config: Pi06Config) -> None:
|
||||
super().__init__()
|
||||
d = config.d_model
|
||||
self.n_heads = config.n_heads
|
||||
self.dropout = config.dropout
|
||||
|
||||
self.vlm_norm1 = RMSNorm(d, config.rms_eps)
|
||||
self.vlm_attention = AttentionWeights(d)
|
||||
self.vlm_norm2 = RMSNorm(d, config.rms_eps)
|
||||
self.vlm_mlp = MLP(d, config.mlp_ratio, config.dropout)
|
||||
|
||||
self.action_norm1 = AdaRMSNorm(d, d, config.rms_eps)
|
||||
self.action_attention = AttentionWeights(d)
|
||||
self.action_norm2 = AdaRMSNorm(d, d, config.rms_eps)
|
||||
self.action_mlp = MLP(d, config.mlp_ratio, config.dropout)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
prefix: Tensor,
|
||||
action: Tensor,
|
||||
time_condition: Tensor,
|
||||
prefix_edges: Tensor,
|
||||
action_edges: Tensor,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
prefix_input = self.vlm_norm1(prefix)
|
||||
p = self.vlm_attention
|
||||
prefix_update = attention(
|
||||
p.q(prefix_input), p.k(prefix_input), p.v(prefix_input),
|
||||
self.n_heads, prefix_edges, self.dropout, self.training,
|
||||
)
|
||||
prefix = prefix + p.out(prefix_update)
|
||||
prefix = prefix + self.vlm_mlp(self.vlm_norm2(prefix))
|
||||
|
||||
action_input = self.action_norm1(action, time_condition)
|
||||
a = self.action_attention
|
||||
# Knowledge Insulation: the flow expert may consume VLM activations,
|
||||
# but the flow loss cannot alter either those activations or the VLM's
|
||||
# K/V projection weights. Text loss still trains the VLM normally.
|
||||
prefix_key = p.k(prefix_input).detach()
|
||||
prefix_value = p.v(prefix_input).detach()
|
||||
key = torch.cat((prefix_key, a.k(action_input)), dim=1)
|
||||
value = torch.cat((prefix_value, a.v(action_input)), dim=1)
|
||||
action_update = attention(
|
||||
a.q(action_input), key, value, self.n_heads,
|
||||
action_edges, self.dropout, self.training,
|
||||
)
|
||||
action = action + a.out(action_update)
|
||||
action = action + self.action_mlp(self.action_norm2(action, time_condition))
|
||||
return prefix, action
|
||||
|
||||
|
||||
class Pi06Mini(nn.Module):
|
||||
r"""Small policy retaining the structural contracts of :math:`\pi^*_{0.6}`."""
|
||||
|
||||
def __init__(self, config: Optional[Pi06Config] = None) -> None:
|
||||
super().__init__()
|
||||
self.config = config or Pi06Config()
|
||||
c = self.config
|
||||
self.patch_encoder = nn.Conv2d(
|
||||
c.image_channels, c.d_model, c.patch_size, stride=c.patch_size
|
||||
)
|
||||
self.token_embedding = nn.Embedding(c.vocab_size, c.d_model)
|
||||
self.action_projection = nn.Linear(c.action_dim, c.d_model)
|
||||
self.prefix_position = nn.Parameter(
|
||||
torch.randn(1, c.max_prefix_tokens, c.d_model) * 0.01
|
||||
)
|
||||
self.action_position = nn.Parameter(
|
||||
torch.randn(1, c.max_action_tokens, c.d_model) * 0.01
|
||||
)
|
||||
self.time_embedding = TimeEmbedding(c.d_model)
|
||||
self.layers = nn.ModuleList(
|
||||
KnowledgeInsulatedBlock(c) for _ in range(c.n_layers)
|
||||
)
|
||||
self.prefix_norm = RMSNorm(c.d_model, c.rms_eps)
|
||||
self.action_norm = AdaRMSNorm(c.d_model, c.d_model, c.rms_eps)
|
||||
self.text_head = nn.Linear(c.d_model, c.vocab_size, bias=False)
|
||||
self.vector_field_head = nn.Linear(c.d_model, c.action_dim)
|
||||
self.text_head.weight = self.token_embedding.weight
|
||||
|
||||
def encode_prefix(self, pixel_values: Tensor, input_ids: Tensor) -> tuple[Tensor, int]:
|
||||
patches = self.patch_encoder(pixel_values).flatten(2).transpose(1, 2)
|
||||
prefix = torch.cat((patches, self.token_embedding(input_ids)), dim=1)
|
||||
if prefix.shape[1] > self.config.max_prefix_tokens:
|
||||
raise ValueError("prefix exceeds configured maximum length")
|
||||
return prefix + self.prefix_position[:, : prefix.shape[1]], patches.shape[1]
|
||||
|
||||
def _derive_action_prefix_mask(self, input_ids: Tensor) -> Tensor:
|
||||
"""Allow flow to read every token strictly before the first ACTION_START."""
|
||||
|
||||
is_action = input_ids.eq(self.config.action_start_token_id)
|
||||
# If there is no teacher-forced FAST stream (normal flow inference),
|
||||
# every supplied observation/subtask/advantage token is readable.
|
||||
return is_action.long().cumsum(dim=1).eq(0)
|
||||
|
||||
def _masks(
|
||||
self,
|
||||
batch: int,
|
||||
prefix_length: int,
|
||||
action_length: int,
|
||||
text_offset: int,
|
||||
attention_mask: Optional[Tensor],
|
||||
action_prefix_mask: Tensor,
|
||||
generated_token_start: Optional[int],
|
||||
device: torch.device,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
valid = torch.ones(batch, prefix_length, dtype=torch.bool, device=device)
|
||||
if attention_mask is not None:
|
||||
valid[:, text_offset:] = attention_mask.bool()
|
||||
|
||||
generated_start = (
|
||||
prefix_length
|
||||
if generated_token_start is None
|
||||
else text_offset + generated_token_start
|
||||
)
|
||||
if not text_offset <= generated_start <= prefix_length:
|
||||
raise ValueError("generated_token_start must index input_ids")
|
||||
q = torch.arange(prefix_length, device=device)[:, None]
|
||||
k = torch.arange(prefix_length, device=device)[None]
|
||||
structural = torch.where(q < generated_start, k < generated_start, k <= q)
|
||||
prefix_edges = structural[None] & valid[:, None]
|
||||
identity = torch.eye(prefix_length, dtype=torch.bool, device=device)[None]
|
||||
prefix_edges = torch.where(valid[:, :, None], prefix_edges, identity)
|
||||
|
||||
supplied = action_prefix_mask.bool()
|
||||
if supplied.shape[1] == prefix_length - text_offset:
|
||||
supplied = torch.cat(
|
||||
(
|
||||
torch.ones(batch, text_offset, dtype=torch.bool, device=device),
|
||||
supplied,
|
||||
),
|
||||
dim=1,
|
||||
)
|
||||
if supplied.shape != valid.shape:
|
||||
raise ValueError("action_prefix_mask must cover text tokens or the full prefix")
|
||||
readable = valid & supplied
|
||||
prefix_keys = readable[:, None].expand(batch, action_length, prefix_length)
|
||||
action_keys = torch.ones(
|
||||
batch, action_length, action_length, dtype=torch.bool, device=device
|
||||
)
|
||||
return prefix_edges, torch.cat((prefix_keys, action_keys), dim=-1)
|
||||
|
||||
def vlm_parameters(self) -> Iterator[nn.Parameter]:
|
||||
"""Iterate unique parameters protected from flow-loss gradients by KI."""
|
||||
|
||||
modules: tuple[nn.Module, ...] = (
|
||||
self.patch_encoder,
|
||||
self.token_embedding,
|
||||
self.prefix_norm,
|
||||
self.text_head,
|
||||
)
|
||||
seen: set[int] = set()
|
||||
for parameter in (self.prefix_position,):
|
||||
seen.add(id(parameter))
|
||||
yield parameter
|
||||
for layer in self.layers:
|
||||
modules += (
|
||||
layer.vlm_norm1,
|
||||
layer.vlm_attention,
|
||||
layer.vlm_norm2,
|
||||
layer.vlm_mlp,
|
||||
)
|
||||
for module in modules:
|
||||
for parameter in module.parameters():
|
||||
if id(parameter) not in seen:
|
||||
seen.add(id(parameter))
|
||||
yield parameter
|
||||
|
||||
def action_parameters(self) -> Iterator[nn.Parameter]:
|
||||
"""Iterate parameters trained by the continuous flow loss."""
|
||||
|
||||
modules: list[nn.Module] = [
|
||||
self.action_projection,
|
||||
self.time_embedding,
|
||||
self.action_norm,
|
||||
self.vector_field_head,
|
||||
]
|
||||
for layer in self.layers:
|
||||
modules.extend(
|
||||
(
|
||||
layer.action_norm1,
|
||||
layer.action_attention,
|
||||
layer.action_norm2,
|
||||
layer.action_mlp,
|
||||
)
|
||||
)
|
||||
yield self.action_position
|
||||
for module in modules:
|
||||
yield from module.parameters()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pixel_values: Tensor,
|
||||
input_ids: Tensor,
|
||||
noisy_actions: Tensor,
|
||||
timesteps: Tensor,
|
||||
attention_mask: Optional[Tensor] = None,
|
||||
action_prefix_mask: Optional[Tensor] = None,
|
||||
generated_token_start: Optional[int] = None,
|
||||
text_labels: Optional[Tensor] = None,
|
||||
target_velocity: Optional[Tensor] = None,
|
||||
text_loss_weight: float = 1.0,
|
||||
flow_loss_weight: float = 1.0,
|
||||
) -> Pi06Output:
|
||||
if noisy_actions.ndim != 3:
|
||||
raise ValueError("noisy_actions must have shape [batch, horizon, action_dim]")
|
||||
if noisy_actions.shape[-1] != self.config.action_dim:
|
||||
raise ValueError("noisy_actions has the wrong action dimension")
|
||||
if noisy_actions.shape[1] > self.config.max_action_tokens:
|
||||
raise ValueError("action horizon exceeds configured maximum")
|
||||
|
||||
prefix, text_offset = self.encode_prefix(pixel_values, input_ids)
|
||||
horizon = noisy_actions.shape[1]
|
||||
action = self.action_projection(noisy_actions)
|
||||
action = action + self.action_position[:, :horizon]
|
||||
condition = self.time_embedding(timesteps).to(action.dtype)
|
||||
if action_prefix_mask is None:
|
||||
action_prefix_mask = self._derive_action_prefix_mask(input_ids)
|
||||
prefix_edges, action_edges = self._masks(
|
||||
prefix.shape[0], prefix.shape[1], horizon, text_offset,
|
||||
attention_mask, action_prefix_mask, generated_token_start, prefix.device,
|
||||
)
|
||||
for layer in self.layers:
|
||||
prefix, action = layer(
|
||||
prefix, action, condition, prefix_edges, action_edges
|
||||
)
|
||||
|
||||
prefix_hidden = self.prefix_norm(prefix)
|
||||
action_hidden = self.action_norm(action, condition)
|
||||
text_logits = self.text_head(prefix_hidden[:, text_offset:])
|
||||
vector_field = self.vector_field_head(action_hidden)
|
||||
|
||||
text_loss = None
|
||||
if text_labels is not None:
|
||||
text_loss = F.cross_entropy(
|
||||
text_logits.reshape(-1, text_logits.shape[-1]),
|
||||
text_labels.reshape(-1), ignore_index=-100,
|
||||
)
|
||||
flow_loss = None
|
||||
if target_velocity is not None:
|
||||
flow_loss = F.mse_loss(vector_field, target_velocity)
|
||||
terms = [
|
||||
term for term in (
|
||||
None if text_loss is None else text_loss_weight * text_loss,
|
||||
None if flow_loss is None else flow_loss_weight * flow_loss,
|
||||
) if term is not None
|
||||
]
|
||||
loss = sum(terms) if terms else None
|
||||
return Pi06Output(
|
||||
text_logits=text_logits,
|
||||
vector_field=vector_field,
|
||||
prefix_hidden=prefix_hidden,
|
||||
action_hidden=action_hidden,
|
||||
text_token_offset=text_offset,
|
||||
loss=loss,
|
||||
text_loss=text_loss,
|
||||
flow_loss=flow_loss,
|
||||
)
|
||||
|
||||
|
||||
def flow_matching_sample(actions: Tensor) -> tuple[Tensor, Tensor, Tensor]:
|
||||
"""Construct the linear conditional-flow target used by the policy."""
|
||||
|
||||
batch = actions.shape[0]
|
||||
timestep = torch.rand(batch, device=actions.device, dtype=actions.dtype)
|
||||
noise = torch.randn_like(actions)
|
||||
broadcast_t = timestep.reshape(batch, *([1] * (actions.ndim - 1)))
|
||||
noisy_actions = (1 - broadcast_t) * noise + broadcast_t * actions
|
||||
return noisy_actions, timestep, actions - noise
|
||||
@@ -0,0 +1,78 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def condition_dropout_mask(
|
||||
batch_size: int,
|
||||
*,
|
||||
drop_probability: float = 0.3,
|
||||
device: torch.device | str | None = None,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Sample the 30% advantage-condition dropout used for CFG training."""
|
||||
|
||||
if batch_size < 1:
|
||||
raise ValueError("batch_size must be positive")
|
||||
if not 0.0 <= drop_probability <= 1.0:
|
||||
raise ValueError("drop_probability must lie in [0, 1]")
|
||||
return torch.rand(batch_size, device=device, generator=generator) < drop_probability
|
||||
|
||||
|
||||
def apply_condition_dropout(
|
||||
condition: torch.Tensor,
|
||||
null_condition: torch.Tensor | float | int = 0,
|
||||
*,
|
||||
drop_probability: float = 0.3,
|
||||
mask: torch.Tensor | None = None,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Replace complete per-example ReCAP conditions with a null condition."""
|
||||
|
||||
condition = torch.as_tensor(condition)
|
||||
if condition.ndim < 1 or condition.shape[0] < 1:
|
||||
raise ValueError("condition must have a non-empty batch dimension")
|
||||
if mask is None:
|
||||
mask = condition_dropout_mask(
|
||||
condition.shape[0],
|
||||
drop_probability=drop_probability,
|
||||
device=condition.device,
|
||||
generator=generator,
|
||||
)
|
||||
else:
|
||||
mask = torch.as_tensor(mask, dtype=torch.bool, device=condition.device)
|
||||
if mask.shape != (condition.shape[0],):
|
||||
raise ValueError("mask must have shape [batch_size]")
|
||||
null = torch.as_tensor(null_condition, dtype=condition.dtype, device=condition.device)
|
||||
batch_mask = mask
|
||||
while mask.ndim < condition.ndim:
|
||||
mask = mask.unsqueeze(-1)
|
||||
dropped = torch.where(mask, null, condition)
|
||||
return dropped, batch_mask
|
||||
|
||||
|
||||
def classifier_free_guidance(
|
||||
conditional_prediction: torch.Tensor,
|
||||
unconditional_prediction: torch.Tensor,
|
||||
*,
|
||||
guidance_scale: float = 1.5,
|
||||
) -> torch.Tensor:
|
||||
"""Combine conditional and null-condition flow fields.
|
||||
|
||||
A scale of zero returns the unconditional prediction and a scale of one
|
||||
returns the ordinary positive-condition prediction.
|
||||
"""
|
||||
|
||||
conditional_prediction = torch.as_tensor(conditional_prediction)
|
||||
unconditional_prediction = torch.as_tensor(
|
||||
unconditional_prediction,
|
||||
dtype=conditional_prediction.dtype,
|
||||
device=conditional_prediction.device,
|
||||
)
|
||||
if conditional_prediction.shape != unconditional_prediction.shape:
|
||||
raise ValueError("conditional and unconditional predictions must have equal shape")
|
||||
if guidance_scale < 0.0:
|
||||
raise ValueError("guidance_scale must be non-negative")
|
||||
return unconditional_prediction + guidance_scale * (
|
||||
conditional_prediction - unconditional_prediction
|
||||
)
|
||||
@@ -0,0 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
|
||||
def load_config(path: str | Path) -> dict:
|
||||
with Path(path).open("r", encoding="utf-8") as handle:
|
||||
return yaml.safe_load(handle)
|
||||
|
||||
|
||||
def seed_everything(seed: int) -> None:
|
||||
random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
|
||||
|
||||
|
||||
def choose_device(name: str = "auto") -> torch.device:
|
||||
if name != "auto": return torch.device(name)
|
||||
if torch.cuda.is_available(): return torch.device("cuda")
|
||||
if torch.backends.mps.is_available(): return torch.device("mps")
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
def write_json(path: str | Path, value: object) -> None:
|
||||
target = Path(path); target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(json.dumps(value, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
@@ -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,
|
||||
)
|
||||
@@ -0,0 +1,153 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
from pathlib import Path
|
||||
import torch
|
||||
|
||||
from .data import Pi06ExperienceDataset, materialize_subset
|
||||
from .flow import FlowMatching
|
||||
from .model import Pi06Config, Pi06Mini
|
||||
from .recap import apply_condition_dropout
|
||||
from .runtime import choose_device, load_config, seed_everything, write_json
|
||||
from .tokenizer import MiniTokenizer
|
||||
from .value import ValueConfig, ValueModel
|
||||
|
||||
|
||||
def build_models(cfg: dict, tokenizer: MiniTokenizer) -> tuple[Pi06Mini, ValueModel]:
|
||||
m = cfg["model"]
|
||||
policy = Pi06Mini(Pi06Config(
|
||||
vocab_size=tokenizer.vocab_size, action_dim=2,
|
||||
image_size=cfg["data"]["image_size"], patch_size=m["patch_size"],
|
||||
d_model=m["policy_d_model"], n_heads=m["n_heads"], n_layers=m["n_layers"],
|
||||
mlp_ratio=m["mlp_ratio"], max_text_tokens=m["max_text_tokens"],
|
||||
max_action_tokens=cfg["data"]["action_horizon"],
|
||||
action_start_token_id=tokenizer.ACTION_START, dropout=m["dropout"],
|
||||
))
|
||||
value = ValueModel(ValueConfig(
|
||||
vocab_size=tokenizer.vocab_size, image_size=cfg["data"]["image_size"],
|
||||
patch_size=m["patch_size"], d_model=m["value_d_model"], n_heads=m["n_heads"],
|
||||
n_layers=2, mlp_ratio=m["mlp_ratio"], max_text_tokens=m["max_text_tokens"],
|
||||
num_bins=m["value_bins"], value_min=m["value_min"], value_max=m["value_max"],
|
||||
))
|
||||
return policy, value
|
||||
|
||||
|
||||
def policy_step(
|
||||
model: Pi06Mini, sample: dict, tokenizer: MiniTokenizer, flow: FlowMatching,
|
||||
device: torch.device, advantage: str | None, cfg: dict,
|
||||
) -> tuple[torch.Tensor, dict[str, float]]:
|
||||
sequence = tokenizer.build_policy_sequence(
|
||||
sample["history_xy"], sample["reasoning"], sample["actions"], advantage,
|
||||
# Outcome/source metadata is reserved for the critic in this synthetic
|
||||
# proxy. The policy must distinguish behavior through advantage alone.
|
||||
metadata=None,
|
||||
)
|
||||
image = sample["image"][None].to(device)
|
||||
ids = sequence.input_ids[None].to(device)
|
||||
labels = sequence.labels[None].to(device)
|
||||
actions = sample["actions"][None].to(device)
|
||||
noisy, timestep, velocity, weight = flow.training_batch(actions)
|
||||
output = model(
|
||||
image, ids, noisy, timestep,
|
||||
action_prefix_mask=sequence.action_prefix_mask[None].to(device),
|
||||
generated_token_start=sequence.generated_token_start,
|
||||
text_labels=labels,
|
||||
)
|
||||
flow_loss = flow.weighted_loss(output.vector_field, velocity, weight)
|
||||
loss = cfg["train"]["token_loss_weight"] * output.text_loss + cfg["train"]["flow_loss_weight"] * flow_loss
|
||||
return loss, {
|
||||
"text_loss": float(output.text_loss.detach().cpu()),
|
||||
"flow_loss": float(flow_loss.detach().cpu()),
|
||||
}
|
||||
|
||||
|
||||
def train(config_path: str) -> Path:
|
||||
cfg = load_config(config_path)
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
seed_everything(cfg["seed"])
|
||||
device = choose_device(cfg.get("device", "auto"))
|
||||
manifest = materialize_subset(cfg["data"]["source_manifest"], root / "data/overfit")
|
||||
dataset = Pi06ExperienceDataset(
|
||||
manifest, image_size=cfg["data"]["image_size"],
|
||||
augmentation_factor=cfg["experience"]["augmentation_factor"],
|
||||
action_horizon=cfg["data"]["action_horizon"], n_step=cfg["experience"]["n_step"],
|
||||
gamma=cfg["experience"]["gamma"], positive_quantile=cfg["experience"]["positive_quantile"],
|
||||
)
|
||||
tokenizer = MiniTokenizer()
|
||||
base, value = build_models(cfg, tokenizer)
|
||||
base, value = base.to(device), value.to(device)
|
||||
flow = FlowMatching(cfg["inference"]["flow_steps"])
|
||||
outdir = root / cfg["output_dir"]; outdir.mkdir(parents=True, exist_ok=True)
|
||||
history: list[dict] = []
|
||||
|
||||
demos = [dataset[i] for i in range(len(dataset)) if dataset[i]["source_type"] == "demo_success"]
|
||||
optimizer = torch.optim.AdamW(base.parameters(), lr=cfg["train"]["learning_rate"])
|
||||
base.train()
|
||||
for step in range(1, cfg["train"]["base_steps"] + 1):
|
||||
sample = demos[(step - 1) % len(demos)]
|
||||
loss, metrics = policy_step(base, sample, tokenizer, flow, device, None, cfg)
|
||||
optimizer.zero_grad(set_to_none=True); loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(base.parameters(), cfg["train"]["grad_clip_norm"]); optimizer.step()
|
||||
if step == 1 or step % cfg["train"]["log_every"] == 0:
|
||||
row={"stage":"base","step":step,"loss":float(loss.detach().cpu()),**metrics};history.append(row);print(json.dumps(row))
|
||||
base_state = copy.deepcopy(base.state_dict())
|
||||
torch.save(base_state, outdir / "pi06_base.pt")
|
||||
|
||||
value_opt = torch.optim.AdamW(value.parameters(), lr=cfg["train"]["value_learning_rate"])
|
||||
value.train()
|
||||
for step in range(1, cfg["train"]["value_steps"] + 1):
|
||||
sample = dataset[(step - 1) % len(dataset)]
|
||||
ids = tokenizer.build_value_input(sample["history_xy"], sample["source_type"])[None].to(device)
|
||||
output = value(sample["image"][None].to(device), ids, target_values=sample["value_target"][None].to(device))
|
||||
value_opt.zero_grad(set_to_none=True); output.loss.backward(); value_opt.step()
|
||||
if step == 1 or step % cfg["train"]["log_every"] == 0:
|
||||
row={"stage":"value","step":step,"loss":float(output.loss.detach().cpu()),"prediction":float(output.expectation.detach().cpu()),"target":float(sample["value_target"])};history.append(row);print(json.dumps(row))
|
||||
|
||||
def fit_policy(stage: str, steps: int, conditioned: bool) -> Pi06Mini:
|
||||
model, _ = build_models(cfg, tokenizer); model.load_state_dict(base_state); model.to(device).train()
|
||||
opt = torch.optim.AdamW(model.parameters(), lr=cfg["train"]["learning_rate"])
|
||||
for step in range(1, steps + 1):
|
||||
sample = dataset[(step - 1) % len(dataset)]
|
||||
advantage = None
|
||||
if conditioned:
|
||||
dropped = bool(apply_condition_dropout(
|
||||
torch.tensor([sample["positive_condition"]]),
|
||||
drop_probability=cfg["recap"]["condition_dropout"],
|
||||
)[1][0])
|
||||
advantage = None if dropped else sample["advantage_label"]
|
||||
loss, metrics = policy_step(model, sample, tokenizer, flow, device, advantage, cfg)
|
||||
opt.zero_grad(set_to_none=True); loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"]); opt.step()
|
||||
if step == 1 or step % cfg["train"]["log_every"] == 0:
|
||||
row={"stage":stage,"step":step,"loss":float(loss.detach().cpu()),**metrics};history.append(row);print(json.dumps(row))
|
||||
return model
|
||||
|
||||
mixed = fit_policy("mixed", cfg["train"]["mixed_steps"], conditioned=False)
|
||||
recap = fit_policy("recap", cfg["train"]["recap_steps"], conditioned=True)
|
||||
torch.save(mixed.state_dict(), outdir / "mixed_sft.pt")
|
||||
torch.save(recap.state_dict(), outdir / "pi06star_recap.pt")
|
||||
checkpoint = {
|
||||
"base": base_state, "mixed": mixed.state_dict(), "recap": recap.state_dict(),
|
||||
"value": value.state_dict(), "policy_config": base.config.__dict__,
|
||||
"value_config": value.config.__dict__, "config": cfg,
|
||||
"action_statistics": dataset.action_normalizer.as_dict(),
|
||||
}
|
||||
torch.save(checkpoint, outdir / "checkpoint.pt")
|
||||
write_json(outdir / "training_history.json", history)
|
||||
write_json(outdir / "summary.json", {
|
||||
"device": str(device), "experience_samples": len(dataset),
|
||||
"policy_parameters": sum(p.numel() for p in base.parameters()),
|
||||
"value_parameters": sum(p.numel() for p in value.parameters()),
|
||||
"sources": [dataset[i]["source_type"] for i in range(len(dataset))],
|
||||
})
|
||||
return outdir / "checkpoint.pt"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser=argparse.ArgumentParser();parser.add_argument("--config",default="configs/overfit.yaml")
|
||||
train(parser.parse_args().config)
|
||||
|
||||
|
||||
if __name__ == "__main__": main()
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Independent 201-bin distributional value model for minimal Recap training."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .model import AttentionWeights, MLP, RMSNorm, attention
|
||||
|
||||
|
||||
@dataclass
|
||||
class ValueConfig:
|
||||
vocab_size: int = 512
|
||||
image_channels: int = 3
|
||||
image_size: int = 64
|
||||
patch_size: int = 16
|
||||
d_model: int = 96
|
||||
n_heads: int = 4
|
||||
n_layers: int = 3
|
||||
mlp_ratio: int = 4
|
||||
max_text_tokens: int = 64
|
||||
num_bins: int = 201
|
||||
value_min: float = -1.0
|
||||
value_max: float = 0.0
|
||||
dropout: float = 0.0
|
||||
rms_eps: float = 1e-6
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.d_model % self.n_heads:
|
||||
raise ValueError("d_model must be divisible by n_heads")
|
||||
if self.image_size % self.patch_size:
|
||||
raise ValueError("image_size must be divisible by patch_size")
|
||||
if self.num_bins < 2:
|
||||
raise ValueError("num_bins must be at least two")
|
||||
if self.value_max <= self.value_min:
|
||||
raise ValueError("value_max must be greater than value_min")
|
||||
|
||||
@property
|
||||
def max_image_tokens(self) -> int:
|
||||
return (self.image_size // self.patch_size) ** 2
|
||||
|
||||
@property
|
||||
def max_sequence_tokens(self) -> int:
|
||||
# One learned pooling token plus image and discrete-state tokens.
|
||||
return 1 + self.max_image_tokens + self.max_text_tokens
|
||||
|
||||
|
||||
@dataclass
|
||||
class ValueOutput:
|
||||
logits: Tensor
|
||||
probabilities: Tensor
|
||||
expectation: Tensor
|
||||
hidden: Tensor
|
||||
loss: Optional[Tensor] = None
|
||||
target_bins: Optional[Tensor] = None
|
||||
|
||||
|
||||
class ValueBlock(nn.Module):
|
||||
def __init__(self, config: ValueConfig) -> None:
|
||||
super().__init__()
|
||||
self.n_heads = config.n_heads
|
||||
self.dropout = config.dropout
|
||||
self.norm1 = RMSNorm(config.d_model, config.rms_eps)
|
||||
self.attention = AttentionWeights(config.d_model)
|
||||
self.norm2 = RMSNorm(config.d_model, config.rms_eps)
|
||||
self.mlp = MLP(config.d_model, config.mlp_ratio, config.dropout)
|
||||
|
||||
def forward(self, x: Tensor, edges: Tensor) -> Tensor:
|
||||
normalized = self.norm1(x)
|
||||
weights = self.attention
|
||||
update = attention(
|
||||
weights.q(normalized),
|
||||
weights.k(normalized),
|
||||
weights.v(normalized),
|
||||
self.n_heads,
|
||||
edges,
|
||||
self.dropout,
|
||||
self.training,
|
||||
)
|
||||
x = x + weights.out(update)
|
||||
return x + self.mlp(self.norm2(x))
|
||||
|
||||
|
||||
class DistributionalValueModel(nn.Module):
|
||||
"""Predict ``p(V | image, discrete state)`` over 201 return bins.
|
||||
|
||||
The value network is deliberately a separate model: it shares no
|
||||
parameters with ``Pi06Mini``. By default its support is the paper's
|
||||
task-normalized interval ``[-1, 0]``.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[ValueConfig] = None) -> None:
|
||||
super().__init__()
|
||||
self.config = config or ValueConfig()
|
||||
c = self.config
|
||||
self.patch_encoder = nn.Conv2d(
|
||||
c.image_channels, c.d_model, c.patch_size, stride=c.patch_size
|
||||
)
|
||||
self.token_embedding = nn.Embedding(c.vocab_size, c.d_model)
|
||||
self.value_token = nn.Parameter(torch.randn(1, 1, c.d_model) * 0.01)
|
||||
self.position = nn.Parameter(
|
||||
torch.randn(1, c.max_sequence_tokens, c.d_model) * 0.01
|
||||
)
|
||||
self.layers = nn.ModuleList(ValueBlock(c) for _ in range(c.n_layers))
|
||||
self.norm = RMSNorm(c.d_model, c.rms_eps)
|
||||
self.value_head = nn.Linear(c.d_model, c.num_bins)
|
||||
self.register_buffer(
|
||||
"support",
|
||||
torch.linspace(c.value_min, c.value_max, c.num_bins),
|
||||
persistent=True,
|
||||
)
|
||||
|
||||
def values_to_bins(self, values: Tensor) -> Tensor:
|
||||
"""Nearest-bin discretization of normalized Monte-Carlo returns."""
|
||||
|
||||
values = torch.as_tensor(values, dtype=self.support.dtype, device=self.support.device)
|
||||
unit = (values.clamp(self.config.value_min, self.config.value_max) - self.config.value_min)
|
||||
unit = unit / (self.config.value_max - self.config.value_min)
|
||||
return torch.round(unit * (self.config.num_bins - 1)).long()
|
||||
|
||||
def bins_to_values(self, bins: Tensor) -> Tensor:
|
||||
bins = torch.as_tensor(bins, dtype=torch.long, device=self.support.device)
|
||||
if bins.numel() and (bins.min() < 0 or bins.max() >= self.config.num_bins):
|
||||
raise ValueError("value bins are outside the configured support")
|
||||
return self.support[bins]
|
||||
|
||||
def expectation_from_logits(self, logits: Tensor) -> Tensor:
|
||||
return (logits.softmax(dim=-1) * self.support.to(logits.dtype)).sum(dim=-1)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pixel_values: Tensor,
|
||||
input_ids: Tensor,
|
||||
attention_mask: Optional[Tensor] = None,
|
||||
target_values: Optional[Tensor] = None,
|
||||
target_bins: Optional[Tensor] = None,
|
||||
) -> ValueOutput:
|
||||
if target_values is not None and target_bins is not None:
|
||||
raise ValueError("provide target_values or target_bins, not both")
|
||||
patches = self.patch_encoder(pixel_values).flatten(2).transpose(1, 2)
|
||||
text = self.token_embedding(input_ids)
|
||||
batch = input_ids.shape[0]
|
||||
pooling = self.value_token.expand(batch, -1, -1)
|
||||
hidden = torch.cat((pooling, patches, text), dim=1)
|
||||
if hidden.shape[1] > self.config.max_sequence_tokens:
|
||||
raise ValueError("value-model sequence exceeds configured maximum length")
|
||||
hidden = hidden + self.position[:, : hidden.shape[1]]
|
||||
|
||||
valid = torch.ones(batch, hidden.shape[1], dtype=torch.bool, device=hidden.device)
|
||||
if attention_mask is not None:
|
||||
if attention_mask.shape != input_ids.shape:
|
||||
raise ValueError("attention_mask must have the same shape as input_ids")
|
||||
valid[:, 1 + patches.shape[1] :] = attention_mask.bool()
|
||||
edges = valid[:, None, :].expand(batch, hidden.shape[1], hidden.shape[1])
|
||||
identity = torch.eye(hidden.shape[1], dtype=torch.bool, device=hidden.device)[None]
|
||||
edges = torch.where(valid[:, :, None], edges, identity)
|
||||
for layer in self.layers:
|
||||
hidden = layer(hidden, edges)
|
||||
hidden = self.norm(hidden)
|
||||
|
||||
logits = self.value_head(hidden[:, 0])
|
||||
probabilities = logits.softmax(dim=-1)
|
||||
expectation = (probabilities * self.support.to(logits.dtype)).sum(dim=-1)
|
||||
|
||||
bins = None
|
||||
if target_values is not None:
|
||||
bins = self.values_to_bins(target_values).to(logits.device)
|
||||
elif target_bins is not None:
|
||||
bins = torch.as_tensor(target_bins, dtype=torch.long, device=logits.device)
|
||||
if bins.numel() and (bins.min() < 0 or bins.max() >= self.config.num_bins):
|
||||
raise ValueError("target_bins are outside the configured support")
|
||||
loss = None if bins is None else F.cross_entropy(logits, bins.reshape(-1))
|
||||
return ValueOutput(
|
||||
logits=logits,
|
||||
probabilities=probabilities,
|
||||
expectation=expectation,
|
||||
hidden=hidden,
|
||||
loss=loss,
|
||||
target_bins=bins,
|
||||
)
|
||||
|
||||
|
||||
# Short name convenient in configs and training scripts.
|
||||
ValueModel = DistributionalValueModel
|
||||
Reference in New Issue
Block a user