seeker
618 lines
26 KiB
Plaintext
618 lines
26 KiB
Plaintext
{
|
||
"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×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": [
|
||
"<a id=\"environment\"></a>\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": [
|
||
"<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×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": [
|
||
"<a id=\"mask\"></a>\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": [
|
||
"<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，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": [
|
||
"<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：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": [
|
||
"<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，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": [
|
||
"<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
|
||
}
|