@@ -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": [
"<a id=\"paper\"></a>\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": [
"<a id=\"architecture\"></a>\n",
"## 2. π₀ 架构与阶段输入输出\n",
"\n",
" 64×
" 固定语言 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": [
"<a id=\"environment\"></a>\n",
"## 3. 环境安装\n",
"\n",
"推荐在仓库根目录先创建 kernel：
"\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": [
"<a id=\"data\"></a>\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": [
"<a id=\"sample\"></a>\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×
"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": [
"<a id=\"mask\"></a>\n",
"## 6. 三块 attention mask\n",
"\n",
"Prefix 内部双向可见；state 可以读取 prefix 和自身，但不能偷看 noisy future action；
]
},
{
"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": [
"<a id=\"flow\"></a>\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，
]
},
{
"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": [
"<a id=\"model\"></a>\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": [
"<a id=\"train\"></a>\n",
"## 9. 最小过拟合训练\n",
"\n",
"- full：默认 600 步，用于验证 loss 下降和轨迹拟合。\n",
"- smoke：
"- reuse：优先复用 PI0_CHECKPOINT 或 outputs/notebook/overfit/checkpoint.pt；找不到时仍会训练。\n",
"\n",
"如果显式设置 PI0_CHECKPOINT，
]
},
{
"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": [
"<a id=\"results\"></a>\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，
]
},
{
"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": [
"<a id=\"checklist\"></a>\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
}
seeker