commit 4fe0099f3f24b99a6ffa2db535efc69368b87735 Author: seeker Date: Sat Jul 11 20:04:18 2026 +0800 Initial pi0 minimal validation diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ff92181 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.ruff_cache/ +*.egg-info/ +.venv/ +data/ +outputs/**/*.pt +.DS_Store diff --git a/README.md b/README.md new file mode 100644 index 0000000..beba5ec --- /dev/null +++ b/README.md @@ -0,0 +1,76 @@ +# π0 Minimal Validation + +这是一个 **Physical Intelligence π0 风格的最小、可运行过拟合实验**。它用小型视觉语言前缀和 action expert 学习连续动作 chunk，通过 flow matching 从噪声一次生成未来 `16×2` 轨迹。 + +项目不是官方代码，也不复制官方规模、数据和机器人能力；目标只是验证核心数学与信息流能否端到端训练。 + +## 最小架构 + +```text +图像 + 语言指令 → 小型 VLM prefix + │ blockwise attention +连续本体/历史 state ─────┤ +noisy action + timestep ─┘ + ▼ + action expert + │ + ▼ + flow velocity → action chunk +``` + +默认从相邻 Alpamayo 项目读取 1 个 PAI 驾驶样本，生成 4 个确定性轻微增强变体，输入 `64×64` 图像，预测未来 16 个 XY 动作点。 + +## π0 与 π0.5 的关键差异 + +| 方面 | 本项目对应的 π0 | π0.5 最小复现中的额外能力 | +|---|---|---| +| 控制层级 | 语言/视觉条件直接生成连续 action chunk | 先生成高层 subtask，再以 subtask 为条件生成动作 | +| 动作表示 | flow action expert 输出连续动作 | 可同时有 FAST 风格离散动作 token 和连续 flow 动作 | +| VLM 输出 | 主要作为条件语义前缀 | 自回归预测 subtask/离散动作 token | +| 训练目标 | 本实验仅使用 flow velocity MSE | 还可联合文本/动作 token CE | + +与 π0 一致，本项目把连续 state 放在 action suffix，而不是像 π0.5 那样先离散成语言 token。 + +因此，本项目故意不实现 π0.5 的高层 subtask 生成和 VLM 离散轨迹分支。效果图将对比 flow 生成轨迹与真值，而不是三方动作对比。 + +## 安装与测试 + +```bash +cd pi0-minimal-validation +python -m venv .venv +source .venv/bin/activate +python -m pip install -U pip +python -m pip install -e '.[dev]' +pytest -q +``` + +`tests/test_flow.py` 验证直线概率路径 `x_t=(1-t)·noise+t·action`、目标速度 `action-noise`、MSE 损失、随机可复现性和 Euler 积分。 + +## 过拟合训练与效果图 + +```bash +pi0mini-train --config configs/overfit.yaml +pi0mini-infer --checkpoint outputs/overfit/checkpoint.pt +``` + +也可用 `python -m pi0mini.train --config configs/overfit.yaml` 运行。默认训练 600 步，推理使用 10 步 Euler flow integration。训练后应生成 checkpoint、loss 历史、逐样本评估 JSON 和 `outputs/overfit/trajectory_comparison.png`。 + +默认数据源为： + +```text +../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl +``` + +如本地路径不同，修改 `configs/overfit.yaml` 中的 `data.source_manifest`。 + +## 限制 + +- 小型自定义编码器代替预训练 PaliGemma，参数规模与官方模型不可比。 +- 仅对几个增强样本做记忆性训练，loss 下降不代表泛化、鲁棒性或安全性。 +- 二维驾驶轨迹替代多机器人、多相机、高维关节动作和异构数据联合训练。 +- 固定步长 Euler sampler 只用于演示，没有实时控制、碰撞检查、约束或故障降级。 +- 严禁将输出用于真实车辆、机械臂或其他会产生物理影响的设备。 + +## 成功标准 + +成功只表示：该缩小实现的视觉/语言条件、blockwise action expert 和 flow matching 能够端到端反向传播，并记住极小训练集。它不证明与官方 π0 数值等价。 diff --git a/configs/overfit.yaml b/configs/overfit.yaml new file mode 100644 index 0000000..013dae1 --- /dev/null +++ b/configs/overfit.yaml @@ -0,0 +1,31 @@ +seed: 7 +device: auto +output_dir: outputs/overfit + +data: + source_manifest: ../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl + num_variants: 4 + image_size: 64 + action_horizon: 16 + +model: + d_model: 96 + n_heads: 4 + n_layers: 3 + mlp_ratio: 3 + patch_size: 16 + max_text_tokens: 64 + dropout: 0.0 + +train: + steps: 600 + batch_size: 4 + learning_rate: 0.001 + weight_decay: 0.0 + grad_clip_norm: 1.0 + flow_loss_weight: 1.0 + log_every: 20 + +inference: + flow_steps: 10 + noise_seed: 4 diff --git a/outputs/overfit/evaluation.json b/outputs/overfit/evaluation.json new file mode 100644 index 0000000..dbb532a --- /dev/null +++ b/outputs/overfit/evaluation.json @@ -0,0 +1,18 @@ +[ + { + "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug0", + "trajectory_mse": 0.1524142026901245 + }, + { + "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug1", + "trajectory_mse": 0.2658033072948456 + }, + { + "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug2", + "trajectory_mse": 0.011374758556485176 + }, + { + "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug3", + "trajectory_mse": 0.2792850136756897 + } +] \ No newline at end of file diff --git a/outputs/overfit/summary.json b/outputs/overfit/summary.json new file mode 100644 index 0000000..6f0cd69 --- /dev/null +++ b/outputs/overfit/summary.json @@ -0,0 +1,13 @@ +{ + "device": "mps", + "samples": 4, + "parameters": 662786, + "initial": { + "step": 1, + "flow_loss": 2.2617573738098145 + }, + "final": { + "step": 600, + "flow_loss": 0.16418831050395966 + } +} \ No newline at end of file diff --git a/outputs/overfit/training_history.json b/outputs/overfit/training_history.json new file mode 100644 index 0000000..89bc35c --- /dev/null +++ b/outputs/overfit/training_history.json @@ -0,0 +1,126 @@ +[ + { + "step": 1, + "flow_loss": 2.2617573738098145 + }, + { + "step": 20, + "flow_loss": 1.941188097000122 + }, + { + "step": 40, + "flow_loss": 1.479238510131836 + }, + { + "step": 60, + "flow_loss": 0.7427206635475159 + }, + { + "step": 80, + "flow_loss": 0.6053631901741028 + }, + { + "step": 100, + "flow_loss": 0.270697683095932 + }, + { + "step": 120, + "flow_loss": 0.11459995806217194 + }, + { + "step": 140, + "flow_loss": 0.4224633276462555 + }, + { + "step": 160, + "flow_loss": 0.14150527119636536 + }, + { + "step": 180, + "flow_loss": 0.15842989087104797 + }, + { + "step": 200, + "flow_loss": 0.34323638677597046 + }, + { + "step": 220, + "flow_loss": 0.1762494295835495 + }, + { + "step": 240, + "flow_loss": 0.3678283989429474 + }, + { + "step": 260, + "flow_loss": 0.10764116048812866 + }, + { + "step": 280, + "flow_loss": 0.17170406877994537 + }, + { + "step": 300, + "flow_loss": 0.20748832821846008 + }, + { + "step": 320, + "flow_loss": 0.12005910277366638 + }, + { + "step": 340, + "flow_loss": 0.1855630874633789 + }, + { + "step": 360, + "flow_loss": 0.16064667701721191 + }, + { + "step": 380, + "flow_loss": 0.11575131863355637 + }, + { + "step": 400, + "flow_loss": 0.1507677137851715 + }, + { + "step": 420, + "flow_loss": 0.23872917890548706 + }, + { + "step": 440, + "flow_loss": 0.25908488035202026 + }, + { + "step": 460, + "flow_loss": 0.15245740115642548 + }, + { + "step": 480, + "flow_loss": 0.3529265224933624 + }, + { + "step": 500, + "flow_loss": 0.22890940308570862 + }, + { + "step": 520, + "flow_loss": 0.18230678141117096 + }, + { + "step": 540, + "flow_loss": 0.14777007699012756 + }, + { + "step": 560, + "flow_loss": 0.06152040883898735 + }, + { + "step": 580, + "flow_loss": 0.18225648999214172 + }, + { + "step": 600, + "flow_loss": 0.16418831050395966 + } +] \ No newline at end of file diff --git a/outputs/overfit/trajectory_comparison.png b/outputs/overfit/trajectory_comparison.png new file mode 100644 index 0000000..2e64cc2 Binary files /dev/null and b/outputs/overfit/trajectory_comparison.png differ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..20a1f04 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,35 @@ +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "pi0-minimal-validation" +version = "0.1.0" +description = "A tiny educational pi0-style VLA flow-matching overfit experiment" +readme = "README.md" +requires-python = ">=3.10" +license = { text = "MIT" } +dependencies = [ + "matplotlib>=3.8", + "numpy>=1.26", + "Pillow>=10.0", + "PyYAML>=6.0", + "torch>=2.2", +] + +[project.optional-dependencies] +dev = ["pytest>=8.0"] + +[project.scripts] +pi0mini-train = "pi0mini.train:main" +pi0mini-infer = "pi0mini.infer:main" + +[tool.setuptools] +package-dir = { "" = "src" } + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-ra" diff --git a/src/pi0mini/__init__.py b/src/pi0mini/__init__.py new file mode 100644 index 0000000..d1a5b31 --- /dev/null +++ b/src/pi0mini/__init__.py @@ -0,0 +1,5 @@ +"""Minimal original-pi0 architecture validation package.""" + +from .model import Pi0Config, Pi0Mini, Pi0Output, flow_matching_sample + +__all__ = ["Pi0Config", "Pi0Mini", "Pi0Output", "flow_matching_sample"] diff --git a/src/pi0mini/data.py b/src/pi0mini/data.py new file mode 100644 index 0000000..1bcace6 --- /dev/null +++ b/src/pi0mini/data.py @@ -0,0 +1,310 @@ +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 + + +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 + + @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 2-D 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) + # Tiny overfit subsets may have a nearly constant forward velocity. A + # floor prevents an ill-conditioned target without harming round trips. + 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_xy = torch.as_tensor(history_xy, dtype=torch.float32) + future_xy = torch.as_tensor(future_xy, dtype=torch.float32) + if history_xy.ndim != 2 or history_xy.shape[-1] != 2 or len(history_xy) < 1: + raise ValueError("history_xy must have shape [history_steps, 2]") + if future_xy.ndim != 2 or future_xy.shape[-1] != 2 or len(future_xy) < 1: + raise ValueError("future_xy must have shape [action_horizon, 2]") + previous = torch.cat((history_xy[-1:], future_xy[:-1]), dim=0) + return future_xy - 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") + deltas = [ + cls.to_deltas(history, future).reshape(-1, 2) + for history, future in zip(histories, futures, strict=True) + ] + flat = torch.cat(deltas, 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_xy = torch.as_tensor(history_xy, dtype=torch.float32) + if actions.ndim != 2 or actions.shape[-1] != 2: + raise ValueError("actions must have shape [action_horizon, 2]") + if history_xy.ndim != 2 or history_xy.shape[-1] != 2 or len(history_xy) < 1: + raise ValueError("history_xy must have shape [history_steps, 2]") + deltas = actions * self.scale.to(actions.device) + self.mean.to(actions.device) + return history_xy[-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"]), + ) + ) + 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 one self-contained PAI sample from the neighboring Alpamayo project.""" + + if max_samples != 1: + raise ValueError("the pi0 minimal validation intentionally materializes exactly one sample") + records = load_manifest(source_manifest)[:1] + 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) + + record = records[0] + stem = f"sample_00_{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) + row = { + "clip_id": record.clip_id, + "t0_us": record.t0_us, + "image_path": str(image_path.relative_to(destination)), + "trajectory_path": str(trajectory_path.relative_to(destination)), + "history_steps": 8, + "action_horizon": 16, + } + manifest_path = destination / "overfit_manifest.jsonl" + manifest_path.write_text(json.dumps(row) + "\n", encoding="utf-8") + return manifest_path + + +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 + signal = xy.T.unsqueeze(0) + return F.interpolate(signal, size=steps, mode="linear", align_corners=True)[0].T + + +class Pi0OverfitDataset(Dataset[dict[str, Any]]): + """One PAI sample with four deterministic variants for pi0 overfitting.""" + + def __init__( + self, + manifest_path: str | Path, + *, + image_size: int = 64, + augmentation_factor: int = 4, + history_steps: int = 8, + action_horizon: int = 16, + seed: int = 7, + action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None, + ) -> None: + if image_size < 1: + raise ValueError("image_size must be positive") + if augmentation_factor < 1: + raise ValueError("augmentation_factor must be at least one") + if history_steps != 8: + raise ValueError("history_steps must be 8 so continuous state has dimension 16") + if action_horizon != 16: + raise ValueError("action_horizon must be 16 for this minimal validation") + 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 + + if action_statistics is None: + geometries = [ + self._load_geometry(base_index, variant) + for base_index in range(len(self.records)) + for variant in range(self.augmentation_factor) + ] + self.action_normalizer = ActionNormalizer.fit( + [history for history, _ in geometries], + [future for _, future in geometries], + ) + 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.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_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] + history = _resample_xy(history, self.history_steps) + future = _resample_xy(future, 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) + progress = torch.linspace(0.0, 1.0, self.action_horizon, dtype=future.dtype) + future[:, 1] += float(rng.normal(0.0, 0.012)) * 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) + image = ImageEnhance.Brightness(image).enhance(float(rng.uniform(0.97, 1.03))) + image = ImageEnhance.Contrast(image).enhance(float(rng.uniform(0.98, 1.02))) + 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]: + base_index, variant = self._indices(index) + record = self.records[base_index] + history_xy, future_xy = self._load_geometry(base_index, variant) + actions = self.action_normalizer.encode(history_xy, future_xy) + return { + "image": self._load_image(base_index, variant), + "state": history_xy.reshape(-1), + "actions": actions, + "history_xy": history_xy, + "future_xy": future_xy, + "sample_id": f"{record.sample_id}:aug{variant}", + } diff --git a/src/pi0mini/flow.py b/src/pi0mini/flow.py new file mode 100644 index 0000000..18a4121 --- /dev/null +++ b/src/pi0mini/flow.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from collections.abc import Callable +import torch +from torch import Tensor + + +class FlowMatching: + def __init__(self, inference_steps: int = 10, time_cutoff: float = 0.999) -> None: + self.inference_steps = inference_steps + if not 0.0 < time_cutoff <= 1.0: + raise ValueError("time_cutoff must be in (0, 1]") + self.time_cutoff = time_cutoff + + def sample_training_times( + self, + batch: int, + *, + device: torch.device | str, + dtype: torch.dtype, + generator: torch.Generator | None = None, + ) -> Tensor: + """Sample pi0's shifted Beta(1.5, 1) flow times. + + The paper parameterizes time from clean action toward noise and uses + ``0.001 + 0.999 * Beta(1.5, 1)``. This project parameterizes the same + path in the reverse direction, from noise at tau=0 to action at tau=1, + hence ``tau = 0.999 * (1 - Beta(1.5, 1))``. Beta(a, 1) can be sampled + as ``U ** (1/a)``, which preserves ``torch.Generator`` support. + """ + uniform = torch.rand( + batch, device=device, dtype=dtype, generator=generator + ) + beta = uniform.pow(1.0 / 1.5) + return self.time_cutoff * (1.0 - beta) + + def training_batch( + self, target: Tensor, generator: torch.Generator | None = None + ) -> tuple[Tensor, Tensor, Tensor]: + batch = target.shape[0] + t = self.sample_training_times( + batch, + device=target.device, + dtype=target.dtype, + generator=generator, + ) + noise = torch.randn( + target.shape, device=target.device, dtype=target.dtype, generator=generator + ) + noisy = t[:, None, None] * target + (1 - t[:, None, None]) * noise + return noisy, t, 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, 1, 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 diff --git a/src/pi0mini/infer.py b/src/pi0mini/infer.py new file mode 100644 index 0000000..e22f9ed --- /dev/null +++ b/src/pi0mini/infer.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import argparse +from pathlib import Path +import matplotlib.pyplot as plt +import numpy as np +import torch + +from .data import Pi0OverfitDataset +from .flow import FlowMatching +from .model import Pi0Config, Pi0Mini +from .runtime import choose_device, write_json + + +@torch.no_grad() +def evaluate(checkpoint_path: str) -> 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")) + dataset = Pi0OverfitDataset( + root / "data/overfit/overfit_manifest.jsonl", + image_size=cfg["data"]["image_size"], + augmentation_factor=cfg["data"]["num_variants"], + action_horizon=cfg["data"]["action_horizon"], + action_statistics=checkpoint["action_statistics"] + ) + model = Pi0Mini(Pi0Config(**checkpoint["model_config"])) + model.load_state_dict(checkpoint["model"]) + model.to(device).eval() + flow = FlowMatching(cfg["inference"]["flow_steps"]) + figure, axes = plt.subplots(1, len(dataset), figsize=(4 * len(dataset), 4), squeeze=False) + metrics = [] + for index, sample in enumerate(dataset): + image = sample["image"][None].to(device) + state = sample["state"][None].to(device) + input_ids = torch.tensor([[1, 2]], device=device) + def velocity(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + return model(image, input_ids, state, x, t).vector_field + # Fixed noise makes the tiny overfit demonstration reproducible. + torch.manual_seed(cfg["inference"].get("noise_seed", cfg["seed"])) + normalized = flow.sample( + velocity, (1, dataset.action_horizon, 2), device, + cfg["inference"]["flow_steps"] + )[0].cpu() + predicted = dataset.action_normalizer.decode(normalized, sample["history_xy"]).numpy() + target = sample["future_xy"].numpy() + mse = float(np.mean((predicted - target) ** 2)) + metrics.append({"sample_id": sample["sample_id"], "trajectory_mse": mse}) + ax = axes[0, index] + ax.plot(target[:, 0], target[:, 1], "o-", label="target") + ax.plot(predicted[:, 0], predicted[:, 1], "x--", label="pi0 flow") + ax.scatter([0], [0], c="black", marker="s", label="ego") + ax.set_title(f"π0 flow MSE={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 = output_dir / "trajectory_comparison.png" + figure.savefig(plot, dpi=160) + plt.close(figure) + write_json(output_dir / "evaluation.json", metrics) + print(metrics) + return plot + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", default="outputs/overfit/checkpoint.pt") + args = parser.parse_args() + evaluate(args.checkpoint) + + +if __name__ == "__main__": + main() diff --git a/src/pi0mini/model.py b/src/pi0mini/model.py new file mode 100644 index 0000000..77aa261 --- /dev/null +++ b/src/pi0mini/model.py @@ -0,0 +1,458 @@ +"""A compact, pure-PyTorch approximation of the original pi0 model. + +This module keeps the architectural choices that distinguish pi0 from a +plain diffusion Transformer: + +* image patches and language tokens form the VLM prefix; +* continuous proprioception and noisy action tokens form distinct state and + action blocks inside the action-expert stream; +* the VLM and action experts have independent parameters at every layer; +* state queries can read the prefix and state, while action queries can read + the prefix, state, and all action tokens; +* noisy actions are concatenated with a flow-time embedding before an MLP; +* flow time is *not* injected with adaptive normalization (the pi0.5 change). + +The implementation is intentionally small and dependency-free so it can be +overfit on a few examples as an architecture/optimization smoke test. It is +not intended to reproduce the scale or robotics capability of pi0. +""" + +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 Pi0Config: + """Configuration for :class:`Pi0Mini`.""" + + vocab_size: int = 512 + action_dim: int = 2 + state_dim: int = 8 + 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 + time_feature_dim: int = 64 + dropout: float = 0.0 + rms_eps: float = 1e-6 + # The observation/language prefix is bidirectional in pi0. Keeping this + # configurable is useful only for small ablations. + prefix_causal: bool = False + + 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.time_feature_dim < 2 or self.time_feature_dim % 2: + raise ValueError("time_feature_dim must be a positive even number") + + @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 + + @property + def max_suffix_tokens(self) -> int: + # One continuous proprioceptive-state token precedes action tokens. + return 1 + self.max_action_tokens + + +@dataclass +class Pi0Output: + """Outputs of the two-expert network.""" + + vector_field: Tensor + prefix_hidden: Tensor + suffix_hidden: Tensor + state_hidden: Tensor + action_hidden: Tensor + 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: + # Accumulate in fp32 for stable mixed-precision training. + mean_square = x.float().pow(2).mean(dim=-1, keepdim=True) + normalized = x * torch.rsqrt(mean_square.to(x.dtype) + self.eps) + return normalized * self.weight + + +class FourierTimeFeatures(nn.Module): + """Fixed sinusoidal features for a scalar flow time in [0, 1].""" + + def __init__(self, dim: int) -> None: + super().__init__() + half = dim // 2 + frequencies = torch.exp(torch.linspace(math.log(1.0), math.log(1000.0), half)) + self.register_buffer("frequencies", frequencies, persistent=False) + + def forward(self, timesteps: Tensor) -> Tensor: + timesteps = timesteps.reshape(-1).float() + angles = 2 * math.pi * timesteps[:, None] * self.frequencies[None, :] + return torch.cat((angles.sin(), angles.cos()), dim=-1) + + +class ActionTimeProjection(nn.Module): + """pi0-style joint embedding of noisy action values and flow time.""" + + def __init__(self, action_dim: int, time_dim: int, model_dim: int) -> None: + super().__init__() + self.time_features = FourierTimeFeatures(time_dim) + self.mlp = nn.Sequential( + nn.Linear(action_dim + time_dim, model_dim), + nn.SiLU(), + nn.Linear(model_dim, model_dim), + ) + + def forward(self, noisy_actions: Tensor, timesteps: Tensor) -> Tensor: + time = self.time_features(timesteps).to( + device=noisy_actions.device, dtype=noisy_actions.dtype + ) + time = time[:, None, :].expand(-1, noisy_actions.shape[1], -1) + return self.mlp(torch.cat((noisy_actions, time), dim=-1)) + + +class AttentionWeights(nn.Module): + """Independent attention projections belonging to one expert.""" + + 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, query_length, width = query.shape + head_dim = width // n_heads + + def split_heads(x: Tensor) -> Tensor: + return x.reshape(batch, -1, n_heads, head_dim).transpose(1, 2) + + q = split_heads(query) + k = split_heads(key) + v = split_heads(value) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim) + if allowed is not None: + # allowed has shape [B, Q, K]; True denotes a visible attention edge. + 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) + result = torch.matmul(probabilities, v) + return result.transpose(1, 2).contiguous().reshape(batch, query_length, width) + + +class DualExpertBlock(nn.Module): + """One pi0 block with separate VLM and action-expert parameters.""" + + def __init__(self, config: Pi0Config) -> None: + super().__init__() + dim = config.d_model + self.n_heads = config.n_heads + self.dropout = config.dropout + + self.vlm_norm1 = RMSNorm(dim, config.rms_eps) + self.vlm_attention = AttentionWeights(dim) + self.vlm_norm2 = RMSNorm(dim, config.rms_eps) + self.vlm_mlp = MLP(dim, config.mlp_ratio, config.dropout) + + # Unlike pi0.5, the original pi0 action expert uses ordinary norms. + # Time has already entered through ActionTimeProjection. + self.action_norm1 = RMSNorm(dim, config.rms_eps) + self.action_attention = AttentionWeights(dim) + self.action_norm2 = RMSNorm(dim, config.rms_eps) + self.action_mlp = MLP(dim, config.mlp_ratio, config.dropout) + + def forward( + self, + prefix: Tensor, + suffix: Tensor, + prefix_edges: Tensor, + suffix_edges: Tensor, + ) -> tuple[Tensor, Tensor]: + prefix_input = self.vlm_norm1(prefix) + vlm = self.vlm_attention + prefix_update = _attention( + vlm.q(prefix_input), + vlm.k(prefix_input), + vlm.v(prefix_input), + self.n_heads, + prefix_edges, + self.dropout, + self.training, + ) + prefix = prefix + vlm.out(prefix_update) + prefix = prefix + self.vlm_mlp(self.vlm_norm2(prefix)) + + suffix_input = self.action_norm1(suffix) + action = self.action_attention + # Prefix K/V are projected by the VLM expert while state/action K/V + # are projected by the action expert. Suffix queries use action-expert + # Q weights, giving one-way prefix -> suffix information flow. + keys = torch.cat((vlm.k(prefix_input), action.k(suffix_input)), dim=1) + values = torch.cat((vlm.v(prefix_input), action.v(suffix_input)), dim=1) + suffix_update = _attention( + action.q(suffix_input), + keys, + values, + self.n_heads, + suffix_edges, + self.dropout, + self.training, + ) + suffix = suffix + action.out(suffix_update) + suffix = suffix + self.action_mlp(self.action_norm2(suffix)) + return prefix, suffix + + +class Pi0Mini(nn.Module): + """Minimal original-pi0 dual-expert flow policy.""" + + def __init__(self, config: Optional[Pi0Config] = None) -> None: + super().__init__() + self.config = config or Pi0Config() + config = self.config + + self.patch_encoder = nn.Conv2d( + config.image_channels, + config.d_model, + config.patch_size, + stride=config.patch_size, + ) + self.token_embedding = nn.Embedding(config.vocab_size, config.d_model) + self.state_projection = nn.Sequential( + nn.Linear(config.state_dim, config.d_model), + nn.SiLU(), + nn.Linear(config.d_model, config.d_model), + ) + self.action_time_projection = ActionTimeProjection( + config.action_dim, config.time_feature_dim, config.d_model + ) + self.prefix_position = nn.Parameter( + torch.randn(1, config.max_prefix_tokens, config.d_model) * 0.01 + ) + self.suffix_position = nn.Parameter( + torch.randn(1, config.max_suffix_tokens, config.d_model) * 0.01 + ) + self.layers = nn.ModuleList( + DualExpertBlock(config) for _ in range(config.n_layers) + ) + self.prefix_norm = RMSNorm(config.d_model, config.rms_eps) + self.suffix_norm = RMSNorm(config.d_model, config.rms_eps) + self.vector_field_head = nn.Linear(config.d_model, config.action_dim) + + def encode_prefix(self, pixel_values: Tensor, input_ids: Tensor) -> Tensor: + if pixel_values.ndim != 4: + raise ValueError("pixel_values must have shape [batch, channels, height, width]") + if input_ids.ndim != 2: + raise ValueError("input_ids must have shape [batch, text_length]") + patches = self.patch_encoder(pixel_values).flatten(2).transpose(1, 2) + language = self.token_embedding(input_ids) + prefix = torch.cat((patches, language), 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]] + + def encode_suffix( + self, states: Tensor, noisy_actions: Tensor, timesteps: Tensor + ) -> Tensor: + if states.ndim != 2: + raise ValueError("states must have shape [batch, state_dim]") + state_token = self.state_projection(states)[:, None, :] + action_tokens = self.action_time_projection(noisy_actions, timesteps) + suffix = torch.cat((state_token, action_tokens), dim=1) + return suffix + self.suffix_position[:, : suffix.shape[1]] + + def _masks( + self, + batch: int, + prefix_length: int, + suffix_length: int, + image_tokens: int, + attention_mask: Optional[Tensor], + device: torch.device, + ) -> tuple[Tensor, Tensor]: + prefix_valid = torch.ones( + batch, prefix_length, dtype=torch.bool, device=device + ) + if attention_mask is not None: + expected = prefix_length - image_tokens + if attention_mask.shape != (batch, expected): + raise ValueError( + "attention_mask must have shape [batch, input_ids length]" + ) + prefix_valid[:, image_tokens:] = attention_mask.bool() + + prefix_edges = prefix_valid[:, None, :].expand( + batch, prefix_length, prefix_length + ).clone() + if self.config.prefix_causal: + causal = torch.ones( + prefix_length, prefix_length, dtype=torch.bool, device=device + ).tril() + prefix_edges &= causal[None, :, :] + + # Give padded queries a self edge so no softmax row is fully masked. + query_valid = prefix_valid[:, :, None] + identity = torch.eye(prefix_length, dtype=torch.bool, device=device)[None] + prefix_edges = torch.where(query_valid, prefix_edges, identity) + + visible_prefix = prefix_valid[:, None, :].expand( + batch, suffix_length, prefix_length + ) + # pi0 has three ordered attention blocks: + # [image + language] -> prefix only + # [state] -> prefix + state + # [actions] -> prefix + state + all actions + # The state and actions share action-expert weights, but they must not + # be collapsed into one visibility block: doing so leaks noisy future + # actions into the state representation. + visible_suffix = torch.ones( + batch, suffix_length, suffix_length, dtype=torch.bool, device=device + ) + if suffix_length > 1: + visible_suffix[:, 0, 1:] = False + suffix_edges = torch.cat((visible_prefix, visible_suffix), dim=-1) + return prefix_edges, suffix_edges + + def forward( + self, + pixel_values: Tensor, + input_ids: Tensor, + states: Tensor, + noisy_actions: Tensor, + timesteps: Tensor, + attention_mask: Optional[Tensor] = None, + target_velocity: Optional[Tensor] = None, + ) -> Pi0Output: + """Predict the conditional-flow vector field for an action chunk. + + Args: + pixel_values: Images with shape ``[B, C, H, W]``. + input_ids: Language prompt IDs with shape ``[B, L]``. + states: Continuous proprioception with shape ``[B, state_dim]``. + noisy_actions: Flow state with shape ``[B, horizon, action_dim]``. + timesteps: One scalar flow time per batch item. + attention_mask: Optional validity mask for language prompt tokens. + target_velocity: Optional flow-matching target, shaped like actions. + """ + if noisy_actions.ndim != 3: + raise ValueError( + "noisy_actions must have shape [batch, horizon, action_dim]" + ) + batch, horizon, action_dim = noisy_actions.shape + if action_dim != self.config.action_dim: + raise ValueError("noisy_actions has the wrong action dimension") + if horizon > self.config.max_action_tokens: + raise ValueError("action horizon exceeds configured maximum") + if states.shape != (batch, self.config.state_dim): + raise ValueError("states has the wrong batch or state dimension") + if timesteps.numel() != batch: + raise ValueError("timesteps must contain one scalar per batch item") + + prefix = self.encode_prefix(pixel_values, input_ids) + suffix = self.encode_suffix(states, noisy_actions, timesteps) + image_tokens = prefix.shape[1] - input_ids.shape[1] + prefix_edges, suffix_edges = self._masks( + batch, + prefix.shape[1], + suffix.shape[1], + image_tokens, + attention_mask, + prefix.device, + ) + for layer in self.layers: + prefix, suffix = layer(prefix, suffix, prefix_edges, suffix_edges) + + prefix_hidden = self.prefix_norm(prefix) + suffix_hidden = self.suffix_norm(suffix) + state_hidden = suffix_hidden[:, :1] + action_hidden = suffix_hidden[:, 1:] + vector_field = self.vector_field_head(action_hidden) + + flow_loss = None + if target_velocity is not None: + if target_velocity.shape != noisy_actions.shape: + raise ValueError("target_velocity must have the shape of noisy_actions") + flow_loss = F.mse_loss(vector_field, target_velocity) + + return Pi0Output( + vector_field=vector_field, + prefix_hidden=prefix_hidden, + suffix_hidden=suffix_hidden, + state_hidden=state_hidden, + action_hidden=action_hidden, + loss=flow_loss, + flow_loss=flow_loss, + ) + + +def flow_matching_sample(actions: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Sample the straight conditional-flow path used by pi-style policies. + + ``x_t = (1 - t) * noise + t * action`` and the target vector field is the + constant ``action - noise``. Sampling starts from noise at ``t=0`` and + integrates toward an action at ``t=1``. + """ + if actions.ndim < 2: + raise ValueError("actions must include batch and feature dimensions") + batch = actions.shape[0] + # In the paper's noise-time convention, t_noise is sampled as + # 0.001 + 0.999 * Beta(1.5, 1). This implementation integrates in the + # opposite, action-time direction, so tau = 1 - t_noise. + uniform = torch.rand(batch, device=actions.device, dtype=actions.dtype) + timesteps = 0.999 * (1.0 - uniform.pow(1.0 / 1.5)) + noise = torch.randn_like(actions) + broadcast_time = timesteps.reshape(batch, *([1] * (actions.ndim - 1))) + noisy_actions = (1 - broadcast_time) * noise + broadcast_time * actions + target_velocity = actions - noise + return noisy_actions, timesteps, target_velocity diff --git a/src/pi0mini/runtime.py b/src/pi0mini/runtime.py new file mode 100644 index 0000000..167255e --- /dev/null +++ b/src/pi0mini/runtime.py @@ -0,0 +1,35 @@ +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") diff --git a/src/pi0mini/train.py b/src/pi0mini/train.py new file mode 100644 index 0000000..461782e --- /dev/null +++ b/src/pi0mini/train.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import torch +from torch.utils.data import DataLoader + +from .data import Pi0OverfitDataset, materialize_subset +from .flow import FlowMatching +from .model import Pi0Config, Pi0Mini +from .runtime import choose_device, load_config, seed_everything, write_json + + +def _collate(rows: list[dict]) -> dict: + keys = ("image", "state", "actions") + return {key: torch.stack([row[key] for row in rows]) for key in keys} + + +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", max_samples=1 + ) + dataset = Pi0OverfitDataset( + manifest, image_size=cfg["data"]["image_size"], + augmentation_factor=cfg["data"]["num_variants"], + action_horizon=cfg["data"]["action_horizon"] + ) + loader = DataLoader(dataset, batch_size=len(dataset), shuffle=True, collate_fn=_collate) + mc = cfg["model"] + model = Pi0Mini(Pi0Config( + vocab_size=16, action_dim=2, state_dim=dataset.state_dim, + image_size=cfg["data"]["image_size"], patch_size=mc["patch_size"], + d_model=mc["d_model"], n_heads=mc["n_heads"], n_layers=mc["n_layers"], + mlp_ratio=mc["mlp_ratio"], max_text_tokens=8, + max_action_tokens=cfg["data"]["action_horizon"], dropout=mc["dropout"] + )).to(device) + optimizer = torch.optim.AdamW( + model.parameters(), lr=cfg["train"]["learning_rate"], + weight_decay=cfg["train"]["weight_decay"] + ) + flow = FlowMatching(cfg["inference"]["flow_steps"]) + iterator = iter(loader) + history = [] + model.train() + for step in range(1, cfg["train"]["steps"] + 1): + try: + batch = next(iterator) + except StopIteration: + iterator = iter(loader) + batch = next(iterator) + image, state, actions = ( + batch["image"].to(device), batch["state"].to(device), + batch["actions"].to(device) + ) + # A tiny fixed language instruction stands in for the PaliGemma prompt. + input_ids = torch.tensor([[1, 2]], device=device).expand(image.shape[0], -1) + noisy, timestep, target_velocity = flow.training_batch(actions) + output = model( + image, input_ids, state, noisy, timestep, + target_velocity=target_velocity + ) + 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 = {"step": step, "flow_loss": float(output.loss.detach().cpu())} + history.append(record) + print(json.dumps(record)) + output_dir = root / cfg["output_dir"] + output_dir.mkdir(parents=True, exist_ok=True) + torch.save({ + "model": model.state_dict(), "model_config": model.config.__dict__, + "config": cfg, "action_statistics": dataset.action_normalizer.as_dict() + }, 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() diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..8bd0575 --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + +from pi0mini.data import ActionNormalizer, Pi0OverfitDataset, load_manifest, materialize_subset + + +def _make_source_project(root: Path) -> Path: + processed = root / "data" / "processed" + image_dir = processed / "images" + trajectory_dir = processed / "trajectories" + image_dir.mkdir(parents=True) + trajectory_dir.mkdir(parents=True) + + pixels = np.zeros((48, 80, 3), dtype=np.uint8) + pixels[..., 0] = 64 + pixels[..., 1] = np.arange(80, dtype=np.uint8)[None, :] + Image.fromarray(pixels).save(image_dir / "sample.jpg") + + history_x = np.linspace(-7.0, 0.0, 8, dtype=np.float32) + history = np.stack((history_x, 0.02 * history_x, np.zeros(8)), axis=-1) + future_x = np.linspace(1.0, 32.0, 16, dtype=np.float32) + future = np.stack((future_x, 0.002 * future_x**2, np.zeros(16)), axis=-1) + np.savez_compressed( + trajectory_dir / "sample.npz", history_xyz=history, future_xyz=future + ) + + manifest = processed / "overfit_manifest.jsonl" + manifest.write_text( + json.dumps( + { + "clip_id": "clip-a", + "t0_us": 1234, + "image_path": "data/processed/images/sample.jpg", + "trajectory_path": "data/processed/trajectories/sample.npz", + } + ) + + "\n", + encoding="utf-8", + ) + return manifest + + +def test_action_normalizer_is_reversible() -> None: + history = torch.tensor([[-1.0, 0.0], [0.0, 0.1]]) + future = torch.tensor([[1.0, 0.2], [2.5, 0.4], [4.0, 0.7]]) + normalizer = ActionNormalizer.fit([history], [future]) + actions = normalizer.encode(history, future) + + assert actions.shape == (3, 2) + assert torch.allclose(normalizer.decode(actions, history), future, atol=1e-6) + restored = ActionNormalizer.from_dict(normalizer.as_dict()) + assert torch.allclose(restored.decode(actions, history), future, atol=1e-6) + + +def test_materialized_dataset_contract_and_determinism(tmp_path: Path) -> None: + source_manifest = _make_source_project(tmp_path / "source") + copied_manifest = materialize_subset(source_manifest, tmp_path / "target") + records = load_manifest(copied_manifest) + + assert len(records) == 1 + assert records[0].image_path.is_file() + assert records[0].trajectory_path.is_file() + + dataset = Pi0OverfitDataset(copied_manifest, seed=19) + repeated = Pi0OverfitDataset( + copied_manifest, + seed=19, + action_statistics=dataset.action_normalizer.as_dict(), + ) + assert len(dataset) == 4 + assert dataset.state_dim == 16 + + original = dataset[0] + augmented = dataset[1] + repeated_augmented = repeated[1] + assert set(original) == { + "image", + "state", + "actions", + "history_xy", + "future_xy", + "sample_id", + } + assert original["image"].shape == (3, 64, 64) + assert original["state"].shape == (16,) + assert original["actions"].shape == (16, 2) + assert original["history_xy"].shape == (8, 2) + assert original["future_xy"].shape == (16, 2) + assert original["sample_id"].endswith(":aug0") + assert augmented["sample_id"].endswith(":aug1") + assert not torch.equal(original["future_xy"], augmented["future_xy"]) + assert torch.equal(augmented["image"], repeated_augmented["image"]) + assert torch.equal(augmented["state"], repeated_augmented["state"]) + assert torch.equal(augmented["actions"], repeated_augmented["actions"]) + + decoded = dataset.action_normalizer.decode( + original["actions"], original["history_xy"] + ) + assert torch.allclose(decoded, original["future_xy"], atol=1e-5) + + +def test_statistics_can_be_saved(tmp_path: Path) -> None: + source_manifest = _make_source_project(tmp_path / "source") + manifest = materialize_subset(source_manifest, tmp_path / "target") + dataset = Pi0OverfitDataset(manifest) + path = dataset.save_action_statistics(tmp_path / "stats.json") + values = json.loads(path.read_text(encoding="utf-8")) + + assert len(values["mean"]) == 2 + assert len(values["scale"]) == 2 + assert min(values["scale"]) >= 0.05 diff --git a/tests/test_flow.py b/tests/test_flow.py new file mode 100644 index 0000000..36ebb7e --- /dev/null +++ b/tests/test_flow.py @@ -0,0 +1,65 @@ +import torch + +from pi0mini.flow import FlowMatching + + +def _time_view(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + return t.reshape(t.shape[0], *([1] * (target.ndim - 1))) + + +def test_training_batch_obeys_linear_path() -> None: + flow = FlowMatching() + action = torch.tensor([[[1.0, -2.0], [0.5, 3.0]], [[-1.0, 0.0], [2.5, -0.5]]]) + x_t, t, velocity = flow.training_batch( + action, generator=torch.Generator().manual_seed(123) + ) + + noise = action - velocity + expected = (1.0 - _time_view(t, action)) * noise + _time_view(t, action) * action + torch.testing.assert_close(x_t, expected) + assert x_t.shape == action.shape + assert torch.all((t >= 0.0) & (t <= 1.0)) + + +def test_training_batch_is_reproducible() -> None: + action = torch.randn(3, 16, 2) + flow = FlowMatching() + a = flow.training_batch(action, generator=torch.Generator().manual_seed(9)) + b = flow.training_batch(action, generator=torch.Generator().manual_seed(9)) + for left, right in zip(a, b): + torch.testing.assert_close(left, right) + + +def test_training_times_use_low_biased_shifted_beta() -> None: + flow = FlowMatching() + times = flow.sample_training_times( + 100_000, + device="cpu", + dtype=torch.float32, + generator=torch.Generator().manual_seed(17), + ) + + assert torch.all(times >= 0.0) + assert torch.all(times <= 0.999) + # E[0.999 * (1 - Beta(1.5, 1))] = 0.3996. This also guards + # against accidentally reverting to a uniform or high-time-biased draw. + assert abs(float(times.mean()) - 0.3996) < 0.005 + + +def test_loss_is_velocity_mse() -> None: + flow = FlowMatching() + target = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]]) + prediction = target + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]]) + torch.testing.assert_close(flow.loss(prediction, target), torch.tensor(1.5)) + + +def test_exact_vector_field_integrates_to_action() -> None: + flow = FlowMatching(inference_steps=10) + action = torch.tensor([[[0.25, -0.5], [1.0, 2.0]], [[-2.0, 0.75], [0.0, 0.5]]]) + + def exact_field(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + return (action.to(x_t) - x_t) / (1.0 - _time_view(t, x_t)) + + sample = flow.sample(exact_field, tuple(action.shape), action.device, + generator=torch.Generator().manual_seed(5)) + torch.testing.assert_close(sample, action, atol=1e-5, rtol=1e-5) diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..5a14493 --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,65 @@ +import torch + +from pi0mini.model import Pi0Config, Pi0Mini + + +def test_pi0_three_block_attention_mask() -> None: + config = Pi0Config( + image_size=32, + patch_size=16, + max_text_tokens=4, + max_action_tokens=3, + d_model=32, + n_heads=4, + n_layers=1, + time_feature_dim=8, + ) + model = Pi0Mini(config) + + # Two image patches per axis -> 4 image tokens, plus 2 language tokens. + prefix_edges, suffix_edges = model._masks( + batch=1, + prefix_length=6, + suffix_length=4, # one state + three action tokens + image_tokens=4, + attention_mask=torch.tensor([[1, 1]]), + device=torch.device("cpu"), + ) + + assert prefix_edges.shape == (1, 6, 6) + assert torch.all(prefix_edges) + + # Suffix key order is [six prefix keys, state, action0, action1, action2]. + assert suffix_edges.shape == (1, 4, 10) + assert torch.all(suffix_edges[0, 0, :7]) + assert not torch.any(suffix_edges[0, 0, 7:]) + assert torch.all(suffix_edges[0, 1:, :]) + + +def test_pi0_mask_respects_padded_language_keys() -> None: + model = Pi0Mini( + Pi0Config( + image_size=16, + patch_size=16, + max_text_tokens=3, + max_action_tokens=1, + d_model=16, + n_heads=4, + n_layers=1, + time_feature_dim=8, + ) + ) + prefix_edges, suffix_edges = model._masks( + batch=1, + prefix_length=4, + suffix_length=2, + image_tokens=1, + attention_mask=torch.tensor([[1, 0, 0]]), + device=torch.device("cpu"), + ) + + # Padded prefix keys are hidden from valid prefix, state, and action rows. + assert not prefix_edges[0, 0, 2] + assert not prefix_edges[0, 0, 3] + assert not suffix_edges[0, :, 2].any() + assert not suffix_edges[0, :, 3].any()