add end-to-end pi0 reproduction notebook

This commit is contained in:
seeker
2026-07-12 21:30:10 +08:00
parent 4fe0099f3f
commit 462c070dd9
4 changed files with 639 additions and 0 deletions
+16
View File
@@ -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 积分。
## 过拟合训练与效果图