seeker
195 lines
8.9 KiB
Markdown
195 lines
8.9 KiB
Markdown
# π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 | critic inference + n-step estimator + quantile threshold | positive/negative indicator；人工纠正强制 positive | value 更新后重新标注策略数据 |
|
||
| 5. Mixed SFT baseline | 示范与合成经验，不提供 advantage | unconditioned mixed 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）训练原则。
|
||
|
||
critic 只读取 observation/state（本单任务实验隐含固定任务提示），不读取
|
||
`demo_success`、失败类型或 intervention source 等结果标签。value 训练完成后，
|
||
程序会对每条 episode 的所有状态重新运行 critic，并用预测值重新计算 n-step
|
||
advantage；最终 RECAP policy 使用该重标注结果，而不是手工 baseline。
|
||
|
||
## 与 π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，按论文使用无折扣 return（`gamma=1.0`）。训练集的 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]'
|
||
```
|
||
|
||
完整的交互式复现入口见
|
||
[`notebooks/pi06_recap_minimal_reproduction.ipynb`](notebooks/pi06_recap_minimal_reproduction.ipynb)。
|
||
它包含论文链接、结构说明、环境安装、数据/checkpoint 本地导入或 URL 下载、
|
||
测试、训练、评估与结果图展示。启动方式：
|
||
|
||
```bash
|
||
python -m pip install -e '.[dev,notebook]'
|
||
jupyter lab notebooks/pi06_recap_minimal_reproduction.ipynb
|
||
```
|
||
|
||
默认配置读取：
|
||
|
||
```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 | 训练不带 advantage 的全经验 SFT 对照组 |
|
||
| 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 更接近专家轨迹。它不代表模型具备泛化能力。
|
||
推理链会先由 VLM 预测 subtask token，再将该预测与 advantage indicator 一起
|
||
交给连续 action expert，不使用数据集中的真值 subtask 作为低层条件。
|
||
|
||
## 关键配置
|
||
|
||
- 图像大小：`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)
|