π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。

运行测试

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

S
Description
Physical Intelligence π0.5 机器人 VLA 模型的最小复现。
Readme 336 KiB
Languages
Python 100%