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