Files

T

seeker 7e0267be1f add executable pi0.5 reproduction notebook

2026-07-12 21:22:25 +08:00

6.6 KiB

Raw Blame History

π0.5 Minimal Validation

这是一个 π0.5 风格的最小、可运行架构验证项目。它从相邻 Alpamayo 验证项目准备的 PAI 子集中读取 1 个驾驶样本，生成 4 个确定性视觉增强变体做过拟合，目标是检查高层语义预测、blockwise attention、action expert 与 flow matching action chunk 能否组成完整训练和推理链路。

它不是 Physical Intelligence 官方代码，也不声称复现论文中的规模、数据或机器人能力。

验证内容

图像 + 离散本体状态 + 总体任务指令
                  │
                  ▼
         高层语义 subtask 预测
                  │ predicted subtask
                  ▼
图像 + 离散本体状态 + 低层 subtask command
                  │
 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；
同时优化高层语义目标和低层连续动作目标。
高层调用显式读取总体任务指令；低层调用只读取 observation 与预测 subtask，对应论文分解 π(action | observation, subtask) · π(subtask | observation, task)；
VLM 自回归输出高层 subtask 以及 64-bin 离散动作 token；
推理时分别还原 VLM token 轨迹和 flow/diffusion 轨迹，与真值三方比较。

为了单机快速验证而简化的部分：

小型自定义视觉编码器和 Transformer 取代 PaliGemma 2B；
隐藏宽度、层数、图像尺寸和 action horizon 大幅缩小；
使用固定小词表的 task/subtask 分类与生成，而非开放式语言理解；
使用二维驾驶轨迹，取代多机器人、多相机和高维关节控制；
默认使用 1 个 PAI 驾驶样本的 4 个视觉增强版本，没有机器人或网页联合训练数据；
使用正交 DCT-II、低频优先排列和 64-bin 量化实现 FAST-lite；保留全部频率系数，因此只验证表示与训练接口，不复现官方压缩率；
Euler sampler 用固定步长，不包含面向部署的实时控制与安全系统。

因此，本项目适合验证架构接口和数学链路，不适合比较论文指标。

安装

建议使用 Python 3.10 或更新版本：

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，不需要下载预训练权重。默认配置会读取：

../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl

如果该路径不同，请先修改 configs/overfit.yaml 的 data.source_manifest。

一键 Notebook 复现

notebooks/pi05_minimal_reproduction.ipynb 提供从环境安装、论文与架构说明、数据/权重导入、测试、完整训练、推理到结果图的顺序执行流程。Notebook 支持本地 manifest、数据 zip URL 或自动生成的确定性最小样本；默认从零训练，因此不要求预训练权重。

运行测试

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 时采样可复现。

过拟合训练

cd pi05-minimal-validation
pi05mini-train --config configs/overfit.yaml

也可以直接使用模块入口：

python -m pi05mini.train --config configs/overfit.yaml

配置默认先进行 300 步离散 VLM/FAST-lite 预训练，再进行 500 步联合 post-training。Stage 1 不更新 flow action expert；Stage 2 使用同一模型的两个调用联合优化：总体任务条件下的 token CE，以及预测/teacher-forced subtask 条件下的 flow loss。CPU 可以运行；Apple Silicon 可自动选择 MPS，CUDA 可自动选择 GPU。

推理与效果图

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。

项目结构

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 在此实现中能够端到端反向传播并记忆几个样本。它不能证明实现与官方模型数值等价，也不能证明在真实驾驶或机器人任务中有效。

6.6 KiB Raw Blame History Unescape Escape