diff --git a/.gitignore b/.gitignore
index ff92181..55471a9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,4 +6,5 @@ __pycache__/
.venv/
data/
outputs/**/*.pt
+outputs/notebook/
.DS_Store
diff --git a/README.md b/README.md
index beba5ec..424db9b 100644
--- a/README.md
+++ b/README.md
@@ -44,6 +44,22 @@ python -m pip install -e '.[dev]'
pytest -q
```
+## 一键 Notebook 复现
+
+完整的论文链接、架构拆解、数据准备、过拟合训练、checkpoint 重载、推理指标和轨迹图都在：
+
+[`notebooks/pi0_minimal_reproduction.ipynb`](notebooks/pi0_minimal_reproduction.ipynb)
+
+```bash
+python3.12 -m venv .venv
+source .venv/bin/activate
+python -m pip install -U pip
+python -m pip install -e '.[notebook,dev]'
+python -m jupyter lab notebooks/pi0_minimal_reproduction.ipynb
+```
+
+Notebook 默认可直接 Run All：优先读取 `PI0_SOURCE_MANIFEST` 指定的已处理数据，其次尝试相邻 Alpamayo 项目，最后自动生成一个自包含 synthetic 样本。可通过 `PI0_PROFILE=smoke` 将训练缩短到 5 步，仅检查执行链路；默认 `full` 使用 600 步完成过拟合验证。
+
`tests/test_flow.py` 验证直线概率路径 `x_t=(1-t)·noise+t·action`、目标速度 `action-noise`、MSE 损失、随机可复现性和 Euler 积分。
## 过拟合训练与效果图
diff --git a/notebooks/pi0_minimal_reproduction.ipynb b/notebooks/pi0_minimal_reproduction.ipynb
new file mode 100644
index 0000000..6d9188a
--- /dev/null
+++ b/notebooks/pi0_minimal_reproduction.ipynb
@@ -0,0 +1,617 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "pi0-00",
+ "metadata": {},
+ "source": [
+ "# π₀ 最小复现：从论文架构到单样本过拟合\n",
+ "\n",
+ "本 Notebook 是本仓库的完整可执行入口。默认从上到下 Run All 即可完成环境安装、数据导入、结构检查、过拟合训练、checkpoint 重载、推理指标和轨迹绘图。\n",
+ "\n",
+ "## 目录\n",
+ "\n",
+ "- [1. 原论文、官方源码与复现边界](#paper)\n",
+ "- [2. π₀ 架构与阶段输入输出](#architecture)\n",
+ "- [3. 环境安装](#environment)\n",
+ "- [4. 运行配置、数据和权重导入](#data)\n",
+ "- [5. 样本与动作表示](#sample)\n",
+ "- [6. 三块 attention mask](#mask)\n",
+ "- [7. Flow matching 数学](#flow)\n",
+ "- [8. 模型实例化与张量形状](#model)\n",
+ "- [9. 最小过拟合训练](#train)\n",
+ "- [10. Checkpoint 重载、推理和实验结果](#results)\n",
+ "- [11. Run All 验收](#checklist)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-01",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 1. 原论文、官方源码与复现边界\n",
+ "\n",
+ "- 原论文：[π₀: A Vision-Language-Action Flow Model for General Robot Control](https://arxiv.org/abs/2410.24164)\n",
+ "- 官方项目介绍：[Physical Intelligence π₀](https://www.pi.website/blog/pi0)\n",
+ "- 官方源码：[Physical-Intelligence/openpi](https://github.com/Physical-Intelligence/openpi)\n",
+ "- 可选真实数据来源：[NVIDIA PhysicalAI Autonomous Vehicles](https://huggingface.co/datasets/nvidia/PhysicalAI-Autonomous-Vehicles)\n",
+ "\n",
+ "### 复现边界\n",
+ "\n",
+ "这个仓库验证的是 π₀ 最关键的信息流与优化闭环：视觉/语言 prefix、连续 state、独立 action expert、blockwise attention、flow matching 动作生成。它不是官方模型的数值复刻：\n",
+ "\n",
+ "- 小型 patch encoder 和 token embedding 代替预训练 PaliGemma。\n",
+ "- 固定语言 ID [1, 2] 代表一条极简指令，不含真实 tokenizer。\n",
+ "- 二维 16-step 驾驶轨迹代替机器人高维 action chunk。\n",
+ "- 参数量约 66 万，从头在一个样本的 4 个确定性增强版本上过拟合。\n",
+ "- 不下载也不声称使用官方 π₀ 权重。本地 checkpoint 只兼容本项目 toy model。\n",
+ "\n",
+ "成功代表结构和训练链路可运行、可记忆小数据，不代表官方能力、泛化能力或真实机器人安全性。"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-02",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 2. π₀ 架构与阶段输入输出\n",
+ "\n",
+ " 64×64 RGB ── patch encoder ─┐\n",
+ " 固定语言 ID ─ embedding ────┴─ VLM prefix expert ───────────────┐\n",
+ " 历史 XY [8,2] ─ state MLP ───── state token ────────────────────┤\n",
+ " noisy action [16,2] + time ─── action tokens ─ action expert ───┤\n",
+ " ↓\n",
+ " velocity [16,2]\n",
+ " ↓\n",
+ " Euler flow integration\n",
+ " ↓\n",
+ " future XY [16,2]\n",
+ "\n",
+ "| 阶段 | 输入 | 输出 | 本最小实现 |\n",
+ "|---|---|---|---|\n",
+ "| 视觉语言 prefix | RGB 图像、语言 token | image/language hidden tokens | 16 个 patch + 2 个固定语言 token |\n",
+ "| 连续 state 编码 | 历史轨迹 [8,2] | 1 个 state token | flatten 为 [16] 后过 MLP |\n",
+ "| Flow 加噪 | 真值 action [16,2]、噪声、t | noisy action、目标速度 | x_t=(1-t)noise+t action |\n",
+ "| Action/time 编码 | noisy action、Fourier time | 16 个 action tokens | 拼接后过 MLP |\n",
+ "| 双专家 Transformer | prefix、state、action | action hidden states | VLM/action 每层参数独立 |\n",
+ "| Vector-field head | action hidden | velocity [16,2] | flow velocity MSE |\n",
+ "| Euler 推理 | 高斯噪声、10 次预测 | normalized action | t=0 积分到 t=1 |\n",
+ "| 轨迹解码 | normalized XY delta | future XY [16,2] | 反归一化并累计 |\n",
+ "\n",
+ "三块可见性规则：\n",
+ "\n",
+ "| Query \\ Key | Image + Language | State | Action |\n",
+ "|---|---:|---:|---:|\n",
+ "| Image + Language | ✓ | ✗ | ✗ |\n",
+ "| State | ✓ | ✓ | ✗ |\n",
+ "| Action | ✓ | ✓ | ✓ |"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-03",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 3. 环境安装\n",
+ "\n",
+ "推荐在仓库根目录先创建 kernel：\n",
+ "\n",
+ " python3.12 -m venv .venv\n",
+ " source .venv/bin/activate\n",
+ " python -m pip install -U pip\n",
+ " python -m pip install -e '.[notebook,dev]'\n",
+ " python -m jupyter lab notebooks/pi0_minimal_reproduction.ipynb\n",
+ "\n",
+ "下面的第一个代码单元会再次执行 editable install，确保当前 kernel 引用的就是本仓库源码。已有 PyTorch 的环境不会重复下载相同版本。"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-04",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from __future__ import annotations\n",
+ "\n",
+ "import os\n",
+ "import platform\n",
+ "import subprocess\n",
+ "import sys\n",
+ "from pathlib import Path\n",
+ "\n",
+ "\n",
+ "def find_repo_root(start: Path) -> Path:\n",
+ " for candidate in (start.resolve(), *start.resolve().parents):\n",
+ " if (candidate / \"pyproject.toml\").is_file() and (candidate / \"src/pi0mini\").is_dir():\n",
+ " return candidate\n",
+ " raise FileNotFoundError(\"请从 pi0-minimal-validation 仓库内启动 Notebook\")\n",
+ "\n",
+ "\n",
+ "ROOT = find_repo_root(Path.cwd())\n",
+ "os.chdir(ROOT)\n",
+ "subprocess.run(\n",
+ " [sys.executable, \"-m\", \"pip\", \"install\", \"--disable-pip-version-check\", \"-e\", str(ROOT)],\n",
+ " check=True,\n",
+ ")\n",
+ "print(\"Repository:\", ROOT)\n",
+ "print(\"Python:\", sys.version.split()[0])\n",
+ "print(\"Platform:\", platform.platform())"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-05",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import copy\n",
+ "import json\n",
+ "import shutil\n",
+ "\n",
+ "import matplotlib.pyplot as plt\n",
+ "import numpy as np\n",
+ "import torch\n",
+ "import yaml\n",
+ "from IPython.display import Image as NotebookImage, display\n",
+ "from PIL import Image, ImageDraw\n",
+ "\n",
+ "from pi0mini.data import Pi0OverfitDataset, load_manifest, materialize_subset\n",
+ "from pi0mini.flow import FlowMatching\n",
+ "from pi0mini.infer import evaluate\n",
+ "from pi0mini.model import Pi0Config, Pi0Mini\n",
+ "from pi0mini.runtime import choose_device\n",
+ "from pi0mini.train import train\n",
+ "\n",
+ "print(\"torch:\", torch.__version__)\n",
+ "print(\"MPS available:\", torch.backends.mps.is_available())\n",
+ "print(\"CUDA available:\", torch.cuda.is_available())"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-06",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 4. 运行配置、数据和权重导入\n",
+ "\n",
+ "| 变量 | 取值 | 默认行为 |\n",
+ "|---|---|---|\n",
+ "| PI0_PROFILE | full / smoke / reuse | full，训练 600 步 |\n",
+ "| PI0_DATA_MODE | auto / local / synthetic | auto |\n",
+ "| PI0_SOURCE_MANIFEST | manifest 绝对路径 | 未设置时自动探测 |\n",
+ "| PI0_CHECKPOINT | 本项目 checkpoint 路径 | 未设置时从头训练 |\n",
+ "| PI0_DEVICE | auto / cpu / mps / cuda | auto |\n",
+ "\n",
+ "数据优先级为：用户指定 manifest → 相邻 Alpamayo 已处理数据 → 本仓库已物化数据 → synthetic fallback。\n",
+ "\n",
+ "真实 PAI 数据是 gated dataset。请先在数据页面接受许可并在外部完成下载/预处理，再把 PI0_SOURCE_MANIFEST 指向包含 image_path 和 trajectory_path 的 JSONL。Notebook 不打印或保存 HF_TOKEN，也不把原始数据提交 Git。\n",
+ "\n",
+ "checkpoint 必须来自本仓库相同架构。若导入 checkpoint，推理数据也应与训练时一致；Notebook 会至少保证数据格式和固定 data/overfit 路径可用。"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-07",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "PROFILE = os.environ.get(\"PI0_PROFILE\", \"full\").lower()\n",
+ "DATA_MODE = os.environ.get(\"PI0_DATA_MODE\", \"auto\").lower()\n",
+ "LOCAL_MANIFEST = os.environ.get(\"PI0_SOURCE_MANIFEST\", \"\").strip()\n",
+ "LOCAL_CHECKPOINT = os.environ.get(\"PI0_CHECKPOINT\", \"\").strip()\n",
+ "DEVICE_REQUEST = os.environ.get(\"PI0_DEVICE\", \"auto\").lower()\n",
+ "\n",
+ "if PROFILE not in {\"full\", \"smoke\", \"reuse\"}:\n",
+ " raise ValueError(\"PI0_PROFILE must be full, smoke, or reuse\")\n",
+ "if DATA_MODE not in {\"auto\", \"local\", \"synthetic\"}:\n",
+ " raise ValueError(\"PI0_DATA_MODE must be auto, local, or synthetic\")\n",
+ "\n",
+ "TRAIN_STEPS = {\"full\": 600, \"smoke\": 5, \"reuse\": 600}[PROFILE]\n",
+ "NOTEBOOK_OUTPUT = ROOT / \"outputs/notebook\"\n",
+ "STAGING = NOTEBOOK_OUTPUT / \"source_data\"\n",
+ "NOTEBOOK_OUTPUT.mkdir(parents=True, exist_ok=True)\n",
+ "print({\"profile\": PROFILE, \"data_mode\": DATA_MODE, \"steps_if_training\": TRAIN_STEPS,\n",
+ " \"device\": str(choose_device(DEVICE_REQUEST))})"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-08",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def valid_manifest(path: Path) -> bool:\n",
+ " try:\n",
+ " rows = load_manifest(path)\n",
+ " return bool(rows and rows[0].image_path.is_file() and rows[0].trajectory_path.is_file())\n",
+ " except Exception:\n",
+ " return False\n",
+ "\n",
+ "\n",
+ "def create_synthetic_source(root: Path) -> Path:\n",
+ " if root.exists():\n",
+ " shutil.rmtree(root)\n",
+ " image_dir = root / \"images\"\n",
+ " trajectory_dir = root / \"trajectories\"\n",
+ " image_dir.mkdir(parents=True)\n",
+ " trajectory_dir.mkdir(parents=True)\n",
+ " height, width = 96, 160\n",
+ " pixels = np.zeros((height, width, 3), dtype=np.uint8)\n",
+ " pixels[..., 0] = 35\n",
+ " pixels[..., 1] = np.linspace(45, 135, width, dtype=np.uint8)[None, :]\n",
+ " pixels[..., 2] = np.linspace(120, 35, height, dtype=np.uint8)[:, None]\n",
+ " road = Image.fromarray(pixels)\n",
+ " draw = ImageDraw.Draw(road)\n",
+ " draw.polygon([(58, 95), (75, 48), (85, 48), (105, 95)], fill=(55, 55, 60))\n",
+ " draw.line([(80, 90), (80, 53)], fill=(240, 220, 80), width=2)\n",
+ " road.save(image_dir / \"synthetic_front.jpg\")\n",
+ "\n",
+ " history_x = np.linspace(-7.0, 0.0, 8, dtype=np.float32)\n",
+ " history = np.stack([history_x, 0.01 * history_x, np.zeros_like(history_x)], axis=-1)\n",
+ " future_x = np.linspace(1.0, 32.0, 16, dtype=np.float32)\n",
+ " future = np.stack([future_x, 0.002 * future_x**2, np.zeros_like(future_x)], axis=-1)\n",
+ " np.savez_compressed(trajectory_dir / \"synthetic.npz\",\n",
+ " history_xyz=history.astype(np.float32),\n",
+ " future_xyz=future.astype(np.float32))\n",
+ " manifest = root / \"overfit_manifest.jsonl\"\n",
+ " manifest.write_text(json.dumps({\n",
+ " \"clip_id\": \"synthetic-road\", \"t0_us\": 0,\n",
+ " \"image_path\": \"images/synthetic_front.jpg\",\n",
+ " \"trajectory_path\": \"trajectories/synthetic.npz\",\n",
+ " }) + \"\\n\", encoding=\"utf-8\")\n",
+ " return manifest\n",
+ "\n",
+ "\n",
+ "selected = None\n",
+ "selected_label = None\n",
+ "if DATA_MODE != \"synthetic\":\n",
+ " candidates = []\n",
+ " if LOCAL_MANIFEST:\n",
+ " candidates.append((\"PI0_SOURCE_MANIFEST\", Path(LOCAL_MANIFEST).expanduser()))\n",
+ " candidates.extend([\n",
+ " (\"neighbor Alpamayo\", ROOT.parent.parent / \"alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl\"),\n",
+ " (\"existing materialized sample\", ROOT / \"data/overfit/overfit_manifest.jsonl\"),\n",
+ " ])\n",
+ " for label, candidate in candidates:\n",
+ " if valid_manifest(candidate):\n",
+ " selected, selected_label = candidate.resolve(), label\n",
+ " break\n",
+ "\n",
+ "if DATA_MODE == \"local\" and selected is None:\n",
+ " raise FileNotFoundError(\"PI0_DATA_MODE=local requires a valid PI0_SOURCE_MANIFEST\")\n",
+ "if selected is None:\n",
+ " SOURCE_MANIFEST = create_synthetic_source(STAGING)\n",
+ " selected_label = \"synthetic fallback\"\n",
+ "else:\n",
+ " if STAGING.exists():\n",
+ " shutil.rmtree(STAGING)\n",
+ " SOURCE_MANIFEST = materialize_subset(selected, STAGING)\n",
+ "\n",
+ "assert valid_manifest(SOURCE_MANIFEST)\n",
+ "print(\"Data source:\", selected_label)\n",
+ "print(\"Staged manifest:\", SOURCE_MANIFEST)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-09",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 5. 样本与动作表示\n",
+ "\n",
+ "Dataset 输出 image [3,64,64]、state [16]、归一化 actions [16,2]、history_xy [8,2] 和 future_xy [16,2]。下面验证 action encode/decode 可逆，并显示相机图像和轨迹真值。"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-10",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "dataset = Pi0OverfitDataset(SOURCE_MANIFEST, image_size=64, augmentation_factor=4,\n",
+ " history_steps=8, action_horizon=16)\n",
+ "sample = dataset[0]\n",
+ "decoded = dataset.action_normalizer.decode(sample[\"actions\"], sample[\"history_xy\"])\n",
+ "torch.testing.assert_close(decoded, sample[\"future_xy\"], atol=1e-5, rtol=0)\n",
+ "print({key: tuple(value.shape) for key, value in sample.items() if torch.is_tensor(value)})\n",
+ "print(\"action statistics:\", dataset.action_normalizer.as_dict())\n",
+ "\n",
+ "fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))\n",
+ "axes[0].imshow(sample[\"image\"].permute(1, 2, 0))\n",
+ "axes[0].set_title(\"64×64 visual observation\")\n",
+ "axes[0].axis(\"off\")\n",
+ "axes[1].plot(sample[\"history_xy\"][:, 0], sample[\"history_xy\"][:, 1], \"o-\", label=\"history\")\n",
+ "axes[1].plot(sample[\"future_xy\"][:, 0], sample[\"future_xy\"][:, 1], \"o-\", label=\"future target\")\n",
+ "axes[1].scatter([0], [0], c=\"black\", marker=\"s\", label=\"ego\")\n",
+ "axes[1].axis(\"equal\")\n",
+ "axes[1].grid(alpha=0.3)\n",
+ "axes[1].legend()\n",
+ "plt.show()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-11",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 6. 三块 attention mask\n",
+ "\n",
+ "Prefix 内部双向可见；state 可以读取 prefix 和自身，但不能偷看 noisy future action；action 可以读取 prefix、state 和完整 action chunk。"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-12",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "probe = Pi0Mini(Pi0Config(image_size=32, patch_size=16, max_text_tokens=4,\n",
+ " max_action_tokens=3, d_model=32, n_heads=4,\n",
+ " n_layers=1, time_feature_dim=8))\n",
+ "prefix_edges, suffix_edges = probe._masks(\n",
+ " batch=1, prefix_length=6, suffix_length=4, image_tokens=4,\n",
+ " attention_mask=torch.tensor([[1, 1]]), device=torch.device(\"cpu\"))\n",
+ "assert suffix_edges[0, 0, :7].all()\n",
+ "assert not suffix_edges[0, 0, 7:].any()\n",
+ "assert suffix_edges[0, 1:].all()\n",
+ "\n",
+ "plt.figure(figsize=(8, 3))\n",
+ "plt.imshow(suffix_edges[0].numpy(), cmap=\"Greys\", vmin=0, vmax=1)\n",
+ "plt.axvline(5.5, color=\"tab:red\", linewidth=1)\n",
+ "plt.axvline(6.5, color=\"tab:blue\", linewidth=1)\n",
+ "plt.yticks(range(4), [\"state\", \"action 0\", \"action 1\", \"action 2\"])\n",
+ "plt.xlabel(\"key: prefix (0..5) | state (6) | action (7..9)\")\n",
+ "plt.ylabel(\"suffix query\")\n",
+ "plt.title(\"Allowed attention edges\")\n",
+ "plt.colorbar()\n",
+ "plt.show()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-13",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 7. Flow matching 数学\n",
+ "\n",
+ "训练使用直线路径：\n",
+ "\n",
+ "$$x_t=(1-t)\\epsilon+t a, \\qquad v^*(x_t,t)=a-\\epsilon$$\n",
+ "\n",
+ "其中 a 是归一化 action，epsilon 是高斯噪声。模型回归 velocity；推理从 t=0 的噪声出发，用 Euler 积分到 t=1。"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-14",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "flow = FlowMatching(inference_steps=10)\n",
+ "target = torch.randn(4, 16, 2)\n",
+ "generator = torch.Generator().manual_seed(7)\n",
+ "x_t, t, target_velocity = flow.training_batch(target, generator)\n",
+ "noise = target - target_velocity\n",
+ "expected = (1 - t[:, None, None]) * noise + t[:, None, None] * target\n",
+ "torch.testing.assert_close(x_t, expected)\n",
+ "print(\"sampled action-time t:\", t.tolist())\n",
+ "print(\"x_t shape:\", tuple(x_t.shape), \"target velocity shape:\", tuple(target_velocity.shape))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-15",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 8. 模型实例化与张量形状\n",
+ "\n",
+ "先生成 Notebook 专用 runtime config，再用一个 batch 检查图像、语言、state、noisy action 到 vector field 的完整形状。"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-16",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "base_cfg = yaml.safe_load((ROOT / \"configs/overfit.yaml\").read_text(encoding=\"utf-8\"))\n",
+ "cfg = copy.deepcopy(base_cfg)\n",
+ "cfg[\"data\"][\"source_manifest\"] = str(SOURCE_MANIFEST.resolve())\n",
+ "cfg[\"train\"][\"steps\"] = TRAIN_STEPS\n",
+ "cfg[\"train\"][\"log_every\"] = 1 if PROFILE == \"smoke\" else cfg[\"train\"][\"log_every\"]\n",
+ "cfg[\"device\"] = DEVICE_REQUEST\n",
+ "cfg[\"output_dir\"] = \"outputs/notebook/overfit\"\n",
+ "RUNTIME_CONFIG = NOTEBOOK_OUTPUT / \"runtime_config.yaml\"\n",
+ "RUNTIME_CONFIG.write_text(yaml.safe_dump(cfg, sort_keys=False, allow_unicode=True), encoding=\"utf-8\")\n",
+ "print(RUNTIME_CONFIG.read_text(encoding=\"utf-8\"))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-17",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "mc = cfg[\"model\"]\n",
+ "model = Pi0Mini(Pi0Config(\n",
+ " vocab_size=16, action_dim=2, state_dim=dataset.state_dim,\n",
+ " image_size=cfg[\"data\"][\"image_size\"], patch_size=mc[\"patch_size\"],\n",
+ " d_model=mc[\"d_model\"], n_heads=mc[\"n_heads\"], n_layers=mc[\"n_layers\"],\n",
+ " mlp_ratio=mc[\"mlp_ratio\"], max_text_tokens=8,\n",
+ " max_action_tokens=cfg[\"data\"][\"action_horizon\"], dropout=mc[\"dropout\"]))\n",
+ "batch_size = 2\n",
+ "images = torch.stack([dataset[i][\"image\"] for i in range(batch_size)])\n",
+ "states = torch.stack([dataset[i][\"state\"] for i in range(batch_size)])\n",
+ "actions = torch.stack([dataset[i][\"actions\"] for i in range(batch_size)])\n",
+ "noisy, timesteps, velocity_target = flow.training_batch(actions)\n",
+ "out = model(images, torch.tensor([[1, 2]]).expand(batch_size, -1), states,\n",
+ " noisy, timesteps, target_velocity=velocity_target)\n",
+ "print(\"parameters:\", f\"{sum(p.numel() for p in model.parameters()):,}\")\n",
+ "print({\"prefix_hidden\": tuple(out.prefix_hidden.shape),\n",
+ " \"state_hidden\": tuple(out.state_hidden.shape),\n",
+ " \"action_hidden\": tuple(out.action_hidden.shape),\n",
+ " \"vector_field\": tuple(out.vector_field.shape), \"loss\": float(out.loss)})\n",
+ "del model, out"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-18",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 9. 最小过拟合训练\n",
+ "\n",
+ "- full：默认 600 步，用于验证 loss 下降和轨迹拟合。\n",
+ "- smoke：5 步，只验证 Run All 管线，不能代表收敛。\n",
+ "- reuse：优先复用 PI0_CHECKPOINT 或 outputs/notebook/overfit/checkpoint.pt；找不到时仍会训练。\n",
+ "\n",
+ "如果显式设置 PI0_CHECKPOINT，Notebook 不会覆盖它。训练生成的 checkpoint、数据 staging 和 runtime config 都在 gitignore 范围内。"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-19",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "expected_checkpoint = NOTEBOOK_OUTPUT / \"overfit/checkpoint.pt\"\n",
+ "provided_checkpoint = Path(LOCAL_CHECKPOINT).expanduser().resolve() if LOCAL_CHECKPOINT else None\n",
+ "if provided_checkpoint is not None:\n",
+ " if not provided_checkpoint.is_file():\n",
+ " raise FileNotFoundError(provided_checkpoint)\n",
+ " materialize_subset(SOURCE_MANIFEST, ROOT / \"data/overfit\")\n",
+ " checkpoint_path = provided_checkpoint\n",
+ " print(\"Using PI0_CHECKPOINT:\", checkpoint_path)\n",
+ "elif PROFILE == \"reuse\" and expected_checkpoint.is_file():\n",
+ " materialize_subset(SOURCE_MANIFEST, ROOT / \"data/overfit\")\n",
+ " checkpoint_path = expected_checkpoint\n",
+ " print(\"Reusing notebook checkpoint:\", checkpoint_path)\n",
+ "else:\n",
+ " checkpoint_path = train(str(RUNTIME_CONFIG))\n",
+ "\n",
+ "checkpoint_path = Path(checkpoint_path)\n",
+ "assert checkpoint_path.is_file()\n",
+ "print(\"Checkpoint:\", checkpoint_path)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-20",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 10. Checkpoint 重载、推理和实验结果\n",
+ "\n",
+ "evaluate() 会重新构造模型、加载 checkpoint，用固定噪声做 10-step Euler flow sampling，并对 4 个确定性增强样本分别计算 future trajectory MSE。\n",
+ "\n",
+ "本 Notebook 已在本地完成 600-step MPS Run All 验证：flow loss 从 2.262 降至 0.164，4 个增强样本平均 trajectory MSE 约 0.181 m²。具体数值会因设备、数据源和随机算子略有变化；下面显示当前这次 Run All 的真实结果。"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-21",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plot_path = evaluate(str(checkpoint_path))\n",
+ "output_dir = checkpoint_path.parent\n",
+ "summary = json.loads((output_dir / \"summary.json\").read_text(encoding=\"utf-8\"))\n",
+ "history = json.loads((output_dir / \"training_history.json\").read_text(encoding=\"utf-8\"))\n",
+ "evaluation = json.loads((output_dir / \"evaluation.json\").read_text(encoding=\"utf-8\"))\n",
+ "mean_mse = float(np.mean([row[\"trajectory_mse\"] for row in evaluation]))\n",
+ "print(json.dumps(summary, indent=2, ensure_ascii=False))\n",
+ "print(\"mean trajectory MSE:\", mean_mse)\n",
+ "\n",
+ "fig, ax = plt.subplots(figsize=(6, 3.5))\n",
+ "ax.plot([row[\"step\"] for row in history], [row[\"flow_loss\"] for row in history], \"o-\")\n",
+ "ax.set_yscale(\"log\")\n",
+ "ax.set_xlabel(\"training step\")\n",
+ "ax.set_ylabel(\"flow velocity MSE\")\n",
+ "ax.set_title(\"Overfit training curve\")\n",
+ "ax.grid(alpha=0.3)\n",
+ "plt.show()\n",
+ "display(NotebookImage(filename=str(plot_path)))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-22",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "print(\"| sample | trajectory MSE |\")\n",
+ "print(\"|---|---:|\")\n",
+ "for row in evaluation:\n",
+ " print(f\"| {row['sample_id']} | {row['trajectory_mse']:.6f} |\")\n",
+ "print(f\"| mean | {mean_mse:.6f} |\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "pi0-23",
+ "metadata": {},
+ "source": [
+ "\n",
+ "## 11. Run All 验收\n",
+ "\n",
+ "最后一个单元检查数据、配置、checkpoint、训练历史、推理指标和轨迹图。全部通过即表示当前环境完成了最小复现闭环。"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "pi0-24",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "assert SOURCE_MANIFEST.is_file()\n",
+ "assert RUNTIME_CONFIG.is_file()\n",
+ "assert checkpoint_path.is_file()\n",
+ "assert plot_path.is_file()\n",
+ "assert (output_dir / \"training_history.json\").is_file()\n",
+ "assert (output_dir / \"evaluation.json\").is_file()\n",
+ "assert np.isfinite(mean_mse)\n",
+ "assert len(evaluation) == len(dataset)\n",
+ "print(\"✅ π₀ minimal reproduction Run All completed\")\n",
+ "print(\"profile:\", PROFILE)\n",
+ "print(\"data:\", selected_label)\n",
+ "print(\"checkpoint:\", checkpoint_path)\n",
+ "print(\"trajectory plot:\", plot_path)"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python",
+ "version": "3.10"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/pyproject.toml b/pyproject.toml
index 20a1f04..26af438 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,6 +19,11 @@ dependencies = [
[project.optional-dependencies]
dev = ["pytest>=8.0"]
+notebook = [
+ "ipykernel>=6.29",
+ "jupyterlab>=4.2",
+ "nbconvert>=7.16",
+]
[project.scripts]
pi0mini-train = "pi0mini.train:main"