Initial pi0.5 minimal validation
This commit is contained in:
seeker
@@ -0,0 +1,5 @@
|
||||
"""Minimal pi0.5 architecture validation package."""
|
||||
|
||||
from .model import Pi05Config, Pi05Mini, Pi05Output, flow_matching_sample
|
||||
|
||||
__all__ = ["Pi05Config", "Pi05Mini", "Pi05Output", "flow_matching_sample"]
|
||||
@@ -0,0 +1,316 @@
|
||||
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 .tokenizer import MiniTokenizer
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SampleRecord:
|
||||
clip_id: str
|
||||
t0_us: int
|
||||
reasoning: str
|
||||
image_path: Path
|
||||
trajectory_path: Path
|
||||
|
||||
@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:
|
||||
"""Normalize 2-D waypoint deltas and retain invertible dataset statistics."""
|
||||
|
||||
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)
|
||||
|
||||
@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 sequences of equal length")
|
||||
chunks = []
|
||||
for history, future in zip(histories, futures, strict=True):
|
||||
chunks.append(cls.to_deltas(history, future).reshape(-1, 2))
|
||||
flat = torch.cat(chunks, dim=0).float()
|
||||
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())
|
||||
|
||||
@staticmethod
|
||||
def to_deltas(history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor:
|
||||
previous = torch.cat((history_xy[-1:].float(), future_xy[:-1].float()), dim=0)
|
||||
return future_xy.float() - previous
|
||||
|
||||
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:
|
||||
deltas = actions.float() * self.scale + self.mean
|
||||
return history_xy[-1:].float() + 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:
|
||||
output = Path(path).expanduser().resolve()
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(self.as_dict(), indent=2) + "\n", encoding="utf-8")
|
||||
return output
|
||||
|
||||
|
||||
def _resolve_data_path(manifest_path: Path, raw_path: str | Path) -> Path:
|
||||
path = Path(raw_path).expanduser()
|
||||
if path.is_absolute():
|
||||
return path
|
||||
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"]),
|
||||
reasoning=str(row.get("reasoning", "")),
|
||||
image_path=_resolve_data_path(manifest_path, row["image_path"]),
|
||||
trajectory_path=_resolve_data_path(manifest_path, row["trajectory_path"]),
|
||||
)
|
||||
)
|
||||
if not records:
|
||||
raise ValueError(f"no samples found in {manifest_path}")
|
||||
return records
|
||||
|
||||
|
||||
def materialize_subset(
|
||||
source_manifest: str | Path,
|
||||
output_dir: str | Path,
|
||||
max_samples: int = 3,
|
||||
num_variants: int | None = None,
|
||||
) -> Path:
|
||||
"""Copy a self-contained PAI overfit subset out of the Alpamayo project."""
|
||||
|
||||
# ``num_variants`` is accepted as a compatibility alias used by the
|
||||
# training entry point. Actual variants stay virtual and deterministic in
|
||||
# Pi05OverfitDataset, avoiding redundant copies of the large JPEG.
|
||||
if num_variants is not None:
|
||||
if num_variants < 1:
|
||||
raise ValueError("num_variants must be at least one")
|
||||
max_samples = min(max_samples, num_variants)
|
||||
if max_samples < 1:
|
||||
raise ValueError("max_samples must be at least one")
|
||||
records = load_manifest(source_manifest)[:max_samples]
|
||||
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(records):
|
||||
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)),
|
||||
"history_steps": 8,
|
||||
"expert_waypoints": 16,
|
||||
}
|
||||
)
|
||||
|
||||
manifest_path = destination / "overfit_manifest.jsonl"
|
||||
manifest_path.write_text(
|
||||
"".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return manifest_path
|
||||
|
||||
|
||||
def _resample_waypoints(xy: torch.Tensor, steps: int = 16) -> torch.Tensor:
|
||||
if xy.shape[0] == steps:
|
||||
return xy.float()
|
||||
if xy.shape[0] < 2:
|
||||
raise ValueError("future trajectory must contain at least two waypoints")
|
||||
signal = xy.float().T.unsqueeze(0)
|
||||
return F.interpolate(signal, size=steps, mode="linear", align_corners=True)[0].T
|
||||
|
||||
|
||||
class Pi05OverfitDataset(Dataset[dict[str, Any]]):
|
||||
"""Tiny deterministic PAI dataset for validating a pi0.5-style training loop."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path: str | Path,
|
||||
*,
|
||||
image_size: int | tuple[int, int] = (128, 224),
|
||||
augmentation_factor: int = 4,
|
||||
action_horizon: int = 16,
|
||||
seed: int = 7,
|
||||
tokenizer: MiniTokenizer | None = None,
|
||||
action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None,
|
||||
) -> None:
|
||||
if augmentation_factor < 1:
|
||||
raise ValueError("augmentation_factor must be at least one")
|
||||
if isinstance(image_size, int):
|
||||
image_size = (image_size, image_size)
|
||||
if len(image_size) != 2 or min(image_size) < 1:
|
||||
raise ValueError("image_size must be a positive (height, width) pair")
|
||||
if action_horizon < 2:
|
||||
raise ValueError("action_horizon must be at least two")
|
||||
self.records = load_manifest(manifest_path)
|
||||
self.image_size = (int(image_size[0]), int(image_size[1]))
|
||||
self.augmentation_factor = int(augmentation_factor)
|
||||
self.action_horizon = int(action_horizon)
|
||||
self.seed = int(seed)
|
||||
self.tokenizer = tokenizer or MiniTokenizer()
|
||||
self.state_dim = 16
|
||||
|
||||
if action_statistics is None:
|
||||
geometries = [
|
||||
self._load_augmented_geometry(base_index, variant)
|
||||
for base_index in range(len(self.records))
|
||||
for variant in range(self.augmentation_factor)
|
||||
]
|
||||
histories = [geometry[0] for geometry in geometries]
|
||||
futures = [geometry[1] for geometry in geometries]
|
||||
self.action_normalizer = ActionNormalizer.fit(histories, futures)
|
||||
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
|
||||
|
||||
@property
|
||||
def action_mean(self) -> torch.Tensor:
|
||||
return self.action_normalizer.mean.clone()
|
||||
|
||||
@property
|
||||
def action_std(self) -> torch.Tensor:
|
||||
return self.action_normalizer.scale.clone()
|
||||
|
||||
def save_action_statistics(self, path: str | Path) -> Path:
|
||||
return self.action_normalizer.save(path)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.records) * 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 _rng(self, base_index: int, variant: int) -> np.random.Generator:
|
||||
return np.random.default_rng(
|
||||
np.random.SeedSequence([self.seed, base_index, variant])
|
||||
)
|
||||
|
||||
def _load_augmented_geometry(
|
||||
self, base_index: int, variant: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
record = self.records[base_index]
|
||||
with np.load(record.trajectory_path) as arrays:
|
||||
history = torch.from_numpy(np.asarray(arrays["history_xyz"], dtype=np.float32))[:, :2]
|
||||
future = torch.from_numpy(np.asarray(arrays["future_xyz"], dtype=np.float32))[:, :2]
|
||||
future = _resample_waypoints(future, steps=self.action_horizon)
|
||||
if variant == 0:
|
||||
return history, future
|
||||
|
||||
rng = self._rng(base_index, variant)
|
||||
angle = float(rng.normal(0.0, 0.008))
|
||||
scale = float(rng.normal(1.0, 0.006))
|
||||
rotation = history.new_tensor(
|
||||
[[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]
|
||||
)
|
||||
history = scale * (history @ rotation.T)
|
||||
future = scale * (future @ rotation.T)
|
||||
# A tiny smooth lateral perturbation adds distinct trajectories while retaining the label.
|
||||
curvature = float(rng.normal(0.0, 0.012))
|
||||
progress = torch.linspace(0.0, 1.0, future.shape[0], dtype=future.dtype)
|
||||
future[:, 1] += curvature * progress.square()
|
||||
return history, future
|
||||
|
||||
def _load_image(self, base_index: int, variant: int) -> torch.Tensor:
|
||||
with Image.open(self.records[base_index].image_path) as source:
|
||||
image = source.convert("RGB")
|
||||
if variant:
|
||||
rng = self._rng(base_index, variant)
|
||||
brightness = float(rng.uniform(0.97, 1.03))
|
||||
image = ImageEnhance.Brightness(image).enhance(brightness)
|
||||
height, width = self.image_size
|
||||
image = image.resize((width, height), 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]:
|
||||
base_index, variant = self._indices(index)
|
||||
record = self.records[base_index]
|
||||
history_xy, future_xy = self._load_augmented_geometry(base_index, variant)
|
||||
reasoning = record.reasoning
|
||||
target = self.tokenizer.encode_subtask(reasoning)
|
||||
actions = self.action_normalizer.encode(history_xy, future_xy)
|
||||
return {
|
||||
"image": self._load_image(base_index, variant),
|
||||
"prefix_tokens": self.tokenizer.build_prefix(history_xy),
|
||||
"subtask_target": target,
|
||||
"subtask_label": self.tokenizer.subtasks.decode(target),
|
||||
"actions": actions,
|
||||
"action_tokens": self.tokenizer.encode_actions(actions),
|
||||
"state": history_xy.reshape(-1),
|
||||
"history_xy": history_xy,
|
||||
"future_xy": future_xy,
|
||||
"reasoning": reasoning,
|
||||
"sample_id": f"{record.sample_id}:aug{variant}",
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class FlowMatching:
|
||||
"""Rectified-flow objective used by the minimal action expert."""
|
||||
|
||||
def __init__(self, inference_steps: int = 10, cutoff: float = 0.999) -> None:
|
||||
self.inference_steps = inference_steps
|
||||
self.cutoff = cutoff
|
||||
|
||||
def training_batch(
|
||||
self, target: Tensor, generator: torch.Generator | None = None
|
||||
) -> tuple[Tensor, Tensor, Tensor]:
|
||||
batch = target.shape[0]
|
||||
uniform = torch.rand(
|
||||
batch, device=target.device, dtype=target.dtype, generator=generator
|
||||
)
|
||||
# Paper sampling: z~Beta(1.5,1), tau=s*(1-z), emphasizing noisy/low tau.
|
||||
timestep = self.cutoff * (1.0 - uniform.pow(1.0 / 1.5))
|
||||
noise = torch.randn(
|
||||
target.shape,
|
||||
device=target.device,
|
||||
dtype=target.dtype,
|
||||
generator=generator,
|
||||
)
|
||||
t = timestep[:, None, None]
|
||||
noisy = t * target + (1.0 - t) * noise
|
||||
return noisy, timestep, target - noise
|
||||
|
||||
@staticmethod
|
||||
def loss(prediction: Tensor, target_velocity: Tensor) -> Tensor:
|
||||
return torch.nn.functional.mse_loss(prediction.float(), target_velocity.float())
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(
|
||||
self,
|
||||
model_fn: Callable[[Tensor, Tensor], Tensor],
|
||||
shape: tuple[int, ...],
|
||||
device: torch.device | str,
|
||||
steps: int | None = None,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
count = steps or self.inference_steps
|
||||
x = torch.randn(shape, device=device, generator=generator)
|
||||
grid = torch.linspace(0.0, 1.0, count + 1, device=device)
|
||||
for index in range(count):
|
||||
timestep = grid[index].expand(shape[0])
|
||||
x = x + (grid[index + 1] - grid[index]) * model_fn(x, timestep)
|
||||
return x
|
||||
@@ -0,0 +1,129 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .data import Pi05OverfitDataset
|
||||
from .flow import FlowMatching
|
||||
from .model import Pi05Config, Pi05Mini
|
||||
from .runtime import choose_device, write_json
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(checkpoint_path: str, manifest: str | None = None) -> 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"))
|
||||
data_manifest = Path(manifest) if manifest else root / "data/overfit/overfit_manifest.jsonl"
|
||||
dataset = Pi05OverfitDataset(
|
||||
data_manifest, image_size=cfg["data"]["image_size"],
|
||||
action_horizon=cfg["data"]["action_horizon"],
|
||||
action_statistics=checkpoint["action_statistics"]
|
||||
)
|
||||
model = Pi05Mini(Pi05Config(**checkpoint["model_config"]))
|
||||
model.load_state_dict(checkpoint["model"])
|
||||
model.to(device).eval()
|
||||
flow = FlowMatching(cfg["inference"]["flow_steps"])
|
||||
rows = []
|
||||
figure, axes = plt.subplots(1, len(dataset), figsize=(4 * len(dataset), 4), squeeze=False)
|
||||
for index, sample in enumerate(dataset):
|
||||
image = sample["image"][None].to(device)
|
||||
tokens = sample["prefix_tokens"][None].to(device)
|
||||
# Autoregressively predict semantic subtask followed by the complete
|
||||
# discretized action stream (ACTION_START, 32 scalar tokens, ACTION_END).
|
||||
zero_action = torch.zeros(1, dataset.action_horizon, 2, device=device)
|
||||
generated: list[int] = []
|
||||
expected_tokens = 1 + 2 + dataset.action_horizon * 2
|
||||
current_tokens = tokens
|
||||
for _ in range(expected_tokens):
|
||||
semantic = model(
|
||||
image, current_tokens, zero_action,
|
||||
torch.zeros(1, device=device),
|
||||
generated_token_start=tokens.shape[1]
|
||||
)
|
||||
next_token = int(semantic.text_logits[0, -1].argmax().cpu())
|
||||
generated.append(next_token)
|
||||
current_tokens = torch.cat(
|
||||
(current_tokens, torch.tensor([[next_token]], device=device)), dim=1
|
||||
)
|
||||
predicted_token = generated[0]
|
||||
vlm_action_tokens = torch.tensor(generated[1:], dtype=torch.long)
|
||||
try:
|
||||
vlm_normalized = dataset.tokenizer.decode_actions(vlm_action_tokens)
|
||||
vlm_xy = dataset.action_normalizer.decode(
|
||||
vlm_normalized, sample["history_xy"]
|
||||
).numpy()
|
||||
vlm_valid = True
|
||||
except ValueError:
|
||||
vlm_xy = np.full_like(sample["future_xy"].numpy(), np.nan)
|
||||
vlm_valid = False
|
||||
conditioned_tokens = torch.cat(
|
||||
(tokens, torch.tensor([[predicted_token]], device=device)), dim=1
|
||||
)
|
||||
def velocity(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
|
||||
return model(
|
||||
image, conditioned_tokens, x, t,
|
||||
generated_token_start=tokens.shape[1]
|
||||
).vector_field
|
||||
normalized = flow.sample(
|
||||
velocity, (1, dataset.action_horizon, 2), device,
|
||||
steps=cfg["inference"]["flow_steps"],
|
||||
generator=None,
|
||||
)[0].cpu()
|
||||
predicted_xy = dataset.action_normalizer.decode(
|
||||
normalized, sample["history_xy"]
|
||||
).numpy()
|
||||
target_xy = sample["future_xy"].numpy()
|
||||
diffusion_mse = float(np.mean((predicted_xy - target_xy) ** 2))
|
||||
vlm_mse = float(np.mean((vlm_xy - target_xy) ** 2)) if vlm_valid else None
|
||||
target_label = dataset.tokenizer.subtasks.decode(sample["subtask_target"])
|
||||
try:
|
||||
predicted_label = dataset.tokenizer.subtasks.decode(predicted_token)
|
||||
except ValueError:
|
||||
predicted_label = f"token_{predicted_token}"
|
||||
rows.append({
|
||||
"sample_id": sample["sample_id"], "target_subtask": target_label,
|
||||
"predicted_subtask": predicted_label,
|
||||
"vlm_action_tokens_valid": vlm_valid,
|
||||
"vlm_trajectory_mse": vlm_mse,
|
||||
"diffusion_trajectory_mse": diffusion_mse,
|
||||
})
|
||||
ax = axes[0, index]
|
||||
ax.plot(target_xy[:, 0], target_xy[:, 1], "o-", label="target")
|
||||
ax.plot(vlm_xy[:, 0], vlm_xy[:, 1], "s--", label="VLM tokens")
|
||||
ax.plot(predicted_xy[:, 0], predicted_xy[:, 1], "x--", label="diffusion")
|
||||
ax.scatter([0], [0], c="black", marker="s", label="ego")
|
||||
vlm_text = "invalid" if vlm_mse is None else f"{vlm_mse:.3f}"
|
||||
ax.set_title(
|
||||
f"{predicted_label}\nVLM={vlm_text}, diffusion={diffusion_mse:.3f}"
|
||||
)
|
||||
ax.set_xlabel("forward x (m)")
|
||||
ax.set_ylabel("left y (m)")
|
||||
ax.axis("equal")
|
||||
ax.grid(True, alpha=.3)
|
||||
ax.legend(fontsize=8)
|
||||
figure.tight_layout()
|
||||
output_dir = Path(checkpoint_path).parent
|
||||
plot_path = output_dir / "trajectory_comparison.png"
|
||||
figure.savefig(plot_path, dpi=160)
|
||||
plt.close(figure)
|
||||
write_json(output_dir / "evaluation.json", rows)
|
||||
print(rows)
|
||||
return plot_path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--checkpoint", default="outputs/overfit/checkpoint.pt")
|
||||
parser.add_argument("--manifest")
|
||||
args = parser.parse_args()
|
||||
evaluate(args.checkpoint, args.manifest)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,396 @@
|
||||
"""A compact, pure-PyTorch approximation of the pi0.5 dual-expert model.
|
||||
|
||||
The important architectural detail is that prefix and action tokens have
|
||||
different parameters at every Transformer layer. Prefix queries only attend
|
||||
to prefix keys, while action queries attend to the complete prefix and to all
|
||||
action tokens. This is deliberately implemented explicitly (instead of
|
||||
concatenating the streams into an ordinary Transformer) so the information
|
||||
flow is easy to inspect in a minimal reproduction.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
@dataclass
|
||||
class Pi05Config:
|
||||
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 = 64
|
||||
max_action_tokens: int = 32
|
||||
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")
|
||||
|
||||
@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 Pi05Output:
|
||||
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:
|
||||
# Accumulating the variance in fp32 avoids overflows under autocast.
|
||||
rms = x.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
normalized = x * torch.rsqrt(rms.to(x.dtype) + self.eps)
|
||||
return normalized * self.weight
|
||||
|
||||
|
||||
class AdaRMSNorm(nn.Module):
|
||||
"""RMSNorm modulated by a flow-time embedding."""
|
||||
|
||||
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)
|
||||
# Start as an ordinary RMSNorm, which makes tiny-data training stable.
|
||||
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)
|
||||
if fourier_dim % 2:
|
||||
fourier_dim += 1
|
||||
self.fourier_dim = fourier_dim
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(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, :]
|
||||
embedding = torch.cat((angles.sin(), angles.cos()), dim=-1)
|
||||
return self.mlp(embedding)
|
||||
|
||||
|
||||
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
|
||||
k_len = key.shape[1]
|
||||
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:
|
||||
# allowed is [B or 1, Q, K]; True means the edge exists.
|
||||
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 = torch.matmul(probabilities, v)
|
||||
return output.transpose(1, 2).contiguous().reshape(batch, q_len, width)
|
||||
|
||||
|
||||
class DualExpertBlock(nn.Module):
|
||||
def __init__(self, config: Pi05Config) -> None:
|
||||
super().__init__()
|
||||
d = config.d_model
|
||||
self.n_heads = config.n_heads
|
||||
self.dropout = config.dropout
|
||||
|
||||
# The two streams never share Transformer parameters.
|
||||
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
|
||||
# Prefix KV uses the VLM expert's parameters; action KV uses the action
|
||||
# expert's parameters. Thus action can read prefix without allowing a
|
||||
# reverse information path from action into the prefix stream.
|
||||
key = torch.cat((p.k(prefix_input), a.k(action_input)), dim=1)
|
||||
value = torch.cat((p.v(prefix_input), 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 Pi05Mini(nn.Module):
|
||||
"""Minimal pi0.5-style model for architecture validation/overfitting."""
|
||||
|
||||
def __init__(self, config: Optional[Pi05Config] = None) -> None:
|
||||
super().__init__()
|
||||
self.config = config or Pi05Config()
|
||||
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(DualExpertBlock(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)
|
||||
text_tokens = self.token_embedding(input_ids)
|
||||
# pi0.5 represents state in the VLM vocabulary. There is deliberately
|
||||
# no separate continuous-state projection in this model.
|
||||
prefix = torch.cat((patches, text_tokens), dim=1)
|
||||
if prefix.shape[1] > self.config.max_prefix_tokens:
|
||||
raise ValueError("prefix exceeds configured maximum length")
|
||||
prefix = prefix + self.prefix_position[:, : prefix.shape[1]]
|
||||
return prefix, patches.shape[1]
|
||||
|
||||
def _masks(
|
||||
self,
|
||||
batch: int,
|
||||
prefix_length: int,
|
||||
action_length: int,
|
||||
text_offset: int,
|
||||
attention_mask: Optional[Tensor],
|
||||
action_prefix_mask: Optional[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()
|
||||
# Image plus observation/instruction/state tokens form a fully
|
||||
# bidirectional observation block. Tokens generated after that block
|
||||
# (subtask and FAST actions) read the observation and earlier generated
|
||||
# tokens causally; observation queries cannot read generated tokens.
|
||||
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, :]
|
||||
# Padded prefix queries retain a harmless self edge to avoid all -inf
|
||||
# rows; their outputs are not used by the text loss.
|
||||
query_valid = valid[:, :, None]
|
||||
identity = torch.eye(prefix_length, dtype=torch.bool, device=device)[None]
|
||||
prefix_edges = torch.where(query_valid, prefix_edges, identity)
|
||||
|
||||
action_valid = valid
|
||||
if action_prefix_mask is not None:
|
||||
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")
|
||||
action_valid = action_valid & supplied
|
||||
prefix_keys = action_valid[:, None, :].expand(batch, action_length, prefix_length)
|
||||
action_keys = torch.ones(
|
||||
batch, action_length, action_length, dtype=torch.bool, device=device
|
||||
)
|
||||
action_edges = torch.cat((prefix_keys, action_keys), dim=-1)
|
||||
return prefix_edges, action_edges
|
||||
|
||||
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,
|
||||
) -> Pi05Output:
|
||||
"""Run both experts.
|
||||
|
||||
``timesteps`` is a scalar per batch item in [0, 1]. ``text_labels``
|
||||
corresponds only to ``input_ids`` (not image prefix positions)
|
||||
and may contain -100. ``target_velocity`` has the same shape as
|
||||
``noisy_actions``.
|
||||
"""
|
||||
if noisy_actions.ndim != 3:
|
||||
raise ValueError("noisy_actions must have shape [batch, horizon, action_dim]")
|
||||
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)
|
||||
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)
|
||||
# Return logits only for actual text positions, which keeps downstream
|
||||
# next-token loss/generation code uncomplicated.
|
||||
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)
|
||||
loss = None
|
||||
if text_loss is not None or flow_loss is not None:
|
||||
loss = sum(
|
||||
component for component 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 component is not None
|
||||
)
|
||||
return Pi05Output(
|
||||
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 pi-style conditional-flow training targets.
|
||||
|
||||
The path is x_t = (1-t) * noise + t * action, with constant target vector
|
||||
field ``action - noise``.
|
||||
"""
|
||||
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
|
||||
target_velocity = actions - noise
|
||||
return noisy_actions, timestep, target_velocity
|
||||
@@ -0,0 +1,36 @@
|
||||
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(requested: str = "auto") -> torch.device:
|
||||
if requested != "auto":
|
||||
return torch.device(requested)
|
||||
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,285 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CoordinateRange:
|
||||
minimum: float
|
||||
maximum: float
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.maximum <= self.minimum:
|
||||
raise ValueError("coordinate maximum must be greater than minimum")
|
||||
|
||||
|
||||
class CoordinateTokenizer:
|
||||
"""Uniformly discretize XY state values into the literal token range 0..255."""
|
||||
|
||||
bins = 256
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
x_range: tuple[float, float] = (-10.0, 10.0),
|
||||
y_range: tuple[float, float] = (-5.0, 5.0),
|
||||
) -> None:
|
||||
self.ranges = (
|
||||
CoordinateRange(*x_range),
|
||||
CoordinateRange(*y_range),
|
||||
)
|
||||
|
||||
def encode(self, coordinates: torch.Tensor | Iterable[Iterable[float]]) -> torch.Tensor:
|
||||
values = torch.as_tensor(coordinates, dtype=torch.float32)
|
||||
if values.shape[-1] != 2:
|
||||
raise ValueError("coordinates must end in an XY dimension of size 2")
|
||||
low = values.new_tensor([axis.minimum for axis in self.ranges])
|
||||
high = values.new_tensor([axis.maximum for axis in self.ranges])
|
||||
unit = ((values - low) / (high - low)).clamp(0.0, 1.0)
|
||||
return torch.round(unit * (self.bins - 1)).to(torch.long)
|
||||
|
||||
def decode(self, tokens: torch.Tensor | Iterable[Iterable[int]]) -> torch.Tensor:
|
||||
encoded = torch.as_tensor(tokens, dtype=torch.long)
|
||||
if encoded.shape[-1] != 2:
|
||||
raise ValueError("coordinate tokens must end in an XY dimension of size 2")
|
||||
if encoded.numel() and (encoded.min() < 0 or encoded.max() >= self.bins):
|
||||
raise ValueError("coordinate tokens must be in the inclusive range 0..255")
|
||||
low = torch.tensor([axis.minimum for axis in self.ranges], dtype=torch.float32)
|
||||
high = torch.tensor([axis.maximum for axis in self.ranges], dtype=torch.float32)
|
||||
return low + encoded.float() / (self.bins - 1) * (high - low)
|
||||
|
||||
|
||||
class SubtaskVocabulary:
|
||||
"""Small semantic vocabulary standing in for pi0.5's generated language subtask."""
|
||||
|
||||
labels = (
|
||||
"keep_lane",
|
||||
"steer_left",
|
||||
"steer_right",
|
||||
"slow_down",
|
||||
"stop",
|
||||
"accelerate",
|
||||
"unknown",
|
||||
)
|
||||
|
||||
def __init__(self, first_token_id: int = 260) -> None:
|
||||
self.first_token_id = int(first_token_id)
|
||||
self._label_to_token = {
|
||||
label: self.first_token_id + index for index, label in enumerate(self.labels)
|
||||
}
|
||||
self._token_to_label = {token: label for label, token in self._label_to_token.items()}
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
return len(self.labels)
|
||||
|
||||
@property
|
||||
def token_ids(self) -> tuple[int, ...]:
|
||||
return tuple(self._label_to_token[label] for label in self.labels)
|
||||
|
||||
def classify(self, reasoning: str) -> str:
|
||||
text = reasoning.casefold()
|
||||
# More specific commands precede generic lane/follow language.
|
||||
if any(term in text for term in ("stop", "halt", "stationary")):
|
||||
return "stop"
|
||||
if any(term in text for term in ("slow", "deceler", "reduce speed", "brake")):
|
||||
return "slow_down"
|
||||
if any(term in text for term in ("accelerat", "speed up")):
|
||||
return "accelerate"
|
||||
if any(term in text for term in ("left", "counterclockwise")):
|
||||
return "steer_left"
|
||||
if any(term in text for term in ("right", "clockwise")):
|
||||
return "steer_right"
|
||||
if any(term in text for term in ("lane", "straight", "follow", "continue", "proceed")):
|
||||
return "keep_lane"
|
||||
return "unknown"
|
||||
|
||||
def encode(self, label_or_reasoning: str) -> int:
|
||||
label = (
|
||||
label_or_reasoning
|
||||
if label_or_reasoning in self._label_to_token
|
||||
else self.classify(label_or_reasoning)
|
||||
)
|
||||
return self._label_to_token[label]
|
||||
|
||||
def decode(self, token_id: int | torch.Tensor) -> str:
|
||||
token = int(token_id)
|
||||
try:
|
||||
return self._token_to_label[token]
|
||||
except KeyError as error:
|
||||
raise ValueError(f"unknown subtask token: {token}") from error
|
||||
|
||||
|
||||
class MiniTokenizer:
|
||||
"""Compose image/state, semantic, and FAST-lite action tokens.
|
||||
|
||||
Actions are z-score-normalized by the dataset before reaching this
|
||||
tokenizer. A complete ``[horizon, 2]`` chunk is transformed along time
|
||||
with an orthonormal DCT-II. All coefficients are retained and serialized
|
||||
in frequency-major XY order, so low-frequency motion appears first in the
|
||||
token stream. The coefficients are divided by ``sqrt(horizon)`` before
|
||||
being uniformly quantized over ``[-4, 4]`` into 64 shared tokens. This
|
||||
keeps the fixed quantization interval valid for any horizon when the input
|
||||
actions have first been clipped to the same interval.
|
||||
|
||||
This intentionally implements only the central transform/order idea from
|
||||
FAST. It does not learn dataset-specific coefficient bins or drop trailing
|
||||
high-frequency coefficients.
|
||||
"""
|
||||
|
||||
BOS = 256
|
||||
IMAGE = 257
|
||||
STATE = 258
|
||||
SUBTASK_QUERY = 259
|
||||
ACTION_START = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
|
||||
ACTION_END = ACTION_START + 1
|
||||
ACTION_TOKEN_START = ACTION_END + 1
|
||||
ACTION_BINS = 64
|
||||
ACTION_TOKEN_END = ACTION_TOKEN_START + ACTION_BINS - 1
|
||||
ACTION_MINIMUM = -4.0
|
||||
ACTION_MAXIMUM = 4.0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
coordinate_tokenizer: CoordinateTokenizer | None = None,
|
||||
subtasks: SubtaskVocabulary | None = None,
|
||||
) -> None:
|
||||
self.coordinates = coordinate_tokenizer or CoordinateTokenizer()
|
||||
self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1)
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return max(max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1
|
||||
|
||||
@property
|
||||
def action_token_ids(self) -> range:
|
||||
"""Contiguous vocabulary interval reserved for quantized DCT coefficients."""
|
||||
|
||||
return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)
|
||||
|
||||
def build_prefix(self, history_xy: torch.Tensor) -> torch.Tensor:
|
||||
coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1)
|
||||
specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
|
||||
query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
|
||||
return torch.cat((specials, coordinate_tokens.cpu(), query), dim=0)
|
||||
|
||||
def encode_subtask(self, reasoning: str) -> torch.Tensor:
|
||||
return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long)
|
||||
|
||||
@staticmethod
|
||||
def _orthonormal_dct_matrix(
|
||||
horizon: int,
|
||||
*,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Return the orthonormal DCT-II matrix ``[frequency, time]``."""
|
||||
|
||||
if horizon <= 0:
|
||||
raise ValueError("DCT horizon must be positive")
|
||||
time = torch.arange(horizon, device=device, dtype=dtype)
|
||||
frequency = torch.arange(horizon, device=device, dtype=dtype).unsqueeze(1)
|
||||
matrix = torch.cos(math.pi / horizon * (time + 0.5) * frequency)
|
||||
matrix[0] *= math.sqrt(1.0 / horizon)
|
||||
if horizon > 1:
|
||||
matrix[1:] *= math.sqrt(2.0 / horizon)
|
||||
return matrix
|
||||
|
||||
def dct_actions(
|
||||
self, actions: torch.Tensor | Iterable[Iterable[float]]
|
||||
) -> torch.Tensor:
|
||||
"""Apply an orthonormal DCT-II along an action chunk's time axis."""
|
||||
|
||||
values = torch.as_tensor(actions, dtype=torch.float32)
|
||||
if values.ndim != 2 or values.shape[-1] != 2:
|
||||
raise ValueError("actions must have shape [horizon, 2]")
|
||||
if values.shape[0] == 0:
|
||||
raise ValueError("actions must contain at least one timestep")
|
||||
transform = self._orthonormal_dct_matrix(
|
||||
values.shape[0], device=values.device, dtype=values.dtype
|
||||
)
|
||||
return transform @ values
|
||||
|
||||
def idct_actions(
|
||||
self, coefficients: torch.Tensor | Iterable[Iterable[float]]
|
||||
) -> torch.Tensor:
|
||||
"""Invert :meth:`dct_actions` using the transpose of its basis."""
|
||||
|
||||
values = torch.as_tensor(coefficients, dtype=torch.float32)
|
||||
if values.ndim != 2 or values.shape[-1] != 2:
|
||||
raise ValueError("DCT coefficients must have shape [horizon, 2]")
|
||||
if values.shape[0] == 0:
|
||||
raise ValueError("DCT coefficients must contain at least one frequency")
|
||||
transform = self._orthonormal_dct_matrix(
|
||||
values.shape[0], device=values.device, dtype=values.dtype
|
||||
)
|
||||
return transform.transpose(0, 1) @ values
|
||||
|
||||
def encode_actions(
|
||||
self, actions: torch.Tensor | Iterable[Iterable[float]]
|
||||
) -> torch.Tensor:
|
||||
"""DCT-transform and encode normalized actions as a marked token stream."""
|
||||
|
||||
values = torch.as_tensor(actions, dtype=torch.float32)
|
||||
if values.ndim != 2 or values.shape[-1] != 2:
|
||||
raise ValueError("actions must have shape [horizon, 2]")
|
||||
if values.shape[0] == 0:
|
||||
raise ValueError("actions must contain at least one timestep")
|
||||
values = values.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
|
||||
coefficients = self.dct_actions(values) / math.sqrt(values.shape[0])
|
||||
# Exact zero is useful for a stable shared center bin. Tiny non-zero
|
||||
# values here are only floating-point residue from the DCT basis.
|
||||
coefficients = torch.where(
|
||||
coefficients.abs() < 1e-6,
|
||||
torch.zeros_like(coefficients),
|
||||
coefficients,
|
||||
)
|
||||
unit = (
|
||||
(
|
||||
coefficients.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
|
||||
- self.ACTION_MINIMUM
|
||||
)
|
||||
/ (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
|
||||
)
|
||||
quantized = torch.round(unit * (self.ACTION_BINS - 1)).to(torch.long)
|
||||
# Row-major flattening is frequency-major: DC XY, frequency-1 XY, ...
|
||||
action_tokens = quantized.reshape(-1) + self.ACTION_TOKEN_START
|
||||
boundaries = torch.tensor(
|
||||
[self.ACTION_START, self.ACTION_END],
|
||||
dtype=torch.long,
|
||||
device=action_tokens.device,
|
||||
)
|
||||
return torch.cat((boundaries[:1], action_tokens, boundaries[1:]))
|
||||
|
||||
def decode_actions(self, tokens: torch.Tensor | Iterable[int]) -> torch.Tensor:
|
||||
"""Decode FAST-lite tokens back to normalized ``[horizon, 2]`` actions."""
|
||||
|
||||
encoded = torch.as_tensor(tokens, dtype=torch.long)
|
||||
if encoded.ndim != 1:
|
||||
raise ValueError("action tokens must be a one-dimensional sequence")
|
||||
if (
|
||||
encoded.numel() < 4
|
||||
or encoded[0] != self.ACTION_START
|
||||
or encoded[-1] != self.ACTION_END
|
||||
):
|
||||
raise ValueError("action tokens must be enclosed by ACTION_START and ACTION_END")
|
||||
payload = encoded[1:-1]
|
||||
if payload.numel() % 2:
|
||||
raise ValueError("action token payload must contain complete XY pairs")
|
||||
if payload.min() < self.ACTION_TOKEN_START or payload.max() > self.ACTION_TOKEN_END:
|
||||
raise ValueError(
|
||||
"action token payload is outside the reserved action token interval"
|
||||
)
|
||||
quantized = payload - self.ACTION_TOKEN_START
|
||||
normalized_coefficients = (
|
||||
quantized.float()
|
||||
/ (self.ACTION_BINS - 1)
|
||||
* (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
|
||||
+ self.ACTION_MINIMUM
|
||||
)
|
||||
horizon = payload.numel() // 2
|
||||
coefficients = normalized_coefficients.reshape(horizon, 2) * math.sqrt(horizon)
|
||||
return self.idct_actions(coefficients)
|
||||
@@ -0,0 +1,193 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from .data import Pi05OverfitDataset, materialize_subset
|
||||
from .flow import FlowMatching
|
||||
from .model import Pi05Config, Pi05Mini
|
||||
from .runtime import choose_device, load_config, seed_everything, write_json
|
||||
|
||||
|
||||
def _collate(rows: list[dict]) -> dict:
|
||||
tensor_keys = (
|
||||
"image", "prefix_tokens", "subtask_target", "actions", "action_tokens"
|
||||
)
|
||||
result = {key: torch.stack([row[key] for row in rows]) for key in tensor_keys}
|
||||
result["sample_id"] = [row["sample_id"] for row in rows]
|
||||
return result
|
||||
|
||||
|
||||
def build_model(config: dict, dataset: Pi05OverfitDataset) -> Pi05Mini:
|
||||
model_cfg = config["model"]
|
||||
return Pi05Mini(Pi05Config(
|
||||
vocab_size=dataset.tokenizer.vocab_size,
|
||||
action_dim=2,
|
||||
image_size=config["data"]["image_size"],
|
||||
patch_size=model_cfg["patch_size"],
|
||||
d_model=model_cfg["d_model"],
|
||||
n_heads=model_cfg["n_heads"],
|
||||
n_layers=model_cfg["n_layers"],
|
||||
mlp_ratio=model_cfg["mlp_ratio"],
|
||||
max_text_tokens=model_cfg["max_text_tokens"],
|
||||
max_action_tokens=config["data"]["action_horizon"],
|
||||
dropout=model_cfg["dropout"],
|
||||
))
|
||||
|
||||
|
||||
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(
|
||||
source_manifest=cfg["data"]["source_manifest"],
|
||||
output_dir=root / "data" / "overfit",
|
||||
num_variants=cfg["data"]["num_variants"],
|
||||
)
|
||||
dataset = Pi05OverfitDataset(
|
||||
manifest, image_size=cfg["data"]["image_size"],
|
||||
action_horizon=cfg["data"]["action_horizon"]
|
||||
)
|
||||
loader = DataLoader(
|
||||
dataset, batch_size=min(cfg["train"]["batch_size"], len(dataset)),
|
||||
shuffle=True, collate_fn=_collate
|
||||
)
|
||||
model = build_model(cfg, dataset).to(device)
|
||||
flow = FlowMatching(cfg["inference"]["flow_steps"])
|
||||
output_dir = root / cfg["output_dir"]
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
history = []
|
||||
|
||||
def is_action_parameter(name: str) -> bool:
|
||||
return any(part in name for part in (
|
||||
"action_", "time_embedding", "vector_field_head"
|
||||
))
|
||||
|
||||
# Stage 1 follows pi0.5 pre-training: train the VLM/token pathway with
|
||||
# discrete state and FAST-lite actions, while the flow expert is untouched.
|
||||
stage1_params = [p for name, p in model.named_parameters() if not is_action_parameter(name)]
|
||||
optimizer = torch.optim.AdamW(
|
||||
stage1_params, lr=cfg["train"]["learning_rate"],
|
||||
weight_decay=cfg["train"]["weight_decay"]
|
||||
)
|
||||
|
||||
def next_batch(iterator):
|
||||
try:
|
||||
return next(iterator), iterator
|
||||
except StopIteration:
|
||||
iterator = iter(loader)
|
||||
return next(iterator), iterator
|
||||
|
||||
iterator = iter(loader)
|
||||
model.train()
|
||||
for step in range(1, cfg["train"]["pretrain_steps"] + 1):
|
||||
batch, iterator = next_batch(iterator)
|
||||
image = batch["image"].to(device)
|
||||
tokens = batch["prefix_tokens"].to(device)
|
||||
actions = batch["actions"].to(device)
|
||||
subtask = batch["subtask_target"].to(device)[:, None]
|
||||
action_tokens = batch["action_tokens"].to(device)
|
||||
targets = torch.cat((subtask, action_tokens), dim=1)
|
||||
# Teacher-forced autoregressive stream. The final target is predicted
|
||||
# from the preceding token and therefore is not part of model_tokens.
|
||||
model_tokens = torch.cat((tokens, targets[:, :-1]), dim=1)
|
||||
model_labels = torch.full_like(model_tokens, -100)
|
||||
query_index = tokens.shape[1] - 1
|
||||
model_labels[:, query_index:] = targets
|
||||
# Flow tokens may read observation/state and the generated semantic
|
||||
# subtask, but not teacher-forced discrete action tokens.
|
||||
zero_action = torch.zeros_like(actions)
|
||||
output = model(
|
||||
image, model_tokens, zero_action,
|
||||
torch.zeros(image.shape[0], device=device),
|
||||
generated_token_start=tokens.shape[1],
|
||||
text_labels=model_labels,
|
||||
)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
output.loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"])
|
||||
optimizer.step()
|
||||
if step == 1 or step % cfg["train"]["log_every"] == 0:
|
||||
record = {
|
||||
"stage": "pretrain",
|
||||
"step": step,
|
||||
"loss": float(output.loss.detach().cpu()),
|
||||
"text_loss": float(output.text_loss.detach().cpu()),
|
||||
"flow_loss": None,
|
||||
}
|
||||
history.append(record)
|
||||
print(json.dumps(record))
|
||||
|
||||
torch.save(model.state_dict(), output_dir / "stage1_vlm.pt")
|
||||
|
||||
# Stage 2 follows post-training: add the still-random action expert and
|
||||
# jointly retain token prediction while optimizing continuous flow.
|
||||
optimizer = torch.optim.AdamW(
|
||||
model.parameters(), lr=cfg["train"]["learning_rate"],
|
||||
weight_decay=cfg["train"]["weight_decay"]
|
||||
)
|
||||
iterator = iter(loader)
|
||||
for step in range(1, cfg["train"]["posttrain_steps"] + 1):
|
||||
batch, iterator = next_batch(iterator)
|
||||
image = batch["image"].to(device)
|
||||
tokens = batch["prefix_tokens"].to(device)
|
||||
actions = batch["actions"].to(device)
|
||||
subtask = batch["subtask_target"].to(device)[:, None]
|
||||
action_tokens = batch["action_tokens"].to(device)
|
||||
targets = torch.cat((subtask, action_tokens), dim=1)
|
||||
model_tokens = torch.cat((tokens, targets[:, :-1]), dim=1)
|
||||
model_labels = torch.full_like(model_tokens, -100)
|
||||
model_labels[:, tokens.shape[1] - 1:] = targets
|
||||
action_prefix_mask = torch.zeros_like(model_tokens, dtype=torch.bool)
|
||||
action_prefix_mask[:, : tokens.shape[1] + 1] = True
|
||||
noisy, timestep, velocity = flow.training_batch(actions)
|
||||
output = model(
|
||||
image, model_tokens, noisy, timestep,
|
||||
action_prefix_mask=action_prefix_mask,
|
||||
generated_token_start=tokens.shape[1],
|
||||
text_labels=model_labels, target_velocity=velocity,
|
||||
text_loss_weight=cfg["train"]["semantic_loss_weight"],
|
||||
flow_loss_weight=cfg["train"]["flow_loss_weight"],
|
||||
)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
output.loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"])
|
||||
optimizer.step()
|
||||
if step == 1 or step % cfg["train"]["log_every"] == 0:
|
||||
record = {
|
||||
"stage": "posttrain", "step": step,
|
||||
"loss": float(output.loss.detach().cpu()),
|
||||
"text_loss": float(output.text_loss.detach().cpu()),
|
||||
"flow_loss": float(output.flow_loss.detach().cpu()),
|
||||
}
|
||||
history.append(record)
|
||||
print(json.dumps(record))
|
||||
checkpoint = {
|
||||
"model": model.state_dict(), "config": cfg,
|
||||
"model_config": model.config.__dict__,
|
||||
"action_statistics": dataset.action_normalizer.as_dict(),
|
||||
}
|
||||
torch.save(checkpoint, output_dir / "checkpoint.pt")
|
||||
write_json(output_dir / "training_history.json", history)
|
||||
write_json(output_dir / "summary.json", {
|
||||
"device": str(device), "samples": len(dataset),
|
||||
"parameters": sum(p.numel() for p in model.parameters()),
|
||||
"initial": history[0], "final": history[-1],
|
||||
})
|
||||
return output_dir / "checkpoint.pt"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", default="configs/overfit.yaml")
|
||||
args = parser.parse_args()
|
||||
train(args.config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user