seeker
140 lines
5.6 KiB
Markdown
140 lines
5.6 KiB
Markdown
# π0.5 最小复现
|
||
|
||
本项目是 Physical Intelligence π0.5 的最小复现，包含高层 subtask 生成、blockwise attention、FAST-lite 离散动作和 flow action expert。默认使用一个 PAI 驾驶样本及其增强变体，在单机上跑通训练、推理和轨迹对比。
|
||
|
||
## 实现内容
|
||
|
||
```text
|
||
图像 + 本体状态
|
||
│
|
||
▼
|
||
小型视觉/语言主干 ──→ 高层语义 subtask
|
||
│ │
|
||
└──── semantic prefix┘
|
||
│
|
||
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；
|
||
- 同时优化高层语义目标和低层连续动作目标。
|
||
- VLM 自回归输出高层 subtask 以及 64-bin 离散动作 token；
|
||
- 推理时分别还原 VLM token 轨迹和 flow/diffusion 轨迹，与真值三方比较。
|
||
|
||
为了单机快速验证而简化的部分：
|
||
|
||
- 小型自定义视觉编码器和 Transformer 取代 PaliGemma 2B；
|
||
- 隐藏宽度、层数、图像尺寸和 action horizon 大幅缩小；
|
||
- 使用固定小词表的 subtask 分类/生成，而非开放式语言理解；
|
||
- 使用二维驾驶轨迹，取代多机器人、多相机和高维关节控制；
|
||
- 默认使用 1 个 PAI 驾驶样本的 4 个视觉增强版本，没有机器人或网页联合训练数据；
|
||
- 使用正交 DCT-II、低频优先排列和 64-bin 量化实现 FAST-lite；保留全部频率系数，因此只验证表示与训练接口，不复现官方压缩率；
|
||
- Euler sampler 用固定步长，不包含面向部署的实时控制与安全系统。
|
||
|
||
因此，本项目适合验证架构接口和数学链路，不适合比较论文指标。
|
||
|
||
## 安装
|
||
|
||
建议使用 Python 3.10 或更新版本：
|
||
|
||
```bash
|
||
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，不需要下载预训练权重。默认配置会读取：
|
||
|
||
```text
|
||
../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl
|
||
```
|
||
|
||
如果该路径不同，请先修改 `configs/overfit.yaml` 的 `data.source_manifest`。
|
||
|
||
## 运行测试
|
||
|
||
```bash
|
||
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` 时采样可复现。
|
||
|
||
## 过拟合训练
|
||
|
||
```bash
|
||
cd pi05-minimal-validation
|
||
pi05mini-train --config configs/overfit.yaml
|
||
```
|
||
|
||
也可以直接使用模块入口：
|
||
|
||
```bash
|
||
python -m pi05mini.train --config configs/overfit.yaml
|
||
```
|
||
|
||
配置默认先进行 300 步离散 VLM/FAST-lite 预训练，再进行 500 步联合 post-training。Stage 1 不更新 flow action expert；Stage 2 联合优化 token CE 与 flow loss。CPU 可以运行；Apple Silicon 可自动选择 MPS，CUDA 可自动选择 GPU。
|
||
|
||
## 推理与效果图
|
||
|
||
```bash
|
||
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。
|
||
|
||
## 项目结构
|
||
|
||
```text
|
||
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 在此实现中能够端到端反向传播并记忆几个样本。它不能证明实现与官方模型数值等价，也不能证明在真实驾驶或机器人任务中有效。
|