Initial pi0.5 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,141 @@
+# π0.5 Minimal Validation
+
+这是一个 **π0.5 风格的最小、可运行架构验证项目**。它从相邻 Alpamayo 验证项目准备的 PAI 子集中读取 1 个驾驶样本，生成 4 个确定性视觉增强变体做过拟合，目标是检查高层语义预测、blockwise attention、action expert 与 flow matching action chunk 能否组成完整训练和推理链路。
+
+它不是 Physical Intelligence 官方代码，也不声称复现论文中的规模、数据或机器人能力。
+
+## 验证内容
+
+```text
+图像 + 本体状态
+        │
+        ▼
+小型视觉/语言主干 ──→ 高层语义 subtask
+        │                    │
+        └──── semantic prefix┘
+                     │
+ noisy action + time │
+          └──────────▼
+       action expert（双向 action attention）
+                     │
+                     ▼
+          flow velocity → action chunk
+```
+
+默认任务使用 `64×64` 前视相机图像、离散化历史 XY 状态和一句 subtask，预测未来 `16×2` 的二维轨迹。样本少且重复训练，因此成功标准是“能够记忆训练集”，而不是泛化。
+
+## 与 π0.5 的对应关系
+
+保留的关键思想：
+
+- 视觉/语言语义主干与独立 action expert；
+- observation/text 与 action token 采用 blockwise attention 信息流；
+- action token 内部采用双向 attention，一次预测整个 action chunk；
+- flow matching 训练连续动作，推理时从高斯噪声做 Euler 积分；
+- timestep 通过 adaptive RMSNorm 调制 action expert；
+- 同时优化高层语义目标和低层连续动作目标。
+- VLM 自回归输出高层 subtask 以及 64-bin 离散动作 token；
+- 推理时分别还原 VLM token 轨迹和 flow/diffusion 轨迹，与真值三方比较。
+
+为了单机快速验证而简化的部分：
+
+- 小型自定义视觉编码器和 Transformer 取代 PaliGemma 2B；
+- 隐藏宽度、层数、图像尺寸和 action horizon 大幅缩小；
+- 使用固定小词表的 subtask 分类/生成，而非开放式语言理解；
+- 使用二维驾驶轨迹，取代多机器人、多相机和高维关节控制；
+- 默认使用 1 个 PAI 驾驶样本的 4 个视觉增强版本，没有机器人或网页联合训练数据；
+- 使用正交 DCT-II、低频优先排列和 64-bin 量化实现 FAST-lite；保留全部频率系数，因此只验证表示与训练接口，不复现官方压缩率；
+- Euler sampler 用固定步长，不包含面向部署的实时控制与安全系统。
+
+因此，本项目适合验证架构接口和数学链路，不适合比较论文指标。
+
+## 安装
+
+建议使用 Python 3.10 或更新版本：
+
+```bash
+cd pi05-minimal-validation
+python -m venv .venv
+source .venv/bin/activate
+python -m pip install -U pip
+python -m pip install -e '.[dev]'
+```
+
+默认实现只依赖 PyTorch、NumPy、Pillow、Matplotlib 和 PyYAML，不需要下载预训练权重。默认配置会读取：
+
+```text
+../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl
+```
+
+如果该路径不同，请先修改 `configs/overfit.yaml` 的 `data.source_manifest`。
+
+## 运行测试
+
+```bash
+cd pi05-minimal-validation
+pytest -q
+```
+
+`tests/test_flow.py` 单独验证：
+
+- 直线路径 `x_t=(1-t)·noise+t·action`；
+- 目标速度 `v=action-noise`；
+- velocity MSE；
+- 固定目标的精确 vector field 经 Euler 积分后能到达目标；
+- 使用相同 `torch.Generator` 时采样可复现。
+
+## 过拟合训练
+
+```bash
+cd pi05-minimal-validation
+pi05mini-train --config configs/overfit.yaml
+```
+
+也可以直接使用模块入口：
+
+```bash
+python -m pi05mini.train --config configs/overfit.yaml
+```
+
+配置默认先进行 300 步离散 VLM/FAST-lite 预训练，再进行 500 步联合 post-training。Stage 1 不更新 flow action expert；Stage 2 联合优化 token CE 与 flow loss。CPU 可以运行；Apple Silicon 可自动选择 MPS，CUDA 可自动选择 GPU。
+
+## 推理与效果图
+
+```bash
+pi05mini-infer --checkpoint outputs/overfit/checkpoint.pt
+```
+
+输出图 `outputs/overfit/trajectory_comparison.png` 会叠加 ground truth、VLM 离散动作 token 解码轨迹和 diffusion/flow 轨迹；逐样本指标写入 `evaluation.json`。训练过程另存为 `training_history.json`，训练摘要为 `summary.json`。
+
+## 配置重点
+
+`configs/overfit.yaml` 中最常调整的参数：
+
+- `data.num_variants`：默认 4，刻意保持很小；
+- `model.d_model`、`n_layers`：控制最小模型容量；
+- `train.steps`：过拟合不充分时增加；
+- `train.semantic_loss_weight` 与 `flow_loss_weight`：两个训练目标的权重；
+- `inference.flow_steps`：Euler 积分步数，越大通常越慢且更平滑。
+
+## 安全限制
+
+- 该模型只在合成样本上验证记忆能力，输出不具备现实世界有效性。
+- 不得用于车辆、机器人、机械臂或任何会造成物理影响的设备控制。
+- 不得把 loss 下降或单样本轨迹重合解释为安全性、鲁棒性或泛化能力。
+- 真实部署至少还需要独立验证集、分布外和扰动测试、控制约束、碰撞检查、故障降级、急停以及人工监督。
+- checkpoint 和配置可能加载本地路径；只运行可信来源的文件，不加载未知 PyTorch pickle checkpoint。
+
+## 项目结构
+
+```text
+pi05-minimal-validation/
+├── configs/overfit.yaml      # 4 样本过拟合配置
+├── src/pi05mini/             # 模型、flow、数据、训练与推理
+├── tests/test_flow.py        # flow matching 数学测试
+├── pyproject.toml
+└── README.md
+```
+
+## 预期结论
+
+若最小实验成功，它只能说明：π0.5 风格的语义主干、blockwise 双专家交互和 flow action head 在此实现中能够端到端反向传播并记忆几个样本。它不能证明实现与官方模型数值等价，也不能证明在真实驾驶或机器人任务中有效。

configs/overfit.yaml

@@ -0,0 +1,32 @@
+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:
+  pretrain_steps: 300
+  posttrain_steps: 500
+  batch_size: 4
+  learning_rate: 0.001
+  weight_decay: 0.0
+  grad_clip_norm: 1.0
+  semantic_loss_weight: 1.0
+  flow_loss_weight: 2.0
+  log_every: 20
+
+inference:
+  flow_steps: 10

outputs/overfit/evaluation.json

@@ -0,0 +1,34 @@
+[
+  {
+    "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug0",
+    "target_subtask": "steer_left",
+    "predicted_subtask": "steer_left",
+    "vlm_action_tokens_valid": true,
+    "vlm_trajectory_mse": 0.07854792475700378,
+    "diffusion_trajectory_mse": 0.27604275941848755
+  },
+  {
+    "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug1",
+    "target_subtask": "steer_left",
+    "predicted_subtask": "steer_left",
+    "vlm_action_tokens_valid": true,
+    "vlm_trajectory_mse": 0.025470610707998276,
+    "diffusion_trajectory_mse": 0.20049458742141724
+  },
+  {
+    "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug2",
+    "target_subtask": "steer_left",
+    "predicted_subtask": "steer_left",
+    "vlm_action_tokens_valid": true,
+    "vlm_trajectory_mse": 0.013774838298559189,
+    "diffusion_trajectory_mse": 0.03129641339182854
+  },
+  {
+    "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug3",
+    "target_subtask": "steer_left",
+    "predicted_subtask": "steer_left",
+    "vlm_action_tokens_valid": true,
+    "vlm_trajectory_mse": 0.059195056557655334,
+    "diffusion_trajectory_mse": 0.11321215331554413
+  }
+]

outputs/overfit/summary.json

@@ -0,0 +1,19 @@
+{
+  "device": "mps",
+  "samples": 4,
+  "parameters": 816482,
+  "initial": {
+    "stage": "pretrain",
+    "step": 1,
+    "loss": 78.87449645996094,
+    "text_loss": 78.87449645996094,
+    "flow_loss": null
+  },
+  "final": {
+    "stage": "posttrain",
+    "step": 500,
+    "loss": 0.12353827804327011,
+    "text_loss": 3.1691649837739533e-06,
+    "flow_loss": 0.06176755577325821
+  }
+}

outputs/overfit/training_history.json

@@ -0,0 +1,296 @@
+[
+  {
+    "stage": "pretrain",
+    "step": 1,
+    "loss": 78.87449645996094,
+    "text_loss": 78.87449645996094,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 20,
+    "loss": 2.475086212158203,
+    "text_loss": 2.475086212158203,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 40,
+    "loss": 1.1025534868240356,
+    "text_loss": 1.1025534868240356,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 60,
+    "loss": 0.601564884185791,
+    "text_loss": 0.601564884185791,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 80,
+    "loss": 0.25571808218955994,
+    "text_loss": 0.25571808218955994,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 100,
+    "loss": 0.03314117714762688,
+    "text_loss": 0.03314117714762688,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 120,
+    "loss": 0.005310459993779659,
+    "text_loss": 0.005310459993779659,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 140,
+    "loss": 3.0251239877543412e-05,
+    "text_loss": 3.0251239877543412e-05,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 160,
+    "loss": 1.1354763955750968e-05,
+    "text_loss": 1.1354763955750968e-05,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 180,
+    "loss": 6.638075319642667e-06,
+    "text_loss": 6.638075319642667e-06,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 200,
+    "loss": 5.348983449948719e-06,
+    "text_loss": 5.348983449948719e-06,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 220,
+    "loss": 4.736786195280729e-06,
+    "text_loss": 4.736786195280729e-06,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 240,
+    "loss": 4.316162630857434e-06,
+    "text_loss": 4.316162630857434e-06,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 260,
+    "loss": 3.975574600190157e-06,
+    "text_loss": 3.975574600190157e-06,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 280,
+    "loss": 3.6877775073662633e-06,
+    "text_loss": 3.6877775073662633e-06,
+    "flow_loss": null
+  },
+  {
+    "stage": "pretrain",
+    "step": 300,
+    "loss": 3.4408508327032905e-06,
+    "text_loss": 3.4408508327032905e-06,
+    "flow_loss": null
+  },
+  {
+    "stage": "posttrain",
+    "step": 1,
+    "loss": 8.707193374633789,
+    "text_loss": 3.428078571232618e-06,
+    "flow_loss": 4.353594779968262
+  },
+  {
+    "stage": "posttrain",
+    "step": 20,
+    "loss": 3.533726453781128,
+    "text_loss": 0.14151602983474731,
+    "flow_loss": 1.6961052417755127
+  },
+  {
+    "stage": "posttrain",
+    "step": 40,
+    "loss": 3.0850143432617188,
+    "text_loss": 0.00014649856893811375,
+    "flow_loss": 1.5424339771270752
+  },
+  {
+    "stage": "posttrain",
+    "step": 60,
+    "loss": 1.3556954860687256,
+    "text_loss": 0.00011281041224719957,
+    "flow_loss": 0.6777913570404053
+  },
+  {
+    "stage": "posttrain",
+    "step": 80,
+    "loss": 1.4165724515914917,
+    "text_loss": 6.599419066333212e-06,
+    "flow_loss": 0.7082829475402832
+  },
+  {
+    "stage": "posttrain",
+    "step": 100,
+    "loss": 0.8609042167663574,
+    "text_loss": 8.11900736152893e-06,
+    "flow_loss": 0.430448055267334
+  },
+  {
+    "stage": "posttrain",
+    "step": 120,
+    "loss": 0.5642163753509521,
+    "text_loss": 7.576879397674929e-06,
+    "flow_loss": 0.28210440278053284
+  },
+  {
+    "stage": "posttrain",
+    "step": 140,
+    "loss": 0.7110214233398438,
+    "text_loss": 6.249698344618082e-06,
+    "flow_loss": 0.35550758242607117
+  },
+  {
+    "stage": "posttrain",
+    "step": 160,
+    "loss": 0.18725578486919403,
+    "text_loss": 2.876289045161684e-06,
+    "flow_loss": 0.09362645447254181
+  },
+  {
+    "stage": "posttrain",
+    "step": 180,
+    "loss": 0.33923929929733276,
+    "text_loss": 6.748496616637567e-06,
+    "flow_loss": 0.16961628198623657
+  },
+  {
+    "stage": "posttrain",
+    "step": 200,
+    "loss": 0.3003162443637848,
+    "text_loss": 9.495882295595948e-06,
+    "flow_loss": 0.15015336871147156
+  },
+  {
+    "stage": "posttrain",
+    "step": 220,
+    "loss": 0.23716998100280762,
+    "text_loss": 5.660519946104614e-06,
+    "flow_loss": 0.11858215928077698
+  },
+  {
+    "stage": "posttrain",
+    "step": 240,
+    "loss": 0.42717865109443665,
+    "text_loss": 8.400202204938978e-06,
+    "flow_loss": 0.21358512341976166
+  },
+  {
+    "stage": "posttrain",
+    "step": 260,
+    "loss": 0.11702287942171097,
+    "text_loss": 4.122769951209193e-06,
+    "flow_loss": 0.058509379625320435
+  },
+  {
+    "stage": "posttrain",
+    "step": 280,
+    "loss": 0.19245445728302002,
+    "text_loss": 5.941245035501197e-06,
+    "flow_loss": 0.09622425585985184
+  },
+  {
+    "stage": "posttrain",
+    "step": 300,
+    "loss": 0.18475297093391418,
+    "text_loss": 4.744139005197212e-06,
+    "flow_loss": 0.09237411618232727
+  },
+  {
+    "stage": "posttrain",
+    "step": 320,
+    "loss": 0.38481318950653076,
+    "text_loss": 3.1461640901397914e-06,
+    "flow_loss": 0.19240501523017883
+  },
+  {
+    "stage": "posttrain",
+    "step": 340,
+    "loss": 0.2163764089345932,
+    "text_loss": 4.1934181354008615e-06,
+    "flow_loss": 0.10818611085414886
+  },
+  {
+    "stage": "posttrain",
+    "step": 360,
+    "loss": 0.23533868789672852,
+    "text_loss": 3.3019753118423978e-06,
+    "flow_loss": 0.11766768991947174
+  },
+  {
+    "stage": "posttrain",
+    "step": 380,
+    "loss": 0.08397223055362701,
+    "text_loss": 8.292082384286914e-06,
+    "flow_loss": 0.04198196902871132
+  },
+  {
+    "stage": "posttrain",
+    "step": 400,
+    "loss": 0.13249839842319489,
+    "text_loss": 1.0871653103095014e-05,
+    "flow_loss": 0.06624376028776169
+  },
+  {
+    "stage": "posttrain",
+    "step": 420,
+    "loss": 0.4479130506515503,
+    "text_loss": 1.0128642315976322e-05,
+    "flow_loss": 0.22395145893096924
+  },
+  {
+    "stage": "posttrain",
+    "step": 440,
+    "loss": 0.22416232526302338,
+    "text_loss": 2.6114355478057405e-06,
+    "flow_loss": 0.11207985877990723
+  },
+  {
+    "stage": "posttrain",
+    "step": 460,
+    "loss": 0.06346385926008224,
+    "text_loss": 3.1035808660817565e-06,
+    "flow_loss": 0.031730376183986664
+  },
+  {
+    "stage": "posttrain",
+    "step": 480,
+    "loss": 0.385486900806427,
+    "text_loss": 2.184032382501755e-06,
+    "flow_loss": 0.19274236261844635
+  },
+  {
+    "stage": "posttrain",
+    "step": 500,
+    "loss": 0.12353827804327011,
+    "text_loss": 3.1691649837739533e-06,
+    "flow_loss": 0.06176755577325821
+  }
+]

pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools>=69", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "pi05-minimal-validation"
+version = "0.1.0"
+description = "A tiny, educational pi0.5-style hierarchical VLA overfitting experiment"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [{ name = "pi05 minimal validation contributors" }]
+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]
+pi05mini-train = "pi05mini.train:main"
+pi05mini-infer = "pi05mini.infer:main"
+
+[tool.setuptools]
+package-dir = { "" = "src" }
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-ra"

src/pi05mini/__init__.py

@@ -0,0 +1,5 @@
+"""Minimal pi0.5 architecture validation package."""
+
+from .model import Pi05Config, Pi05Mini, Pi05Output, flow_matching_sample
+
+__all__ = ["Pi05Config", "Pi05Mini", "Pi05Output", "flow_matching_sample"]

src/pi05mini/data.py

@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+import json
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image, ImageEnhance
+from torch.utils.data import Dataset
+
+from .tokenizer import MiniTokenizer
+
+
+@dataclass(frozen=True)
+class SampleRecord:
+    clip_id: str
+    t0_us: int
+    reasoning: str
+    image_path: Path
+    trajectory_path: Path
+
+    @property
+    def sample_id(self) -> str:
+        return f"{self.clip_id}_{self.t0_us}"
+
+
+@dataclass(frozen=True)
+class ActionStatistics:
+    mean: torch.Tensor
+    scale: torch.Tensor
+
+    def as_dict(self) -> dict[str, list[float]]:
+        return {
+            "mean": [float(value) for value in self.mean],
+            "scale": [float(value) for value in self.scale],
+        }
+
+
+class ActionNormalizer:
+    """Normalize 2-D waypoint deltas and retain invertible dataset statistics."""
+
+    def __init__(self, mean: torch.Tensor, scale: torch.Tensor) -> None:
+        self.mean = torch.as_tensor(mean, dtype=torch.float32).reshape(2)
+        self.scale = torch.as_tensor(scale, dtype=torch.float32).reshape(2).clamp_min(0.05)
+
+    @classmethod
+    def fit(cls, histories: Sequence[torch.Tensor], futures: Sequence[torch.Tensor]) -> ActionNormalizer:
+        if not histories or len(histories) != len(futures):
+            raise ValueError("histories and futures must be non-empty sequences of equal length")
+        chunks = []
+        for history, future in zip(histories, futures, strict=True):
+            chunks.append(cls.to_deltas(history, future).reshape(-1, 2))
+        flat = torch.cat(chunks, dim=0).float()
+        return cls(flat.mean(dim=0), flat.std(dim=0, unbiased=False))
+
+    @classmethod
+    def from_dict(cls, values: dict[str, Sequence[float]]) -> ActionNormalizer:
+        return cls(torch.tensor(values["mean"]), torch.tensor(values["scale"]))
+
+    @property
+    def statistics(self) -> ActionStatistics:
+        return ActionStatistics(self.mean.clone(), self.scale.clone())
+
+    @staticmethod
+    def to_deltas(history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor:
+        previous = torch.cat((history_xy[-1:].float(), future_xy[:-1].float()), dim=0)
+        return future_xy.float() - previous
+
+    def encode(self, history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor:
+        return (self.to_deltas(history_xy, future_xy) - self.mean) / self.scale
+
+    def decode(self, actions: torch.Tensor, history_xy: torch.Tensor) -> torch.Tensor:
+        deltas = actions.float() * self.scale + self.mean
+        return history_xy[-1:].float() + torch.cumsum(deltas, dim=0)
+
+    def as_dict(self) -> dict[str, list[float]]:
+        return self.statistics.as_dict()
+
+    def save(self, path: str | Path) -> Path:
+        output = Path(path).expanduser().resolve()
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(json.dumps(self.as_dict(), indent=2) + "\n", encoding="utf-8")
+        return output
+
+
+def _resolve_data_path(manifest_path: Path, raw_path: str | Path) -> Path:
+    path = Path(raw_path).expanduser()
+    if path.is_absolute():
+        return path
+    candidates = [manifest_path.parent / path]
+    if len(manifest_path.parents) >= 3:
+        candidates.insert(0, manifest_path.parents[2] / path)
+    for candidate in candidates:
+        if candidate.exists():
+            return candidate.resolve()
+    return candidates[0].resolve()
+
+
+def load_manifest(path: str | Path) -> list[SampleRecord]:
+    manifest_path = Path(path).expanduser().resolve()
+    records: list[SampleRecord] = []
+    for line in manifest_path.read_text(encoding="utf-8").splitlines():
+        if not line.strip():
+            continue
+        row = json.loads(line)
+        records.append(
+            SampleRecord(
+                clip_id=str(row["clip_id"]),
+                t0_us=int(row["t0_us"]),
+                reasoning=str(row.get("reasoning", "")),
+                image_path=_resolve_data_path(manifest_path, row["image_path"]),
+                trajectory_path=_resolve_data_path(manifest_path, row["trajectory_path"]),
+            )
+        )
+    if not records:
+        raise ValueError(f"no samples found in {manifest_path}")
+    return records
+
+
+def materialize_subset(
+    source_manifest: str | Path,
+    output_dir: str | Path,
+    max_samples: int = 3,
+    num_variants: int | None = None,
+) -> Path:
+    """Copy a self-contained PAI overfit subset out of the Alpamayo project."""
+
+    # ``num_variants`` is accepted as a compatibility alias used by the
+    # training entry point. Actual variants stay virtual and deterministic in
+    # Pi05OverfitDataset, avoiding redundant copies of the large JPEG.
+    if num_variants is not None:
+        if num_variants < 1:
+            raise ValueError("num_variants must be at least one")
+        max_samples = min(max_samples, num_variants)
+    if max_samples < 1:
+        raise ValueError("max_samples must be at least one")
+    records = load_manifest(source_manifest)[:max_samples]
+    destination = Path(output_dir).expanduser().resolve()
+    image_dir = destination / "images"
+    trajectory_dir = destination / "trajectories"
+    image_dir.mkdir(parents=True, exist_ok=True)
+    trajectory_dir.mkdir(parents=True, exist_ok=True)
+
+    rows: list[dict[str, Any]] = []
+    for index, record in enumerate(records):
+        stem = f"sample_{index:02d}_{record.sample_id}"
+        image_path = image_dir / f"{stem}{record.image_path.suffix.lower()}"
+        trajectory_path = trajectory_dir / f"{stem}.npz"
+        shutil.copy2(record.image_path, image_path)
+        shutil.copy2(record.trajectory_path, trajectory_path)
+        rows.append(
+            {
+                "clip_id": record.clip_id,
+                "t0_us": record.t0_us,
+                "reasoning": record.reasoning,
+                "image_path": str(image_path.relative_to(destination)),
+                "trajectory_path": str(trajectory_path.relative_to(destination)),
+                "history_steps": 8,
+                "expert_waypoints": 16,
+            }
+        )
+
+    manifest_path = destination / "overfit_manifest.jsonl"
+    manifest_path.write_text(
+        "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
+        encoding="utf-8",
+    )
+    return manifest_path
+
+
+def _resample_waypoints(xy: torch.Tensor, steps: int = 16) -> torch.Tensor:
+    if xy.shape[0] == steps:
+        return xy.float()
+    if xy.shape[0] < 2:
+        raise ValueError("future trajectory must contain at least two waypoints")
+    signal = xy.float().T.unsqueeze(0)
+    return F.interpolate(signal, size=steps, mode="linear", align_corners=True)[0].T
+
+
+class Pi05OverfitDataset(Dataset[dict[str, Any]]):
+    """Tiny deterministic PAI dataset for validating a pi0.5-style training loop."""
+
+    def __init__(
+        self,
+        manifest_path: str | Path,
+        *,
+        image_size: int | tuple[int, int] = (128, 224),
+        augmentation_factor: int = 4,
+        action_horizon: int = 16,
+        seed: int = 7,
+        tokenizer: MiniTokenizer | None = None,
+        action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None,
+    ) -> None:
+        if augmentation_factor < 1:
+            raise ValueError("augmentation_factor must be at least one")
+        if isinstance(image_size, int):
+            image_size = (image_size, image_size)
+        if len(image_size) != 2 or min(image_size) < 1:
+            raise ValueError("image_size must be a positive (height, width) pair")
+        if action_horizon < 2:
+            raise ValueError("action_horizon must be at least two")
+        self.records = load_manifest(manifest_path)
+        self.image_size = (int(image_size[0]), int(image_size[1]))
+        self.augmentation_factor = int(augmentation_factor)
+        self.action_horizon = int(action_horizon)
+        self.seed = int(seed)
+        self.tokenizer = tokenizer or MiniTokenizer()
+        self.state_dim = 16
+
+        if action_statistics is None:
+            geometries = [
+                self._load_augmented_geometry(base_index, variant)
+                for base_index in range(len(self.records))
+                for variant in range(self.augmentation_factor)
+            ]
+            histories = [geometry[0] for geometry in geometries]
+            futures = [geometry[1] for geometry in geometries]
+            self.action_normalizer = ActionNormalizer.fit(histories, futures)
+        elif isinstance(action_statistics, ActionStatistics):
+            self.action_normalizer = ActionNormalizer(
+                action_statistics.mean, action_statistics.scale
+            )
+        else:
+            self.action_normalizer = ActionNormalizer.from_dict(action_statistics)
+
+    @property
+    def action_statistics(self) -> ActionStatistics:
+        return self.action_normalizer.statistics
+
+    @property
+    def action_mean(self) -> torch.Tensor:
+        return self.action_normalizer.mean.clone()
+
+    @property
+    def action_std(self) -> torch.Tensor:
+        return self.action_normalizer.scale.clone()
+
+    def save_action_statistics(self, path: str | Path) -> Path:
+        return self.action_normalizer.save(path)
+
+    def __len__(self) -> int:
+        return len(self.records) * self.augmentation_factor
+
+    def _indices(self, index: int) -> tuple[int, int]:
+        if index < 0:
+            index += len(self)
+        if index < 0 or index >= len(self):
+            raise IndexError(index)
+        return divmod(index, self.augmentation_factor)
+
+    def _rng(self, base_index: int, variant: int) -> np.random.Generator:
+        return np.random.default_rng(
+            np.random.SeedSequence([self.seed, base_index, variant])
+        )
+
+    def _load_augmented_geometry(
+        self, base_index: int, variant: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        record = self.records[base_index]
+        with np.load(record.trajectory_path) as arrays:
+            history = torch.from_numpy(np.asarray(arrays["history_xyz"], dtype=np.float32))[:, :2]
+            future = torch.from_numpy(np.asarray(arrays["future_xyz"], dtype=np.float32))[:, :2]
+        future = _resample_waypoints(future, steps=self.action_horizon)
+        if variant == 0:
+            return history, future
+
+        rng = self._rng(base_index, variant)
+        angle = float(rng.normal(0.0, 0.008))
+        scale = float(rng.normal(1.0, 0.006))
+        rotation = history.new_tensor(
+            [[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]
+        )
+        history = scale * (history @ rotation.T)
+        future = scale * (future @ rotation.T)
+        # A tiny smooth lateral perturbation adds distinct trajectories while retaining the label.
+        curvature = float(rng.normal(0.0, 0.012))
+        progress = torch.linspace(0.0, 1.0, future.shape[0], dtype=future.dtype)
+        future[:, 1] += curvature * progress.square()
+        return history, future
+
+    def _load_image(self, base_index: int, variant: int) -> torch.Tensor:
+        with Image.open(self.records[base_index].image_path) as source:
+            image = source.convert("RGB")
+        if variant:
+            rng = self._rng(base_index, variant)
+            brightness = float(rng.uniform(0.97, 1.03))
+            image = ImageEnhance.Brightness(image).enhance(brightness)
+        height, width = self.image_size
+        image = image.resize((width, height), resample=Image.Resampling.BILINEAR)
+        array = np.asarray(image, dtype=np.float32).copy() / 255.0
+        return torch.from_numpy(array).permute(2, 0, 1).contiguous()
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        base_index, variant = self._indices(index)
+        record = self.records[base_index]
+        history_xy, future_xy = self._load_augmented_geometry(base_index, variant)
+        reasoning = record.reasoning
+        target = self.tokenizer.encode_subtask(reasoning)
+        actions = self.action_normalizer.encode(history_xy, future_xy)
+        return {
+            "image": self._load_image(base_index, variant),
+            "prefix_tokens": self.tokenizer.build_prefix(history_xy),
+            "subtask_target": target,
+            "subtask_label": self.tokenizer.subtasks.decode(target),
+            "actions": actions,
+            "action_tokens": self.tokenizer.encode_actions(actions),
+            "state": history_xy.reshape(-1),
+            "history_xy": history_xy,
+            "future_xy": future_xy,
+            "reasoning": reasoning,
+            "sample_id": f"{record.sample_id}:aug{variant}",
+        }

src/pi05mini/flow.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+from torch import Tensor
+
+
+class FlowMatching:
+    """Rectified-flow objective used by the minimal action expert."""
+
+    def __init__(self, inference_steps: int = 10, cutoff: float = 0.999) -> None:
+        self.inference_steps = inference_steps
+        self.cutoff = cutoff
+
+    def training_batch(
+        self, target: Tensor, generator: torch.Generator | None = None
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        batch = target.shape[0]
+        uniform = torch.rand(
+            batch, device=target.device, dtype=target.dtype, generator=generator
+        )
+        # Paper sampling: z~Beta(1.5,1), tau=s*(1-z), emphasizing noisy/low tau.
+        timestep = self.cutoff * (1.0 - uniform.pow(1.0 / 1.5))
+        noise = torch.randn(
+            target.shape,
+            device=target.device,
+            dtype=target.dtype,
+            generator=generator,
+        )
+        t = timestep[:, None, None]
+        noisy = t * target + (1.0 - t) * noise
+        return noisy, timestep, target - noise
+
+    @staticmethod
+    def loss(prediction: Tensor, target_velocity: Tensor) -> Tensor:
+        return torch.nn.functional.mse_loss(prediction.float(), target_velocity.float())
+
+    @torch.no_grad()
+    def sample(
+        self,
+        model_fn: Callable[[Tensor, Tensor], Tensor],
+        shape: tuple[int, ...],
+        device: torch.device | str,
+        steps: int | None = None,
+        generator: torch.Generator | None = None,
+    ) -> Tensor:
+        count = steps or self.inference_steps
+        x = torch.randn(shape, device=device, generator=generator)
+        grid = torch.linspace(0.0, 1.0, count + 1, device=device)
+        for index in range(count):
+            timestep = grid[index].expand(shape[0])
+            x = x + (grid[index + 1] - grid[index]) * model_fn(x, timestep)
+        return x

src/pi05mini/infer.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from .data import Pi05OverfitDataset
+from .flow import FlowMatching
+from .model import Pi05Config, Pi05Mini
+from .runtime import choose_device, write_json
+
+
+@torch.no_grad()
+def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path:
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    cfg = checkpoint["config"]
+    root = Path(__file__).resolve().parents[2]
+    device = choose_device(cfg.get("device", "auto"))
+    data_manifest = Path(manifest) if manifest else root / "data/overfit/overfit_manifest.jsonl"
+    dataset = Pi05OverfitDataset(
+        data_manifest, image_size=cfg["data"]["image_size"],
+        action_horizon=cfg["data"]["action_horizon"],
+        action_statistics=checkpoint["action_statistics"]
+    )
+    model = Pi05Mini(Pi05Config(**checkpoint["model_config"]))
+    model.load_state_dict(checkpoint["model"])
+    model.to(device).eval()
+    flow = FlowMatching(cfg["inference"]["flow_steps"])
+    rows = []
+    figure, axes = plt.subplots(1, len(dataset), figsize=(4 * len(dataset), 4), squeeze=False)
+    for index, sample in enumerate(dataset):
+        image = sample["image"][None].to(device)
+        tokens = sample["prefix_tokens"][None].to(device)
+        # Autoregressively predict semantic subtask followed by the complete
+        # discretized action stream (ACTION_START, 32 scalar tokens, ACTION_END).
+        zero_action = torch.zeros(1, dataset.action_horizon, 2, device=device)
+        generated: list[int] = []
+        expected_tokens = 1 + 2 + dataset.action_horizon * 2
+        current_tokens = tokens
+        for _ in range(expected_tokens):
+            semantic = model(
+                image, current_tokens, zero_action,
+                torch.zeros(1, device=device),
+                generated_token_start=tokens.shape[1]
+            )
+            next_token = int(semantic.text_logits[0, -1].argmax().cpu())
+            generated.append(next_token)
+            current_tokens = torch.cat(
+                (current_tokens, torch.tensor([[next_token]], device=device)), dim=1
+            )
+        predicted_token = generated[0]
+        vlm_action_tokens = torch.tensor(generated[1:], dtype=torch.long)
+        try:
+            vlm_normalized = dataset.tokenizer.decode_actions(vlm_action_tokens)
+            vlm_xy = dataset.action_normalizer.decode(
+                vlm_normalized, sample["history_xy"]
+            ).numpy()
+            vlm_valid = True
+        except ValueError:
+            vlm_xy = np.full_like(sample["future_xy"].numpy(), np.nan)
+            vlm_valid = False
+        conditioned_tokens = torch.cat(
+            (tokens, torch.tensor([[predicted_token]], device=device)), dim=1
+        )
+        def velocity(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+            return model(
+                image, conditioned_tokens, x, t,
+                generated_token_start=tokens.shape[1]
+            ).vector_field
+        normalized = flow.sample(
+            velocity, (1, dataset.action_horizon, 2), device,
+            steps=cfg["inference"]["flow_steps"],
+            generator=None,
+        )[0].cpu()
+        predicted_xy = dataset.action_normalizer.decode(
+            normalized, sample["history_xy"]
+        ).numpy()
+        target_xy = sample["future_xy"].numpy()
+        diffusion_mse = float(np.mean((predicted_xy - target_xy) ** 2))
+        vlm_mse = float(np.mean((vlm_xy - target_xy) ** 2)) if vlm_valid else None
+        target_label = dataset.tokenizer.subtasks.decode(sample["subtask_target"])
+        try:
+            predicted_label = dataset.tokenizer.subtasks.decode(predicted_token)
+        except ValueError:
+            predicted_label = f"token_{predicted_token}"
+        rows.append({
+            "sample_id": sample["sample_id"], "target_subtask": target_label,
+            "predicted_subtask": predicted_label,
+            "vlm_action_tokens_valid": vlm_valid,
+            "vlm_trajectory_mse": vlm_mse,
+            "diffusion_trajectory_mse": diffusion_mse,
+        })
+        ax = axes[0, index]
+        ax.plot(target_xy[:, 0], target_xy[:, 1], "o-", label="target")
+        ax.plot(vlm_xy[:, 0], vlm_xy[:, 1], "s--", label="VLM tokens")
+        ax.plot(predicted_xy[:, 0], predicted_xy[:, 1], "x--", label="diffusion")
+        ax.scatter([0], [0], c="black", marker="s", label="ego")
+        vlm_text = "invalid" if vlm_mse is None else f"{vlm_mse:.3f}"
+        ax.set_title(
+            f"{predicted_label}\nVLM={vlm_text}, diffusion={diffusion_mse:.3f}"
+        )
+        ax.set_xlabel("forward x (m)")
+        ax.set_ylabel("left y (m)")
+        ax.axis("equal")
+        ax.grid(True, alpha=.3)
+        ax.legend(fontsize=8)
+    figure.tight_layout()
+    output_dir = Path(checkpoint_path).parent
+    plot_path = output_dir / "trajectory_comparison.png"
+    figure.savefig(plot_path, dpi=160)
+    plt.close(figure)
+    write_json(output_dir / "evaluation.json", rows)
+    print(rows)
+    return plot_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", default="outputs/overfit/checkpoint.pt")
+    parser.add_argument("--manifest")
+    args = parser.parse_args()
+    evaluate(args.checkpoint, args.manifest)
+
+
+if __name__ == "__main__":
+    main()

src/pi05mini/model.py

@@ -0,0 +1,396 @@
+"""A compact, pure-PyTorch approximation of the pi0.5 dual-expert model.
+
+The important architectural detail is that prefix and action tokens have
+different parameters at every Transformer layer.  Prefix queries only attend
+to prefix keys, while action queries attend to the complete prefix and to all
+action tokens.  This is deliberately implemented explicitly (instead of
+concatenating the streams into an ordinary Transformer) so the information
+flow is easy to inspect in a minimal reproduction.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+@dataclass
+class Pi05Config:
+    vocab_size: int = 512
+    action_dim: int = 2
+    image_channels: int = 3
+    image_size: int = 64
+    patch_size: int = 16
+    d_model: int = 128
+    n_heads: int = 4
+    n_layers: int = 4
+    mlp_ratio: int = 4
+    max_text_tokens: int = 64
+    max_action_tokens: int = 32
+    dropout: float = 0.0
+    rms_eps: float = 1e-6
+
+    def __post_init__(self) -> None:
+        if self.d_model % self.n_heads:
+            raise ValueError("d_model must be divisible by n_heads")
+        if self.image_size % self.patch_size:
+            raise ValueError("image_size must be divisible by patch_size")
+
+    @property
+    def max_image_tokens(self) -> int:
+        return (self.image_size // self.patch_size) ** 2
+
+    @property
+    def max_prefix_tokens(self) -> int:
+        return self.max_image_tokens + self.max_text_tokens
+
+
+@dataclass
+class Pi05Output:
+    text_logits: Tensor
+    vector_field: Tensor
+    prefix_hidden: Tensor
+    action_hidden: Tensor
+    text_token_offset: int
+    loss: Optional[Tensor] = None
+    text_loss: Optional[Tensor] = None
+    flow_loss: Optional[Tensor] = None
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        # Accumulating the variance in fp32 avoids overflows under autocast.
+        rms = x.float().pow(2).mean(dim=-1, keepdim=True)
+        normalized = x * torch.rsqrt(rms.to(x.dtype) + self.eps)
+        return normalized * self.weight
+
+
+class AdaRMSNorm(nn.Module):
+    """RMSNorm modulated by a flow-time embedding."""
+
+    def __init__(self, dim: int, cond_dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.norm = RMSNorm(dim, eps)
+        self.modulation = nn.Linear(cond_dim, 2 * dim)
+        # Start as an ordinary RMSNorm, which makes tiny-data training stable.
+        nn.init.zeros_(self.modulation.weight)
+        nn.init.zeros_(self.modulation.bias)
+
+    def forward(self, x: Tensor, condition: Tensor) -> Tensor:
+        shift, scale = self.modulation(condition).chunk(2, dim=-1)
+        return self.norm(x) * (1 + scale[:, None, :]) + shift[:, None, :]
+
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        fourier_dim = max(8, dim // 2)
+        if fourier_dim % 2:
+            fourier_dim += 1
+        self.fourier_dim = fourier_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(fourier_dim, dim), nn.SiLU(), nn.Linear(dim, dim)
+        )
+
+    def forward(self, timestep: Tensor) -> Tensor:
+        timestep = timestep.reshape(-1).float()
+        half = self.fourier_dim // 2
+        frequencies = torch.exp(
+            torch.linspace(math.log(1.0), math.log(1000.0), half,
+                           device=timestep.device, dtype=timestep.dtype)
+        )
+        angles = 2 * math.pi * timestep[:, None] * frequencies[None, :]
+        embedding = torch.cat((angles.sin(), angles.cos()), dim=-1)
+        return self.mlp(embedding)
+
+
+class AttentionWeights(nn.Module):
+    """One expert's independent Q/K/V/output projections."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.q = nn.Linear(dim, dim, bias=False)
+        self.k = nn.Linear(dim, dim, bias=False)
+        self.v = nn.Linear(dim, dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim: int, ratio: int, dropout: float) -> None:
+        super().__init__()
+        hidden = dim * ratio
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden),
+            nn.GELU(approximate="tanh"),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.net(x)
+
+
+def _attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    n_heads: int,
+    allowed: Optional[Tensor],
+    dropout: float,
+    training: bool,
+) -> Tensor:
+    batch, q_len, width = query.shape
+    k_len = key.shape[1]
+    head_dim = width // n_heads
+
+    def split(x: Tensor) -> Tensor:
+        return x.reshape(batch, -1, n_heads, head_dim).transpose(1, 2)
+
+    q, k, v = split(query), split(key), split(value)
+    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim)
+    if allowed is not None:
+        # allowed is [B or 1, Q, K]; True means the edge exists.
+        scores = scores.masked_fill(~allowed[:, None, :, :], torch.finfo(scores.dtype).min)
+    probabilities = scores.float().softmax(dim=-1).to(scores.dtype)
+    probabilities = F.dropout(probabilities, p=dropout, training=training)
+    output = torch.matmul(probabilities, v)
+    return output.transpose(1, 2).contiguous().reshape(batch, q_len, width)
+
+
+class DualExpertBlock(nn.Module):
+    def __init__(self, config: Pi05Config) -> None:
+        super().__init__()
+        d = config.d_model
+        self.n_heads = config.n_heads
+        self.dropout = config.dropout
+
+        # The two streams never share Transformer parameters.
+        self.vlm_norm1 = RMSNorm(d, config.rms_eps)
+        self.vlm_attention = AttentionWeights(d)
+        self.vlm_norm2 = RMSNorm(d, config.rms_eps)
+        self.vlm_mlp = MLP(d, config.mlp_ratio, config.dropout)
+
+        self.action_norm1 = AdaRMSNorm(d, d, config.rms_eps)
+        self.action_attention = AttentionWeights(d)
+        self.action_norm2 = AdaRMSNorm(d, d, config.rms_eps)
+        self.action_mlp = MLP(d, config.mlp_ratio, config.dropout)
+
+    def forward(
+        self,
+        prefix: Tensor,
+        action: Tensor,
+        time_condition: Tensor,
+        prefix_edges: Tensor,
+        action_edges: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        prefix_input = self.vlm_norm1(prefix)
+        p = self.vlm_attention
+        prefix_update = _attention(
+            p.q(prefix_input), p.k(prefix_input), p.v(prefix_input),
+            self.n_heads, prefix_edges, self.dropout, self.training
+        )
+        prefix = prefix + p.out(prefix_update)
+        prefix = prefix + self.vlm_mlp(self.vlm_norm2(prefix))
+
+        action_input = self.action_norm1(action, time_condition)
+        a = self.action_attention
+        # Prefix KV uses the VLM expert's parameters; action KV uses the action
+        # expert's parameters.  Thus action can read prefix without allowing a
+        # reverse information path from action into the prefix stream.
+        key = torch.cat((p.k(prefix_input), a.k(action_input)), dim=1)
+        value = torch.cat((p.v(prefix_input), a.v(action_input)), dim=1)
+        action_update = _attention(
+            a.q(action_input), key, value, self.n_heads, action_edges,
+            self.dropout, self.training
+        )
+        action = action + a.out(action_update)
+        action = action + self.action_mlp(self.action_norm2(action, time_condition))
+        return prefix, action
+
+
+class Pi05Mini(nn.Module):
+    """Minimal pi0.5-style model for architecture validation/overfitting."""
+
+    def __init__(self, config: Optional[Pi05Config] = None) -> None:
+        super().__init__()
+        self.config = config or Pi05Config()
+        c = self.config
+        self.patch_encoder = nn.Conv2d(
+            c.image_channels, c.d_model, c.patch_size, stride=c.patch_size
+        )
+        self.token_embedding = nn.Embedding(c.vocab_size, c.d_model)
+        self.action_projection = nn.Linear(c.action_dim, c.d_model)
+        self.prefix_position = nn.Parameter(
+            torch.randn(1, c.max_prefix_tokens, c.d_model) * 0.01
+        )
+        self.action_position = nn.Parameter(
+            torch.randn(1, c.max_action_tokens, c.d_model) * 0.01
+        )
+        self.time_embedding = TimeEmbedding(c.d_model)
+        self.layers = nn.ModuleList(DualExpertBlock(c) for _ in range(c.n_layers))
+        self.prefix_norm = RMSNorm(c.d_model, c.rms_eps)
+        self.action_norm = AdaRMSNorm(c.d_model, c.d_model, c.rms_eps)
+        self.text_head = nn.Linear(c.d_model, c.vocab_size, bias=False)
+        self.vector_field_head = nn.Linear(c.d_model, c.action_dim)
+        self.text_head.weight = self.token_embedding.weight
+
+    def encode_prefix(self, pixel_values: Tensor, input_ids: Tensor) -> tuple[Tensor, int]:
+        patches = self.patch_encoder(pixel_values).flatten(2).transpose(1, 2)
+        text_tokens = self.token_embedding(input_ids)
+        # pi0.5 represents state in the VLM vocabulary.  There is deliberately
+        # no separate continuous-state projection in this model.
+        prefix = torch.cat((patches, text_tokens), dim=1)
+        if prefix.shape[1] > self.config.max_prefix_tokens:
+            raise ValueError("prefix exceeds configured maximum length")
+        prefix = prefix + self.prefix_position[:, : prefix.shape[1]]
+        return prefix, patches.shape[1]
+
+    def _masks(
+        self,
+        batch: int,
+        prefix_length: int,
+        action_length: int,
+        text_offset: int,
+        attention_mask: Optional[Tensor],
+        action_prefix_mask: Optional[Tensor],
+        generated_token_start: Optional[int],
+        device: torch.device,
+    ) -> tuple[Tensor, Tensor]:
+        valid = torch.ones(batch, prefix_length, dtype=torch.bool, device=device)
+        if attention_mask is not None:
+            valid[:, text_offset:] = attention_mask.bool()
+        # Image plus observation/instruction/state tokens form a fully
+        # bidirectional observation block.  Tokens generated after that block
+        # (subtask and FAST actions) read the observation and earlier generated
+        # tokens causally; observation queries cannot read generated tokens.
+        generated_start = prefix_length if generated_token_start is None else text_offset + generated_token_start
+        if not text_offset <= generated_start <= prefix_length:
+            raise ValueError("generated_token_start must index input_ids")
+        q = torch.arange(prefix_length, device=device)[:, None]
+        k = torch.arange(prefix_length, device=device)[None, :]
+        structural = torch.where(q < generated_start, k < generated_start, k <= q)
+        prefix_edges = structural[None] & valid[:, None, :]
+        # Padded prefix queries retain a harmless self edge to avoid all -inf
+        # rows; their outputs are not used by the text loss.
+        query_valid = valid[:, :, None]
+        identity = torch.eye(prefix_length, dtype=torch.bool, device=device)[None]
+        prefix_edges = torch.where(query_valid, prefix_edges, identity)
+
+        action_valid = valid
+        if action_prefix_mask is not None:
+            supplied = action_prefix_mask.bool()
+            if supplied.shape[1] == prefix_length - text_offset:
+                supplied = torch.cat(
+                    (torch.ones(batch, text_offset, dtype=torch.bool, device=device), supplied),
+                    dim=1,
+                )
+            if supplied.shape != valid.shape:
+                raise ValueError("action_prefix_mask must cover text tokens or the full prefix")
+            action_valid = action_valid & supplied
+        prefix_keys = action_valid[:, None, :].expand(batch, action_length, prefix_length)
+        action_keys = torch.ones(
+            batch, action_length, action_length, dtype=torch.bool, device=device
+        )
+        action_edges = torch.cat((prefix_keys, action_keys), dim=-1)
+        return prefix_edges, action_edges
+
+    def forward(
+        self,
+        pixel_values: Tensor,
+        input_ids: Tensor,
+        noisy_actions: Tensor,
+        timesteps: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        action_prefix_mask: Optional[Tensor] = None,
+        generated_token_start: Optional[int] = None,
+        text_labels: Optional[Tensor] = None,
+        target_velocity: Optional[Tensor] = None,
+        text_loss_weight: float = 1.0,
+        flow_loss_weight: float = 1.0,
+    ) -> Pi05Output:
+        """Run both experts.
+
+        ``timesteps`` is a scalar per batch item in [0, 1].  ``text_labels``
+        corresponds only to ``input_ids`` (not image prefix positions)
+        and may contain -100.  ``target_velocity`` has the same shape as
+        ``noisy_actions``.
+        """
+        if noisy_actions.ndim != 3:
+            raise ValueError("noisy_actions must have shape [batch, horizon, action_dim]")
+        if noisy_actions.shape[1] > self.config.max_action_tokens:
+            raise ValueError("action horizon exceeds configured maximum")
+
+        prefix, text_offset = self.encode_prefix(pixel_values, input_ids)
+        horizon = noisy_actions.shape[1]
+        action = self.action_projection(noisy_actions)
+        action = action + self.action_position[:, :horizon]
+        condition = self.time_embedding(timesteps).to(action.dtype)
+        prefix_edges, action_edges = self._masks(
+            prefix.shape[0], prefix.shape[1], horizon, text_offset,
+            attention_mask, action_prefix_mask, generated_token_start, prefix.device
+        )
+        for layer in self.layers:
+            prefix, action = layer(
+                prefix, action, condition, prefix_edges, action_edges
+            )
+
+        prefix_hidden = self.prefix_norm(prefix)
+        action_hidden = self.action_norm(action, condition)
+        # Return logits only for actual text positions, which keeps downstream
+        # next-token loss/generation code uncomplicated.
+        text_logits = self.text_head(prefix_hidden[:, text_offset:])
+        vector_field = self.vector_field_head(action_hidden)
+
+        text_loss = None
+        if text_labels is not None:
+            text_loss = F.cross_entropy(
+                text_logits.reshape(-1, text_logits.shape[-1]),
+                text_labels.reshape(-1), ignore_index=-100
+            )
+        flow_loss = None
+        if target_velocity is not None:
+            flow_loss = F.mse_loss(vector_field, target_velocity)
+        loss = None
+        if text_loss is not None or flow_loss is not None:
+            loss = sum(
+                component for component in (
+                    None if text_loss is None else text_loss_weight * text_loss,
+                    None if flow_loss is None else flow_loss_weight * flow_loss,
+                ) if component is not None
+            )
+        return Pi05Output(
+            text_logits=text_logits,
+            vector_field=vector_field,
+            prefix_hidden=prefix_hidden,
+            action_hidden=action_hidden,
+            text_token_offset=text_offset,
+            loss=loss,
+            text_loss=text_loss,
+            flow_loss=flow_loss,
+        )
+
+
+def flow_matching_sample(actions: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+    """Construct pi-style conditional-flow training targets.
+
+    The path is x_t = (1-t) * noise + t * action, with constant target vector
+    field ``action - noise``.
+    """
+    batch = actions.shape[0]
+    timestep = torch.rand(batch, device=actions.device, dtype=actions.dtype)
+    noise = torch.randn_like(actions)
+    broadcast_t = timestep.reshape(batch, *([1] * (actions.ndim - 1)))
+    noisy_actions = (1 - broadcast_t) * noise + broadcast_t * actions
+    target_velocity = actions - noise
+    return noisy_actions, timestep, target_velocity

src/pi05mini/runtime.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+
+def load_config(path: str | Path) -> dict:
+    with Path(path).open("r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def choose_device(requested: str = "auto") -> torch.device:
+    if requested != "auto":
+        return torch.device(requested)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def write_json(path: str | Path, value: object) -> None:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(json.dumps(value, indent=2, ensure_ascii=False), encoding="utf-8")

src/pi05mini/tokenizer.py

@@ -0,0 +1,285 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Iterable
+
+import torch
+
+
+@dataclass(frozen=True)
+class CoordinateRange:
+    minimum: float
+    maximum: float
+
+    def __post_init__(self) -> None:
+        if self.maximum <= self.minimum:
+            raise ValueError("coordinate maximum must be greater than minimum")
+
+
+class CoordinateTokenizer:
+    """Uniformly discretize XY state values into the literal token range 0..255."""
+
+    bins = 256
+
+    def __init__(
+        self,
+        x_range: tuple[float, float] = (-10.0, 10.0),
+        y_range: tuple[float, float] = (-5.0, 5.0),
+    ) -> None:
+        self.ranges = (
+            CoordinateRange(*x_range),
+            CoordinateRange(*y_range),
+        )
+
+    def encode(self, coordinates: torch.Tensor | Iterable[Iterable[float]]) -> torch.Tensor:
+        values = torch.as_tensor(coordinates, dtype=torch.float32)
+        if values.shape[-1] != 2:
+            raise ValueError("coordinates must end in an XY dimension of size 2")
+        low = values.new_tensor([axis.minimum for axis in self.ranges])
+        high = values.new_tensor([axis.maximum for axis in self.ranges])
+        unit = ((values - low) / (high - low)).clamp(0.0, 1.0)
+        return torch.round(unit * (self.bins - 1)).to(torch.long)
+
+    def decode(self, tokens: torch.Tensor | Iterable[Iterable[int]]) -> torch.Tensor:
+        encoded = torch.as_tensor(tokens, dtype=torch.long)
+        if encoded.shape[-1] != 2:
+            raise ValueError("coordinate tokens must end in an XY dimension of size 2")
+        if encoded.numel() and (encoded.min() < 0 or encoded.max() >= self.bins):
+            raise ValueError("coordinate tokens must be in the inclusive range 0..255")
+        low = torch.tensor([axis.minimum for axis in self.ranges], dtype=torch.float32)
+        high = torch.tensor([axis.maximum for axis in self.ranges], dtype=torch.float32)
+        return low + encoded.float() / (self.bins - 1) * (high - low)
+
+
+class SubtaskVocabulary:
+    """Small semantic vocabulary standing in for pi0.5's generated language subtask."""
+
+    labels = (
+        "keep_lane",
+        "steer_left",
+        "steer_right",
+        "slow_down",
+        "stop",
+        "accelerate",
+        "unknown",
+    )
+
+    def __init__(self, first_token_id: int = 260) -> None:
+        self.first_token_id = int(first_token_id)
+        self._label_to_token = {
+            label: self.first_token_id + index for index, label in enumerate(self.labels)
+        }
+        self._token_to_label = {token: label for label, token in self._label_to_token.items()}
+
+    @property
+    def size(self) -> int:
+        return len(self.labels)
+
+    @property
+    def token_ids(self) -> tuple[int, ...]:
+        return tuple(self._label_to_token[label] for label in self.labels)
+
+    def classify(self, reasoning: str) -> str:
+        text = reasoning.casefold()
+        # More specific commands precede generic lane/follow language.
+        if any(term in text for term in ("stop", "halt", "stationary")):
+            return "stop"
+        if any(term in text for term in ("slow", "deceler", "reduce speed", "brake")):
+            return "slow_down"
+        if any(term in text for term in ("accelerat", "speed up")):
+            return "accelerate"
+        if any(term in text for term in ("left", "counterclockwise")):
+            return "steer_left"
+        if any(term in text for term in ("right", "clockwise")):
+            return "steer_right"
+        if any(term in text for term in ("lane", "straight", "follow", "continue", "proceed")):
+            return "keep_lane"
+        return "unknown"
+
+    def encode(self, label_or_reasoning: str) -> int:
+        label = (
+            label_or_reasoning
+            if label_or_reasoning in self._label_to_token
+            else self.classify(label_or_reasoning)
+        )
+        return self._label_to_token[label]
+
+    def decode(self, token_id: int | torch.Tensor) -> str:
+        token = int(token_id)
+        try:
+            return self._token_to_label[token]
+        except KeyError as error:
+            raise ValueError(f"unknown subtask token: {token}") from error
+
+
+class MiniTokenizer:
+    """Compose image/state, semantic, and FAST-lite action tokens.
+
+    Actions are z-score-normalized by the dataset before reaching this
+    tokenizer.  A complete ``[horizon, 2]`` chunk is transformed along time
+    with an orthonormal DCT-II.  All coefficients are retained and serialized
+    in frequency-major XY order, so low-frequency motion appears first in the
+    token stream.  The coefficients are divided by ``sqrt(horizon)`` before
+    being uniformly quantized over ``[-4, 4]`` into 64 shared tokens.  This
+    keeps the fixed quantization interval valid for any horizon when the input
+    actions have first been clipped to the same interval.
+
+    This intentionally implements only the central transform/order idea from
+    FAST.  It does not learn dataset-specific coefficient bins or drop trailing
+    high-frequency coefficients.
+    """
+
+    BOS = 256
+    IMAGE = 257
+    STATE = 258
+    SUBTASK_QUERY = 259
+    ACTION_START = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
+    ACTION_END = ACTION_START + 1
+    ACTION_TOKEN_START = ACTION_END + 1
+    ACTION_BINS = 64
+    ACTION_TOKEN_END = ACTION_TOKEN_START + ACTION_BINS - 1
+    ACTION_MINIMUM = -4.0
+    ACTION_MAXIMUM = 4.0
+
+    def __init__(
+        self,
+        coordinate_tokenizer: CoordinateTokenizer | None = None,
+        subtasks: SubtaskVocabulary | None = None,
+    ) -> None:
+        self.coordinates = coordinate_tokenizer or CoordinateTokenizer()
+        self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1)
+
+    @property
+    def vocab_size(self) -> int:
+        return max(max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1
+
+    @property
+    def action_token_ids(self) -> range:
+        """Contiguous vocabulary interval reserved for quantized DCT coefficients."""
+
+        return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)
+
+    def build_prefix(self, history_xy: torch.Tensor) -> torch.Tensor:
+        coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1)
+        specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
+        query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
+        return torch.cat((specials, coordinate_tokens.cpu(), query), dim=0)
+
+    def encode_subtask(self, reasoning: str) -> torch.Tensor:
+        return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long)
+
+    @staticmethod
+    def _orthonormal_dct_matrix(
+        horizon: int,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Return the orthonormal DCT-II matrix ``[frequency, time]``."""
+
+        if horizon <= 0:
+            raise ValueError("DCT horizon must be positive")
+        time = torch.arange(horizon, device=device, dtype=dtype)
+        frequency = torch.arange(horizon, device=device, dtype=dtype).unsqueeze(1)
+        matrix = torch.cos(math.pi / horizon * (time + 0.5) * frequency)
+        matrix[0] *= math.sqrt(1.0 / horizon)
+        if horizon > 1:
+            matrix[1:] *= math.sqrt(2.0 / horizon)
+        return matrix
+
+    def dct_actions(
+        self, actions: torch.Tensor | Iterable[Iterable[float]]
+    ) -> torch.Tensor:
+        """Apply an orthonormal DCT-II along an action chunk's time axis."""
+
+        values = torch.as_tensor(actions, dtype=torch.float32)
+        if values.ndim != 2 or values.shape[-1] != 2:
+            raise ValueError("actions must have shape [horizon, 2]")
+        if values.shape[0] == 0:
+            raise ValueError("actions must contain at least one timestep")
+        transform = self._orthonormal_dct_matrix(
+            values.shape[0], device=values.device, dtype=values.dtype
+        )
+        return transform @ values
+
+    def idct_actions(
+        self, coefficients: torch.Tensor | Iterable[Iterable[float]]
+    ) -> torch.Tensor:
+        """Invert :meth:`dct_actions` using the transpose of its basis."""
+
+        values = torch.as_tensor(coefficients, dtype=torch.float32)
+        if values.ndim != 2 or values.shape[-1] != 2:
+            raise ValueError("DCT coefficients must have shape [horizon, 2]")
+        if values.shape[0] == 0:
+            raise ValueError("DCT coefficients must contain at least one frequency")
+        transform = self._orthonormal_dct_matrix(
+            values.shape[0], device=values.device, dtype=values.dtype
+        )
+        return transform.transpose(0, 1) @ values
+
+    def encode_actions(
+        self, actions: torch.Tensor | Iterable[Iterable[float]]
+    ) -> torch.Tensor:
+        """DCT-transform and encode normalized actions as a marked token stream."""
+
+        values = torch.as_tensor(actions, dtype=torch.float32)
+        if values.ndim != 2 or values.shape[-1] != 2:
+            raise ValueError("actions must have shape [horizon, 2]")
+        if values.shape[0] == 0:
+            raise ValueError("actions must contain at least one timestep")
+        values = values.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
+        coefficients = self.dct_actions(values) / math.sqrt(values.shape[0])
+        # Exact zero is useful for a stable shared center bin.  Tiny non-zero
+        # values here are only floating-point residue from the DCT basis.
+        coefficients = torch.where(
+            coefficients.abs() < 1e-6,
+            torch.zeros_like(coefficients),
+            coefficients,
+        )
+        unit = (
+            (
+                coefficients.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
+                - self.ACTION_MINIMUM
+            )
+            / (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
+        )
+        quantized = torch.round(unit * (self.ACTION_BINS - 1)).to(torch.long)
+        # Row-major flattening is frequency-major: DC XY, frequency-1 XY, ...
+        action_tokens = quantized.reshape(-1) + self.ACTION_TOKEN_START
+        boundaries = torch.tensor(
+            [self.ACTION_START, self.ACTION_END],
+            dtype=torch.long,
+            device=action_tokens.device,
+        )
+        return torch.cat((boundaries[:1], action_tokens, boundaries[1:]))
+
+    def decode_actions(self, tokens: torch.Tensor | Iterable[int]) -> torch.Tensor:
+        """Decode FAST-lite tokens back to normalized ``[horizon, 2]`` actions."""
+
+        encoded = torch.as_tensor(tokens, dtype=torch.long)
+        if encoded.ndim != 1:
+            raise ValueError("action tokens must be a one-dimensional sequence")
+        if (
+            encoded.numel() < 4
+            or encoded[0] != self.ACTION_START
+            or encoded[-1] != self.ACTION_END
+        ):
+            raise ValueError("action tokens must be enclosed by ACTION_START and ACTION_END")
+        payload = encoded[1:-1]
+        if payload.numel() % 2:
+            raise ValueError("action token payload must contain complete XY pairs")
+        if payload.min() < self.ACTION_TOKEN_START or payload.max() > self.ACTION_TOKEN_END:
+            raise ValueError(
+                "action token payload is outside the reserved action token interval"
+            )
+        quantized = payload - self.ACTION_TOKEN_START
+        normalized_coefficients = (
+            quantized.float()
+            / (self.ACTION_BINS - 1)
+            * (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
+            + self.ACTION_MINIMUM
+        )
+        horizon = payload.numel() // 2
+        coefficients = normalized_coefficients.reshape(horizon, 2) * math.sqrt(horizon)
+        return self.idct_actions(coefficients)

src/pi05mini/train.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+from .data import Pi05OverfitDataset, materialize_subset
+from .flow import FlowMatching
+from .model import Pi05Config, Pi05Mini
+from .runtime import choose_device, load_config, seed_everything, write_json
+
+
+def _collate(rows: list[dict]) -> dict:
+    tensor_keys = (
+        "image", "prefix_tokens", "subtask_target", "actions", "action_tokens"
+    )
+    result = {key: torch.stack([row[key] for row in rows]) for key in tensor_keys}
+    result["sample_id"] = [row["sample_id"] for row in rows]
+    return result
+
+
+def build_model(config: dict, dataset: Pi05OverfitDataset) -> Pi05Mini:
+    model_cfg = config["model"]
+    return Pi05Mini(Pi05Config(
+        vocab_size=dataset.tokenizer.vocab_size,
+        action_dim=2,
+        image_size=config["data"]["image_size"],
+        patch_size=model_cfg["patch_size"],
+        d_model=model_cfg["d_model"],
+        n_heads=model_cfg["n_heads"],
+        n_layers=model_cfg["n_layers"],
+        mlp_ratio=model_cfg["mlp_ratio"],
+        max_text_tokens=model_cfg["max_text_tokens"],
+        max_action_tokens=config["data"]["action_horizon"],
+        dropout=model_cfg["dropout"],
+    ))
+
+
+def train(config_path: str) -> Path:
+    cfg = load_config(config_path)
+    root = Path(__file__).resolve().parents[2]
+    seed_everything(cfg["seed"])
+    device = choose_device(cfg.get("device", "auto"))
+    manifest = materialize_subset(
+        source_manifest=cfg["data"]["source_manifest"],
+        output_dir=root / "data" / "overfit",
+        num_variants=cfg["data"]["num_variants"],
+    )
+    dataset = Pi05OverfitDataset(
+        manifest, image_size=cfg["data"]["image_size"],
+        action_horizon=cfg["data"]["action_horizon"]
+    )
+    loader = DataLoader(
+        dataset, batch_size=min(cfg["train"]["batch_size"], len(dataset)),
+        shuffle=True, collate_fn=_collate
+    )
+    model = build_model(cfg, dataset).to(device)
+    flow = FlowMatching(cfg["inference"]["flow_steps"])
+    output_dir = root / cfg["output_dir"]
+    output_dir.mkdir(parents=True, exist_ok=True)
+    history = []
+
+    def is_action_parameter(name: str) -> bool:
+        return any(part in name for part in (
+            "action_", "time_embedding", "vector_field_head"
+        ))
+
+    # Stage 1 follows pi0.5 pre-training: train the VLM/token pathway with
+    # discrete state and FAST-lite actions, while the flow expert is untouched.
+    stage1_params = [p for name, p in model.named_parameters() if not is_action_parameter(name)]
+    optimizer = torch.optim.AdamW(
+        stage1_params, lr=cfg["train"]["learning_rate"],
+        weight_decay=cfg["train"]["weight_decay"]
+    )
+
+    def next_batch(iterator):
+        try:
+            return next(iterator), iterator
+        except StopIteration:
+            iterator = iter(loader)
+            return next(iterator), iterator
+
+    iterator = iter(loader)
+    model.train()
+    for step in range(1, cfg["train"]["pretrain_steps"] + 1):
+        batch, iterator = next_batch(iterator)
+        image = batch["image"].to(device)
+        tokens = batch["prefix_tokens"].to(device)
+        actions = batch["actions"].to(device)
+        subtask = batch["subtask_target"].to(device)[:, None]
+        action_tokens = batch["action_tokens"].to(device)
+        targets = torch.cat((subtask, action_tokens), dim=1)
+        # Teacher-forced autoregressive stream. The final target is predicted
+        # from the preceding token and therefore is not part of model_tokens.
+        model_tokens = torch.cat((tokens, targets[:, :-1]), dim=1)
+        model_labels = torch.full_like(model_tokens, -100)
+        query_index = tokens.shape[1] - 1
+        model_labels[:, query_index:] = targets
+        # Flow tokens may read observation/state and the generated semantic
+        # subtask, but not teacher-forced discrete action tokens.
+        zero_action = torch.zeros_like(actions)
+        output = model(
+            image, model_tokens, zero_action,
+            torch.zeros(image.shape[0], device=device),
+            generated_token_start=tokens.shape[1],
+            text_labels=model_labels,
+        )
+        optimizer.zero_grad(set_to_none=True)
+        output.loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"])
+        optimizer.step()
+        if step == 1 or step % cfg["train"]["log_every"] == 0:
+            record = {
+                "stage": "pretrain",
+                "step": step,
+                "loss": float(output.loss.detach().cpu()),
+                "text_loss": float(output.text_loss.detach().cpu()),
+                "flow_loss": None,
+            }
+            history.append(record)
+            print(json.dumps(record))
+
+    torch.save(model.state_dict(), output_dir / "stage1_vlm.pt")
+
+    # Stage 2 follows post-training: add the still-random action expert and
+    # jointly retain token prediction while optimizing continuous flow.
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=cfg["train"]["learning_rate"],
+        weight_decay=cfg["train"]["weight_decay"]
+    )
+    iterator = iter(loader)
+    for step in range(1, cfg["train"]["posttrain_steps"] + 1):
+        batch, iterator = next_batch(iterator)
+        image = batch["image"].to(device)
+        tokens = batch["prefix_tokens"].to(device)
+        actions = batch["actions"].to(device)
+        subtask = batch["subtask_target"].to(device)[:, None]
+        action_tokens = batch["action_tokens"].to(device)
+        targets = torch.cat((subtask, action_tokens), dim=1)
+        model_tokens = torch.cat((tokens, targets[:, :-1]), dim=1)
+        model_labels = torch.full_like(model_tokens, -100)
+        model_labels[:, tokens.shape[1] - 1:] = targets
+        action_prefix_mask = torch.zeros_like(model_tokens, dtype=torch.bool)
+        action_prefix_mask[:, : tokens.shape[1] + 1] = True
+        noisy, timestep, velocity = flow.training_batch(actions)
+        output = model(
+            image, model_tokens, noisy, timestep,
+            action_prefix_mask=action_prefix_mask,
+            generated_token_start=tokens.shape[1],
+            text_labels=model_labels, target_velocity=velocity,
+            text_loss_weight=cfg["train"]["semantic_loss_weight"],
+            flow_loss_weight=cfg["train"]["flow_loss_weight"],
+        )
+        optimizer.zero_grad(set_to_none=True)
+        output.loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"])
+        optimizer.step()
+        if step == 1 or step % cfg["train"]["log_every"] == 0:
+            record = {
+                "stage": "posttrain", "step": step,
+                "loss": float(output.loss.detach().cpu()),
+                "text_loss": float(output.text_loss.detach().cpu()),
+                "flow_loss": float(output.flow_loss.detach().cpu()),
+            }
+            history.append(record)
+            print(json.dumps(record))
+    checkpoint = {
+        "model": model.state_dict(), "config": cfg,
+        "model_config": model.config.__dict__,
+        "action_statistics": dataset.action_normalizer.as_dict(),
+    }
+    torch.save(checkpoint, output_dir / "checkpoint.pt")
+    write_json(output_dir / "training_history.json", history)
+    write_json(output_dir / "summary.json", {
+        "device": str(device), "samples": len(dataset),
+        "parameters": sum(p.numel() for p in model.parameters()),
+        "initial": history[0], "final": history[-1],
+    })
+    return output_dir / "checkpoint.pt"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/overfit.yaml")
+    args = parser.parse_args()
+    train(args.config)
+
+
+if __name__ == "__main__":
+    main()

tests/test_action_tokens.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from pi05mini.tokenizer import MiniTokenizer
+
+
+def test_orthonormal_dct_round_trip() -> None:
+    tokenizer = MiniTokenizer()
+    actions = torch.tensor(
+        [
+            [-4.0, -2.75],
+            [-0.25, 0.0],
+            [1.375, 3.9],
+            [4.0, 2.125],
+        ]
+    )
+
+    coefficients = tokenizer.dct_actions(actions)
+    reconstructed = tokenizer.idct_actions(coefficients)
+
+    assert coefficients.shape == actions.shape
+    torch.testing.assert_close(reconstructed, actions, atol=1e-5, rtol=1e-5)
+
+
+def test_fast_lite_tokens_round_trip_with_bounded_quantization_error() -> None:
+    tokenizer = MiniTokenizer()
+    actions = torch.tensor(
+        [
+            [-3.5, -2.75],
+            [-0.25, 0.0],
+            [1.375, 3.5],
+            [3.75, 2.125],
+        ]
+    )
+
+    tokens = tokenizer.encode_actions(actions)
+    reconstructed = tokenizer.decode_actions(tokens)
+
+    bin_width = (
+        tokenizer.ACTION_MAXIMUM - tokenizer.ACTION_MINIMUM
+    ) / (tokenizer.ACTION_BINS - 1)
+    assert reconstructed.shape == actions.shape
+    # Orthonormal IDCT preserves total coefficient error.  Coefficients are
+    # quantized after division by sqrt(horizon), yielding this RMSE bound.
+    rmse = torch.sqrt(torch.mean((reconstructed - actions).square()))
+    assert float(rmse) <= bin_width / 2 * actions.shape[0] ** 0.5 + 1e-6
+
+
+def test_action_tokens_are_serialized_low_frequency_first() -> None:
+    tokenizer = MiniTokenizer()
+    horizon = 4
+    coefficients = torch.zeros(horizon, 2)
+    coefficients[0, 0] = 1.5 * horizon**0.5  # positive X DC coefficient
+    coefficients[1, 1] = -1.5 * horizon**0.5  # negative Y first harmonic
+    actions = tokenizer.idct_actions(coefficients)
+
+    payload = tokenizer.encode_actions(actions)[1:-1]
+    zero_token = tokenizer.ACTION_TOKEN_START + round(
+        (0.0 - tokenizer.ACTION_MINIMUM)
+        / (tokenizer.ACTION_MAXIMUM - tokenizer.ACTION_MINIMUM)
+        * (tokenizer.ACTION_BINS - 1)
+    )
+
+    # Frequency-major XY order is [DC-x, DC-y, f1-x, f1-y, ...].
+    assert payload[0].item() > zero_token
+    assert payload[1].item() == zero_token
+    assert payload[2].item() == zero_token
+    assert payload[3].item() < zero_token
+    assert torch.all(payload[4:] == zero_token)
+
+
+def test_action_token_ids_are_contiguous_reserved_vocabulary() -> None:
+    tokenizer = MiniTokenizer()
+    edge_actions = torch.tensor([[-4.0, 4.0], [0.0, 1.0]])
+    tokens = tokenizer.encode_actions(edge_actions)
+
+    assert tokens[0].item() == tokenizer.ACTION_START
+    assert tokens[-1].item() == tokenizer.ACTION_END
+    assert tokens[1:-1].min().item() >= tokenizer.ACTION_TOKEN_START
+    assert tokens[1:-1].max().item() <= tokenizer.ACTION_TOKEN_END
+    assert tuple(tokenizer.action_token_ids) == tuple(
+        range(tokenizer.ACTION_TOKEN_START, tokenizer.ACTION_TOKEN_END + 1)
+    )
+    assert tokenizer.vocab_size == tokenizer.ACTION_TOKEN_END + 1
+
+
+def test_decode_actions_rejects_missing_boundaries_and_non_action_ids() -> None:
+    tokenizer = MiniTokenizer()
+
+    with pytest.raises(ValueError, match="enclosed"):
+        tokenizer.decode_actions(torch.tensor([tokenizer.ACTION_TOKEN_START] * 2))
+    with pytest.raises(ValueError, match="reserved"):
+        tokenizer.decode_actions(
+            torch.tensor(
+                [
+                    tokenizer.ACTION_START,
+                    tokenizer.BOS,
+                    tokenizer.ACTION_TOKEN_START,
+                    tokenizer.ACTION_END,
+                ]
+            )
+        )

tests/test_data.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from pi05mini.data import Pi05OverfitDataset, load_manifest, materialize_subset
+from pi05mini.tokenizer import CoordinateTokenizer, MiniTokenizer, SubtaskVocabulary
+
+
+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)
+
+    image_path = image_dir / "sample.jpg"
+    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_path)
+
+    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)
+    trajectory_path = trajectory_dir / "sample.npz"
+    np.savez_compressed(trajectory_path, history_xyz=history, future_xyz=future)
+
+    manifest = processed / "overfit_manifest.jsonl"
+    row = {
+        "clip_id": "clip-a",
+        "t0_us": 1234,
+        "reasoning": "Steer left through the construction lane.",
+        "image_path": "data/processed/images/sample.jpg",
+        "trajectory_path": "data/processed/trajectories/sample.npz",
+    }
+    manifest.write_text(json.dumps(row) + "\n", encoding="utf-8")
+    return manifest
+
+
+def test_coordinate_tokens_are_literal_bytes_and_round_trip() -> None:
+    tokenizer = CoordinateTokenizer(x_range=(-10.0, 10.0), y_range=(-5.0, 5.0))
+    coordinates = torch.tensor([[-10.0, -5.0], [0.0, 0.0], [10.0, 5.0]])
+    tokens = tokenizer.encode(coordinates)
+
+    assert tokens.dtype == torch.long
+    assert tokens.tolist()[0] == [0, 0]
+    assert tokens.tolist()[-1] == [255, 255]
+    assert 0 <= int(tokens.min()) <= int(tokens.max()) <= 255
+    decoded = tokenizer.decode(tokens)
+    assert torch.allclose(decoded, coordinates, atol=10.1 / 255.0)
+
+
+def test_subtask_vocabulary_and_prefix_contract() -> None:
+    tokenizer = MiniTokenizer()
+    history = torch.zeros(8, 2)
+    prefix = tokenizer.build_prefix(history)
+
+    assert prefix.shape == (20,)
+    assert prefix[:3].tolist() == [tokenizer.BOS, tokenizer.IMAGE, tokenizer.STATE]
+    assert prefix[-1].item() == tokenizer.SUBTASK_QUERY
+    assert prefix[3:-1].min() >= 0 and prefix[3:-1].max() <= 255
+    target = tokenizer.encode_subtask("Steer left around the cones")
+    assert tokenizer.subtasks.decode(target) == "steer_left"
+    assert SubtaskVocabulary().classify("Slow down for a pedestrian") == "slow_down"
+
+
+def test_materialize_and_dataset_are_self_contained_and_deterministic(tmp_path: Path) -> None:
+    source_manifest = _make_source_project(tmp_path / "source")
+    copied_manifest = materialize_subset(source_manifest, tmp_path / "target", max_samples=1)
+    records = load_manifest(copied_manifest)
+
+    assert len(records) == 1
+    assert records[0].image_path.is_file()
+    assert records[0].trajectory_path.is_file()
+
+    dataset = Pi05OverfitDataset(
+        copied_manifest,
+        image_size=(32, 48),
+        augmentation_factor=3,
+        seed=19,
+    )
+    same_dataset = Pi05OverfitDataset(
+        copied_manifest,
+        image_size=(32, 48),
+        augmentation_factor=3,
+        seed=19,
+        action_statistics=dataset.action_normalizer.as_dict(),
+    )
+
+    assert len(dataset) == 3
+    original = dataset[0]
+    augmented = dataset[1]
+    repeated = same_dataset[1]
+    assert original["image"].shape == (3, 32, 48)
+    assert original["prefix_tokens"].shape == (20,)
+    assert original["actions"].shape == (16, 2)
+    assert original["subtask_target"].shape == ()
+    assert original["subtask_label"] == "steer_left"
+    assert original["reasoning"].startswith("Steer left")
+    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["future_xy"], repeated["future_xy"])
+    assert torch.equal(augmented["image"], repeated["image"])
+
+    reconstructed = dataset.action_normalizer.decode(
+        original["actions"], original["history_xy"]
+    )
+    assert torch.allclose(reconstructed, original["future_xy"], atol=1e-5)
+
+    statistics_path = dataset.save_action_statistics(tmp_path / "action_stats.json")
+    statistics = json.loads(statistics_path.read_text(encoding="utf-8"))
+    assert len(statistics["mean"]) == 2
+    assert len(statistics["scale"]) == 2
+    assert min(statistics["scale"]) >= 0.05

tests/test_flow.py

@@ -0,0 +1,85 @@
+import torch
+
+from pi05mini.flow import FlowMatching
+
+
+def _expand_batch_time(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """Expand a per-example time tensor over non-batch dimensions."""
+    return t.reshape(t.shape[0], *([1] * (target.ndim - 1)))
+
+
+def test_training_batch_obeys_linear_flow_equations() -> None:
+    flow = FlowMatching()
+    target = torch.tensor(
+        [
+            [[1.0, -2.0], [0.5, 3.0]],
+            [[-1.0, 0.0], [2.5, -0.5]],
+        ]
+    )
+    generator = torch.Generator().manual_seed(123)
+
+    x_t, t, target_velocity = flow.training_batch(target, generator=generator)
+
+    assert x_t.shape == target.shape
+    assert target_velocity.shape == target.shape
+    assert t.shape[0] == target.shape[0]
+    assert torch.all((t >= 0.0) & (t <= 1.0))
+
+    t_view = _expand_batch_time(t, target)
+    # v = target - noise, so the sampled noise is exactly recoverable.
+    noise = target - target_velocity
+    expected_x_t = (1.0 - t_view) * noise + t_view * target
+    torch.testing.assert_close(x_t, expected_x_t)
+
+
+def test_training_batch_is_reproducible_with_generator() -> None:
+    flow = FlowMatching()
+    target = torch.randn(3, 5, 2)
+
+    batch_a = flow.training_batch(
+        target, generator=torch.Generator().manual_seed(99)
+    )
+    batch_b = flow.training_batch(
+        target, generator=torch.Generator().manual_seed(99)
+    )
+
+    for tensor_a, tensor_b in zip(batch_a, batch_b):
+        torch.testing.assert_close(tensor_a, tensor_b)
+
+
+def test_loss_is_mean_squared_velocity_error() -> None:
+    flow = FlowMatching()
+    target_velocity = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]])
+    prediction = target_velocity + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]])
+
+    # Squared errors are [1, 1, 4, 0], whose mean is 1.5.
+    torch.testing.assert_close(flow.loss(prediction, target_velocity), torch.tensor(1.5))
+    torch.testing.assert_close(
+        flow.loss(target_velocity, target_velocity), torch.tensor(0.0)
+    )
+
+
+def test_euler_sampler_reaches_fixed_target_for_exact_vector_field() -> None:
+    flow = FlowMatching()
+    target = torch.tensor(
+        [
+            [[0.25, -0.5], [1.0, 2.0]],
+            [[-2.0, 0.75], [0.0, 0.5]],
+        ]
+    )
+
+    def exact_model(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        t_view = _expand_batch_time(t, x_t)
+        # Straight conditional path to target. The sampler evaluates t < 1.
+        return (target.to(x_t) - x_t) / (1.0 - t_view)
+
+    torch.manual_seed(5)
+    sampled = flow.sample(
+        exact_model,
+        shape=tuple(target.shape),
+        device=target.device,
+        steps=8,
+    )
+
+    torch.testing.assert_close(sampled, target, atol=1e-5, rtol=1e-5)
+

tests/test_model.py

@@ -0,0 +1,61 @@
+import torch
+
+from pi05mini.model import Pi05Config, Pi05Mini
+
+
+def tiny_model() -> Pi05Mini:
+    return Pi05Mini(Pi05Config(
+        vocab_size=300,
+        image_size=16,
+        patch_size=16,
+        d_model=32,
+        n_heads=4,
+        n_layers=1,
+        max_text_tokens=16,
+        max_action_tokens=4,
+    ))
+
+
+def test_pi05_has_no_continuous_state_projection() -> None:
+    model = tiny_model()
+    assert not hasattr(model, "state_projection")
+    output = model(
+        torch.randn(2, 3, 16, 16),
+        torch.randint(0, 300, (2, 7)),
+        torch.randn(2, 4, 2),
+        torch.rand(2),
+        generated_token_start=4,
+    )
+    assert output.text_logits.shape == (2, 7, 300)
+    assert output.vector_field.shape == (2, 4, 2)
+
+
+def test_observation_generated_and_flow_attention_partition() -> None:
+    model = tiny_model()
+    # One image patch followed by four observation tokens and three generated
+    # tokens. Flow may read observation plus the first generated (subtask)
+    # token, but not the two teacher-forced FAST tokens.
+    action_prefix_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 0, 0]], dtype=torch.bool)
+    prefix, action = model._masks(
+        batch=1,
+        prefix_length=8,
+        action_length=2,
+        text_offset=1,
+        attention_mask=torch.ones(1, 7, dtype=torch.bool),
+        action_prefix_mask=action_prefix_mask,
+        generated_token_start=4,
+        device=torch.device("cpu"),
+    )
+
+    generated_start = 5  # image patch + four observation text tokens
+    # Observation block is fully bidirectional and cannot see generated tokens.
+    assert prefix[0, 0, :generated_start].all()
+    assert not prefix[0, 0, generated_start:].any()
+    assert prefix[0, 4, :generated_start].all()
+    # Generated tokens see all observation tokens and only prior/self generated tokens.
+    assert prefix[0, 5, :6].all() and not prefix[0, 5, 6:].any()
+    assert prefix[0, 6, :7].all() and not prefix[0, 6, 7]
+    # Flow reads allowed prefix keys and all flow-action keys bidirectionally.
+    assert action[0, :, :6].all()
+    assert not action[0, :, 6:8].any()
+    assert action[0, :, 8:].all()