Initial pi0 minimal validation

.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+*.egg-info/
+.venv/
+data/
+outputs/**/*.pt
+.DS_Store

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 数值等价。

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

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
+  }
+]

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
+  }
+}

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
+  }
+]

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"

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"]

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}",
+        }

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

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()

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

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")

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()

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

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)

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()