# π0.6 / RECAP 最小复现

本项目是 Physical Intelligence π0.6 / `π*0.6` RECAP 的最小复现。它在 π0.5 的高层 subtask、FAST 离散动作和 flow action expert 之上，加入示范、失败与人工纠正经验，完成 value function、advantage relabeling、条件策略训练和 CFG 推理的最小闭环。

## 复现范围

本项目以公开的 **“π*0.6: a VLA That Learns From Experience”** RECAP
论文和 π0.6 model card 为主要依据。公开资料没有提供完整的
π0.6/`π*0.6` 训练源码、权重和 RECAP 数据流水线，未公开的部分在本项目中使用小型可运行结构替代：

- 论文明确给出的 RECAP 顺序、201-bin value distribution、二值 advantage
  条件、30% 条件 dropout 和 CFG 被保留；
- 没有公开的部分采用最小且可测试的替代设计；
- 本项目只验证信息流和优化目标能否在几个样本上工作，不声称数值等价。

## 模型阶段与输入输出

| 阶段 | 输入 | 模块 | 输出 | 监督目标 |
|---|---|---|---|---|
| 1. Base policy | `64×64` 图像、离散历史 XY、任务提示 | 小型 VLM policy + stop-gradient flow expert | subtask、FAST-lite tokens、连续 `16×2` action chunk | token CE + flow matching |
| 2. Experience | 一条真实示范轨迹 | 确定性经验构造 | 成功示范、understeer、提前停止、人工纠正 4 类 episode | 无梯度，仅构造 buffer |
| 3. Value | 图像、离散状态、任务、episode return | 小型独立 value VLM | `201` 个 return bins，期望 value 位于 `[-1,0]` | distributional CE |
| 4. Advantage | reward、当前/未来 value、intervention mask | n-step estimator + quantile threshold | positive/negative indicator；人工纠正强制 positive | 无梯度，生成策略条件 |
| 5. Mixed warm-up | 示范与合成经验、advantage 条件 | advantage-conditioned policy | 条件化离散动作和连续 flow | token CE + flow matching |
| 6. RECAP extraction | 更新后的 value、重新标注的全部经验 | policy 从 base checkpoint 再训练 | 改进后的 `π*0.6`-style policy | 条件化 token CE + flow matching |
| 7. CFG inference | 同一 observation 分别使用 positive 和 omitted 条件 | 两次 flow 前向 + CFG，默认 `β=2` | guided 连续 action chunk | 与真值及 unguided 输出比较 |

最关键的信息顺序是：

```text
observation → predicted subtask → advantage indicator → actions
```

advantage token 位于 subtask 之后、动作之前，因此它只改变动作分布，不应改变
高层 subtask。flow expert 可以读取 observation、subtask 和 advantage，但不能
读取 teacher-forced FAST-lite 真值动作。flow loss 对 policy/VLM prefix 使用
stop-gradient，以最小形式模拟论文的 Knowledge Insulation（KI）训练原则。

## 与 π0.5 最小项目的重点差异

| π0.5 minimal | π0.6 / `π*0.6` minimal |
|---|---|
| 只使用示范样本 | 混合示范、autonomous failures 和 intervention corrections |
| 没有 critic | 独立的 201-bin distributional value model |
| 动作只由 observation/subtask 条件化 | 动作额外由 positive/negative/omitted advantage 条件化 |
| 一次离散预训练 + flow post-training | base → value → mixed → RECAP policy extraction |
| 普通 flow sampling | 支持 conditional/unconditional 双前向 CFG |

官方 π0.6 还把 VLM 更新为 Gemma 3 4B、action expert 扩展到约 860M，并增加
更多机器人数据。本项目按“最小结构验证”的目标使用从零训练的小模型：policy
width 为 96，value width 为 64；没有复现这些规模和数据改进。

## 最小经验与 RECAP 近似

一个 PAI 轨迹被确定性地转换为四种经验：

- `demo_success`：原始专家轨迹；
- `autonomous_understeer`：横向转向不足；
- `autonomous_stop`：中途停止；
- `intervention_correction`：先偏离，再由“人工”轨迹拉回。

默认用 4-step bootstrapped advantage，`gamma=0.99`。训练集的 advantage
第 70 百分位作为 positive 阈值，即约 30% 普通样本为 positive；intervention
动作无条件标为 positive。这个合成 buffer 只用于检验 RECAP 接口，不能替代
真实机器人 rollout、任务成功标签或在线人工干预。

## 安装

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

```bash
cd pi06-minimal-validation
python -m venv .venv
source .venv/bin/activate
python -m pip install -U pip
python -m pip install -e '.[dev]'
```

默认配置读取：

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

若本地数据路径不同，请修改 `configs/overfit.yaml` 中的
`data.source_manifest`。仓库不包含或重新分发原始 PAI 数据。

## 测试

```bash
pytest -q
```

测试应覆盖 flow matching、FAST-lite 往返、三块/防泄漏 attention、return 与
n-step advantage、intervention 强制 positive、condition dropout、CFG 公式以及
value/policy 前向形状。

## 四阶段过拟合训练

```bash
pi06mini-train --config configs/overfit.yaml
```

或：

```bash
python -m pi06mini.train --config configs/overfit.yaml
```

默认训练预算刻意很小：

| 阶段 | 步数 | 目的 |
|---|---:|---|
| base | 300 | 先记住示范的 subtask、FAST-lite 和连续动作 |
| value | 400 | 记住成功、失败和纠正 episode 的 201-bin return |
| mixed | 300 | 让策略学会 positive/negative/omitted 三种条件 |
| recap | 700 | 使用合成经验中的 n-step advantage 标签与 30% 条件 dropout 提取最终策略 |

训练应保存 checkpoint、每阶段 loss、value/advantage 诊断和最终评估。CPU 可以
运行；Apple Silicon 自动使用 MPS，CUDA 环境自动使用 GPU。

## 推理与效果图

```bash
pi06mini-infer --checkpoint outputs/overfit/checkpoint.pt
```

或：

```bash
python -m pi06mini.infer --checkpoint outputs/overfit/checkpoint.pt
```

效果图应至少同时显示：

- ground truth；
- base/unguided flow；
- positive advantage flow；
- `β=2` CFG-guided flow；
- 必要时附 FAST-lite 离散轨迹作为独立校验。

在极小数据过拟合场景中，目标是看到 value 能区分合成成功/失败，positive 或
CFG 策略比 negative/unguided 更接近专家轨迹。它不代表模型具备泛化能力。

## 关键配置

- 图像大小：`64`；action horizon：`16`；
- policy width：`96`；value width：`64`；
- value distribution：`201` bins，范围 `[-1,0]`；
- advantage condition dropout：`0.3`；
- CFG `beta`：`2.0`，位于论文建议的温和范围内；
- positive quantile：`0.7`；n-step：`4`。

## 安全限制

- 本模型只在几个训练样本及人为构造的失败轨迹上验证记忆能力。
- 严禁连接车辆、机械臂、移动机器人或任何能造成物理影响的执行器。
- 轨迹重合、value 排序正确或 loss 下降均不代表安全性、鲁棒性或任务成功率。
- 高 CFG 会把动作推向数据支持边界并产生激进行为；本项目即便使用温和
  `β=2` 也只允许离线可视化。
- 真实部署至少需要真实 rollout、独立验证集、分布外测试、动作约束、碰撞
  检查、故障降级、急停和人工监督。
- PyTorch checkpoint 使用 pickle 机制；只加载可信来源的 checkpoint。

## 预期结论

成功的实验最多能说明：这个独立实现中的 π0.5-style 双输出 policy、
distributional value function、advantage relabeling、条件 dropout 与 CFG 可以
组成一条可训练的最小 RECAP 链路，并在合成小经验集上过拟合。它不能证明与
官方 π0.6/`π*0.6` 数值等价，也不能证明强化学习改进能够迁移到真实系统。

## 公开参考资料

- [π*0.6 / RECAP 研究页面](https://pi.website/blog/pistar06)
- [π*0.6 / RECAP 论文 PDF](https://pi.website/download/pistar06.pdf)
- [π0.6 model card](https://website.pi-asset.com/pi06star/PI06_model_card.pdf)
- [Physical Intelligence/openpi](https://github.com/Physical-Intelligence/openpi)