Initial pi0.6 minimal validation
This commit is contained in:
seeker
@@ -0,0 +1,9 @@
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
.pytest_cache/
|
||||
.ruff_cache/
|
||||
*.egg-info/
|
||||
.venv/
|
||||
data/
|
||||
outputs/**/*.pt
|
||||
.DS_Store
|
||||
@@ -0,0 +1,184 @@
|
||||
# π0.6 / RECAP Minimal Validation
|
||||
|
||||
这是一个用于**小样本过拟合验证**的 π0.6 / `π*0.6` 风格项目。它把 π0.5
|
||||
已有的“高层 subtask + FAST 离散动作 + flow action expert”扩展为一个最小
|
||||
RECAP 闭环：收集示范、失败与人工纠正经验，训练分布式 value function，估计
|
||||
advantage，再用正/负 advantage 条件训练策略，并在推理时比较普通策略与
|
||||
classifier-free guidance（CFG）策略。
|
||||
|
||||
它不是 Physical Intelligence 官方实现，也不能用于真实车辆或机器人控制。
|
||||
|
||||
## 公开资料边界
|
||||
|
||||
本项目以公开的 **“π*0.6: a VLA That Learns From Experience”** RECAP
|
||||
论文和 π0.6 model card 为主要依据。π0.6 没有一篇单独公开、包含全部训练细节的
|
||||
基础模型论文；截至本项目实现时，公开 `openpi` 仓库中也没有对应的官方
|
||||
π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)
|
||||
@@ -0,0 +1,32 @@
|
||||
# Third-party notices
|
||||
|
||||
This repository is an independent educational implementation. It does not
|
||||
contain code, weights, datasets, or other assets copied from Physical
|
||||
Intelligence.
|
||||
|
||||
The design is informed by the following public research artifacts:
|
||||
|
||||
- Physical Intelligence, **“π*0.6: a VLA That Learns From Experience”**
|
||||
(RECAP paper): <https://pi.website/download/pistar06.pdf>.
|
||||
- Physical Intelligence, **π0.6 model card**:
|
||||
<https://website.pi-asset.com/pi06star/PI06_model_card.pdf>.
|
||||
- Physical Intelligence, **openpi**, used only as a public reference for the
|
||||
earlier π0/π0.5 family; the inspected repository did not provide an official
|
||||
π0.6/RECAP training implementation or checkpoint.
|
||||
- Pertsch et al., **“FAST: Efficient Action Tokenization for
|
||||
Vision-Language-Action Models.”** This project implements a small
|
||||
DCT-and-quantization `FAST-lite` teaching approximation, not the official
|
||||
tokenizer implementation.
|
||||
|
||||
Runtime dependencies retain their own licenses and copyright notices:
|
||||
|
||||
- PyTorch
|
||||
- NumPy
|
||||
- Pillow
|
||||
- Matplotlib
|
||||
- PyYAML
|
||||
- pytest (development only)
|
||||
|
||||
Users are responsible for complying with the licenses of any external data or
|
||||
checkpoints that they choose to supply. The default configuration points to a
|
||||
local manifest and does not redistribute its source images.
|
||||
@@ -0,0 +1,54 @@
|
||||
seed: 7
|
||||
device: auto
|
||||
output_dir: outputs/overfit
|
||||
|
||||
data:
|
||||
source_manifest: ../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl
|
||||
num_variants: 4
|
||||
image_size: 64
|
||||
action_horizon: 16
|
||||
|
||||
experience:
|
||||
augmentation_factor: 1
|
||||
gamma: 0.99
|
||||
n_step: 4
|
||||
positive_quantile: 0.7
|
||||
|
||||
model:
|
||||
patch_size: 16
|
||||
n_heads: 4
|
||||
n_layers: 3
|
||||
mlp_ratio: 3
|
||||
policy_d_model: 96
|
||||
value_d_model: 64
|
||||
max_text_tokens: 64
|
||||
max_action_tokens: 32
|
||||
action_start_token_id: 269
|
||||
dropout: 0.0
|
||||
value_bins: 201
|
||||
value_min: -1.0
|
||||
value_max: 0.0
|
||||
|
||||
train:
|
||||
base_steps: 300
|
||||
value_steps: 400
|
||||
mixed_steps: 300
|
||||
recap_steps: 700
|
||||
batch_size: 4
|
||||
learning_rate: 0.001
|
||||
value_learning_rate: 0.001
|
||||
weight_decay: 0.0
|
||||
grad_clip_norm: 1.0
|
||||
token_loss_weight: 1.0
|
||||
flow_loss_weight: 2.0
|
||||
value_loss_weight: 1.0
|
||||
log_every: 20
|
||||
|
||||
recap:
|
||||
condition_dropout: 0.3
|
||||
cfg_beta: 2.0
|
||||
force_interventions_positive: true
|
||||
|
||||
inference:
|
||||
flow_steps: 10
|
||||
cfg_beta: 2.0
|
||||
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"trajectory_metrics": {
|
||||
"mixed": 100.86174011230469,
|
||||
"unconditional": 21.51112937927246,
|
||||
"positive": 0.061239250004291534,
|
||||
"negative": 55.823570251464844,
|
||||
"cfg": 0.24478793144226074,
|
||||
"negative_to_failure_mse": 48.71984100341797
|
||||
},
|
||||
"values": [
|
||||
{
|
||||
"source_type": "demo_success",
|
||||
"target": 0.0,
|
||||
"prediction": -0.007052943576127291
|
||||
},
|
||||
{
|
||||
"source_type": "autonomous_understeer",
|
||||
"target": -0.03799404203891754,
|
||||
"prediction": -0.04714635759592056
|
||||
},
|
||||
{
|
||||
"source_type": "autonomous_stop",
|
||||
"target": -0.9509902000427246,
|
||||
"prediction": -0.9405871629714966
|
||||
},
|
||||
{
|
||||
"source_type": "intervention_correction",
|
||||
"target": -0.0051981499418616295,
|
||||
"prediction": -0.0118478462100029
|
||||
}
|
||||
]
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 87 KiB |
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"device": "mps",
|
||||
"experience_samples": 4,
|
||||
"policy_parameters": 817058,
|
||||
"value_parameters": 171977,
|
||||
"sources": [
|
||||
"demo_success",
|
||||
"autonomous_understeer",
|
||||
"autonomous_stop",
|
||||
"intervention_correction"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,625 @@
|
||||
[
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 1,
|
||||
"loss": 72.41638946533203,
|
||||
"text_loss": 68.23114013671875,
|
||||
"flow_loss": 2.0926246643066406
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 20,
|
||||
"loss": 4.305421829223633,
|
||||
"text_loss": 1.9023208618164062,
|
||||
"flow_loss": 1.2015506029129028
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 40,
|
||||
"loss": 2.912196159362793,
|
||||
"text_loss": 0.8887173533439636,
|
||||
"flow_loss": 1.0117393732070923
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 60,
|
||||
"loss": 2.184612512588501,
|
||||
"text_loss": 0.5364341735839844,
|
||||
"flow_loss": 0.8240891695022583
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 80,
|
||||
"loss": 1.578565239906311,
|
||||
"text_loss": 0.23946531116962433,
|
||||
"flow_loss": 0.6695499420166016
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 100,
|
||||
"loss": 1.5856558084487915,
|
||||
"text_loss": 0.04227600246667862,
|
||||
"flow_loss": 0.7716898918151855
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 120,
|
||||
"loss": 0.3423895239830017,
|
||||
"text_loss": 0.0004414900322444737,
|
||||
"flow_loss": 0.1709740161895752
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 140,
|
||||
"loss": 0.24176472425460815,
|
||||
"text_loss": 1.8706961782299913e-05,
|
||||
"flow_loss": 0.12087301164865494
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 160,
|
||||
"loss": 0.3000774383544922,
|
||||
"text_loss": 7.637506496394053e-05,
|
||||
"flow_loss": 0.15000052750110626
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 180,
|
||||
"loss": 0.6410971879959106,
|
||||
"text_loss": 1.6521895304322243e-05,
|
||||
"flow_loss": 0.32054033875465393
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 200,
|
||||
"loss": 0.6132829189300537,
|
||||
"text_loss": 9.676185072748922e-06,
|
||||
"flow_loss": 0.30663663148880005
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 220,
|
||||
"loss": 0.3431224524974823,
|
||||
"text_loss": 7.356796231761109e-06,
|
||||
"flow_loss": 0.17155754566192627
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 240,
|
||||
"loss": 0.28047603368759155,
|
||||
"text_loss": 5.987634267512476e-06,
|
||||
"flow_loss": 0.1402350217103958
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 260,
|
||||
"loss": 0.13324393332004547,
|
||||
"text_loss": 4.9011541705112904e-06,
|
||||
"flow_loss": 0.06661951541900635
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 280,
|
||||
"loss": 0.1635851263999939,
|
||||
"text_loss": 4.0462705328536686e-06,
|
||||
"flow_loss": 0.08179053664207458
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 300,
|
||||
"loss": 0.15081387758255005,
|
||||
"text_loss": 3.4570459774840856e-06,
|
||||
"flow_loss": 0.07540521025657654
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 1,
|
||||
"loss": 5.427414894104004,
|
||||
"prediction": -0.5049723982810974,
|
||||
"target": 0.0
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 20,
|
||||
"loss": 2.24937105178833,
|
||||
"prediction": -0.3925527334213257,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 40,
|
||||
"loss": 1.728512167930603,
|
||||
"prediction": -0.32079368829727173,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 60,
|
||||
"loss": 1.4823004007339478,
|
||||
"prediction": -0.3057250678539276,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 80,
|
||||
"loss": 1.1117331981658936,
|
||||
"prediction": -0.3553512394428253,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 100,
|
||||
"loss": 0.2580283582210541,
|
||||
"prediction": -0.11316870898008347,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 120,
|
||||
"loss": 0.12432077527046204,
|
||||
"prediction": -0.054820768535137177,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 140,
|
||||
"loss": 0.08916354924440384,
|
||||
"prediction": -0.04263461381196976,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 160,
|
||||
"loss": 0.07165656983852386,
|
||||
"prediction": -0.034443676471710205,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 180,
|
||||
"loss": 0.05974685773253441,
|
||||
"prediction": -0.029927225783467293,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 200,
|
||||
"loss": 0.05075548216700554,
|
||||
"prediction": -0.026268303394317627,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 220,
|
||||
"loss": 0.043695785105228424,
|
||||
"prediction": -0.023351771757006645,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 240,
|
||||
"loss": 0.03802856057882309,
|
||||
"prediction": -0.021020548418164253,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 260,
|
||||
"loss": 0.03340449184179306,
|
||||
"prediction": -0.019105376675724983,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 280,
|
||||
"loss": 0.02958041802048683,
|
||||
"prediction": -0.017514187842607498,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 300,
|
||||
"loss": 0.02638067863881588,
|
||||
"prediction": -0.016178548336029053,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 320,
|
||||
"loss": 0.023675711825489998,
|
||||
"prediction": -0.01504577323794365,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 340,
|
||||
"loss": 0.021368399262428284,
|
||||
"prediction": -0.014076616615056992,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 360,
|
||||
"loss": 0.019383536651730537,
|
||||
"prediction": -0.013240991160273552,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 380,
|
||||
"loss": 0.017663555219769478,
|
||||
"prediction": -0.01251534465700388,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 400,
|
||||
"loss": 0.016163267195224762,
|
||||
"prediction": -0.011881118640303612,
|
||||
"target": -0.0051981499418616295
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 1,
|
||||
"loss": 0.11788154393434525,
|
||||
"text_loss": 3.4366105410299497e-06,
|
||||
"flow_loss": 0.05893905460834503
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 20,
|
||||
"loss": 4.7238688468933105,
|
||||
"text_loss": 2.6195578575134277,
|
||||
"flow_loss": 1.0521554946899414
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 40,
|
||||
"loss": 5.051268577575684,
|
||||
"text_loss": 1.2949998378753662,
|
||||
"flow_loss": 1.8781343698501587
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 60,
|
||||
"loss": 1.6957576274871826,
|
||||
"text_loss": 0.8319522142410278,
|
||||
"flow_loss": 0.4319027364253998
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 80,
|
||||
"loss": 2.1900432109832764,
|
||||
"text_loss": 0.8061568140983582,
|
||||
"flow_loss": 0.6919432282447815
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 100,
|
||||
"loss": 2.2935876846313477,
|
||||
"text_loss": 0.5968648791313171,
|
||||
"flow_loss": 0.8483614325523376
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 120,
|
||||
"loss": 1.4990849494934082,
|
||||
"text_loss": 0.6506502032279968,
|
||||
"flow_loss": 0.4242174029350281
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 140,
|
||||
"loss": 1.6056029796600342,
|
||||
"text_loss": 0.5977785587310791,
|
||||
"flow_loss": 0.5039122104644775
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 160,
|
||||
"loss": 1.3810169696807861,
|
||||
"text_loss": 0.48245665431022644,
|
||||
"flow_loss": 0.44928014278411865
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 180,
|
||||
"loss": 1.1724852323532104,
|
||||
"text_loss": 0.3958587944507599,
|
||||
"flow_loss": 0.3883132338523865
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 200,
|
||||
"loss": 1.2829020023345947,
|
||||
"text_loss": 0.18238645792007446,
|
||||
"flow_loss": 0.5502578020095825
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 220,
|
||||
"loss": 0.7328482866287231,
|
||||
"text_loss": 0.1418272703886032,
|
||||
"flow_loss": 0.29551050066947937
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 240,
|
||||
"loss": 1.196341872215271,
|
||||
"text_loss": 0.28424081206321716,
|
||||
"flow_loss": 0.4560505449771881
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 260,
|
||||
"loss": 4.298240661621094,
|
||||
"text_loss": 0.257877379655838,
|
||||
"flow_loss": 2.020181655883789
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 280,
|
||||
"loss": 1.3184027671813965,
|
||||
"text_loss": 0.07380203157663345,
|
||||
"flow_loss": 0.622300386428833
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 300,
|
||||
"loss": 2.4886393547058105,
|
||||
"text_loss": 0.30716001987457275,
|
||||
"flow_loss": 1.0907396078109741
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 1,
|
||||
"loss": 10.11633586883545,
|
||||
"text_loss": 9.866013526916504,
|
||||
"flow_loss": 0.12516121566295624
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 20,
|
||||
"loss": 3.1665196418762207,
|
||||
"text_loss": 2.23919415473938,
|
||||
"flow_loss": 0.4636627733707428
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 40,
|
||||
"loss": 3.1297008991241455,
|
||||
"text_loss": 1.4351346492767334,
|
||||
"flow_loss": 0.847283124923706
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 60,
|
||||
"loss": 2.608584403991699,
|
||||
"text_loss": 1.171854853630066,
|
||||
"flow_loss": 0.7183647751808167
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 80,
|
||||
"loss": 3.5118634700775146,
|
||||
"text_loss": 0.9280703067779541,
|
||||
"flow_loss": 1.2918965816497803
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 100,
|
||||
"loss": 2.7196712493896484,
|
||||
"text_loss": 0.8991937041282654,
|
||||
"flow_loss": 0.9102387428283691
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 120,
|
||||
"loss": 1.1335762739181519,
|
||||
"text_loss": 0.7848832011222839,
|
||||
"flow_loss": 0.17434655129909515
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 140,
|
||||
"loss": 2.5729737281799316,
|
||||
"text_loss": 0.5885628461837769,
|
||||
"flow_loss": 0.9922055006027222
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 160,
|
||||
"loss": 1.9800522327423096,
|
||||
"text_loss": 0.6485287547111511,
|
||||
"flow_loss": 0.6657617688179016
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 180,
|
||||
"loss": 1.5937684774398804,
|
||||
"text_loss": 0.6764931678771973,
|
||||
"flow_loss": 0.45863765478134155
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 200,
|
||||
"loss": 1.588905692100525,
|
||||
"text_loss": 1.1673603057861328,
|
||||
"flow_loss": 0.21077267825603485
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 220,
|
||||
"loss": 1.1653995513916016,
|
||||
"text_loss": 0.5840165019035339,
|
||||
"flow_loss": 0.2906915545463562
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 240,
|
||||
"loss": 2.6636130809783936,
|
||||
"text_loss": 1.3158458471298218,
|
||||
"flow_loss": 0.6738836169242859
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 260,
|
||||
"loss": 0.5980206727981567,
|
||||
"text_loss": 0.42834123969078064,
|
||||
"flow_loss": 0.08483970910310745
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 280,
|
||||
"loss": 1.0666983127593994,
|
||||
"text_loss": 0.47541406750679016,
|
||||
"flow_loss": 0.2956421375274658
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 300,
|
||||
"loss": 0.8535268306732178,
|
||||
"text_loss": 0.4074874520301819,
|
||||
"flow_loss": 0.22301967442035675
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 320,
|
||||
"loss": 5.488016605377197,
|
||||
"text_loss": 0.851451575756073,
|
||||
"flow_loss": 2.3182826042175293
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 340,
|
||||
"loss": 0.7862600088119507,
|
||||
"text_loss": 0.37299302220344543,
|
||||
"flow_loss": 0.20663349330425262
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 360,
|
||||
"loss": 1.3418667316436768,
|
||||
"text_loss": 0.7838777899742126,
|
||||
"flow_loss": 0.27899447083473206
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 380,
|
||||
"loss": 0.7067349553108215,
|
||||
"text_loss": 0.15938280522823334,
|
||||
"flow_loss": 0.2736760675907135
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 400,
|
||||
"loss": 0.4985463619232178,
|
||||
"text_loss": 0.2566893398761749,
|
||||
"flow_loss": 0.12092851102352142
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 420,
|
||||
"loss": 0.5810034871101379,
|
||||
"text_loss": 0.24579037725925446,
|
||||
"flow_loss": 0.16760656237602234
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 440,
|
||||
"loss": 0.40272825956344604,
|
||||
"text_loss": 0.19410930573940277,
|
||||
"flow_loss": 0.10430946946144104
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 460,
|
||||
"loss": 0.7170483469963074,
|
||||
"text_loss": 0.32609468698501587,
|
||||
"flow_loss": 0.19547683000564575
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 480,
|
||||
"loss": 0.3233756422996521,
|
||||
"text_loss": 0.08512496948242188,
|
||||
"flow_loss": 0.11912534385919571
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 500,
|
||||
"loss": 0.625529408454895,
|
||||
"text_loss": 0.23650097846984863,
|
||||
"flow_loss": 0.1945142149925232
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 520,
|
||||
"loss": 1.645598292350769,
|
||||
"text_loss": 0.874957263469696,
|
||||
"flow_loss": 0.3853205144405365
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 540,
|
||||
"loss": 0.8086838722229004,
|
||||
"text_loss": 0.31192541122436523,
|
||||
"flow_loss": 0.24837924540042877
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 560,
|
||||
"loss": 1.9226100444793701,
|
||||
"text_loss": 0.8500325083732605,
|
||||
"flow_loss": 0.5362887382507324
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 580,
|
||||
"loss": 2.138998031616211,
|
||||
"text_loss": 0.7942134141921997,
|
||||
"flow_loss": 0.6723923683166504
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 600,
|
||||
"loss": 0.504305362701416,
|
||||
"text_loss": 0.13201656937599182,
|
||||
"flow_loss": 0.1861443817615509
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 620,
|
||||
"loss": 0.5794067978858948,
|
||||
"text_loss": 0.26564836502075195,
|
||||
"flow_loss": 0.1568792164325714
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 640,
|
||||
"loss": 0.8541183471679688,
|
||||
"text_loss": 0.154631569981575,
|
||||
"flow_loss": 0.34974339604377747
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 660,
|
||||
"loss": 0.16945119202136993,
|
||||
"text_loss": 0.056921038776636124,
|
||||
"flow_loss": 0.056265074759721756
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 680,
|
||||
"loss": 0.47458726167678833,
|
||||
"text_loss": 0.04364287480711937,
|
||||
"flow_loss": 0.21547219157218933
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 700,
|
||||
"loss": 0.18534913659095764,
|
||||
"text_loss": 0.05316398665308952,
|
||||
"flow_loss": 0.06609257310628891
|
||||
}
|
||||
]
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 46 KiB |
@@ -0,0 +1,38 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=69", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "pi06-minimal-validation"
|
||||
version = "0.1.0"
|
||||
description = "A tiny educational pi0.6 and RECAP overfitting experiment"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
license = { text = "MIT" }
|
||||
authors = [{ name = "pi06 minimal validation contributors" }]
|
||||
dependencies = [
|
||||
"matplotlib>=3.8",
|
||||
"numpy>=1.26",
|
||||
"Pillow>=10.0",
|
||||
"PyYAML>=6.0",
|
||||
"torch>=2.2",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=8.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
pi06mini-train = "pi06mini.train:main"
|
||||
pi06mini-infer = "pi06mini.infer:main"
|
||||
|
||||
[tool.setuptools]
|
||||
package-dir = { "" = "src" }
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
addopts = "-ra"
|
||||
@@ -0,0 +1,32 @@
|
||||
"""Minimal pi0.6 / Recap architecture validation package."""
|
||||
|
||||
from .model import Pi06Config, Pi06Mini, Pi06Output, flow_matching_sample
|
||||
from .tokenizer import (
|
||||
Advantage,
|
||||
CoordinateTokenizer,
|
||||
MiniTokenizer,
|
||||
PolicySequence,
|
||||
SubtaskVocabulary,
|
||||
)
|
||||
from .value import (
|
||||
DistributionalValueModel,
|
||||
ValueConfig,
|
||||
ValueModel,
|
||||
ValueOutput,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Advantage",
|
||||
"CoordinateTokenizer",
|
||||
"DistributionalValueModel",
|
||||
"MiniTokenizer",
|
||||
"Pi06Config",
|
||||
"Pi06Mini",
|
||||
"Pi06Output",
|
||||
"PolicySequence",
|
||||
"SubtaskVocabulary",
|
||||
"ValueConfig",
|
||||
"ValueModel",
|
||||
"ValueOutput",
|
||||
"flow_matching_sample",
|
||||
]
|
||||
@@ -0,0 +1,354 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Sequence
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image, ImageEnhance
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from .experience import (
|
||||
ExperienceEpisode,
|
||||
build_synthetic_experience,
|
||||
discounted_returns,
|
||||
n_step_advantages,
|
||||
normalize_returns_to_unit_interval,
|
||||
quantile_advantage_labels,
|
||||
stack_episode_field,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_SOURCE_MANIFEST = (
|
||||
Path(__file__).resolve().parents[4]
|
||||
/ "alpamayo"
|
||||
/ "alpamayo-qwen05b-validation"
|
||||
/ "data"
|
||||
/ "processed"
|
||||
/ "overfit_manifest.jsonl"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SampleRecord:
|
||||
clip_id: str
|
||||
t0_us: int
|
||||
image_path: Path
|
||||
trajectory_path: Path
|
||||
reasoning: str = ""
|
||||
|
||||
@property
|
||||
def sample_id(self) -> str:
|
||||
return f"{self.clip_id}_{self.t0_us}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ActionStatistics:
|
||||
mean: torch.Tensor
|
||||
scale: torch.Tensor
|
||||
|
||||
def as_dict(self) -> dict[str, list[float]]:
|
||||
return {
|
||||
"mean": [float(value) for value in self.mean],
|
||||
"scale": [float(value) for value in self.scale],
|
||||
}
|
||||
|
||||
|
||||
class ActionNormalizer:
|
||||
"""Invertibly normalize XY waypoint deltas for flow matching."""
|
||||
|
||||
def __init__(self, mean: torch.Tensor, scale: torch.Tensor) -> None:
|
||||
self.mean = torch.as_tensor(mean, dtype=torch.float32).reshape(2)
|
||||
self.scale = torch.as_tensor(scale, dtype=torch.float32).reshape(2).clamp_min(0.05)
|
||||
|
||||
@staticmethod
|
||||
def to_deltas(history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor:
|
||||
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||
future = torch.as_tensor(future_xy, dtype=torch.float32)
|
||||
if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 1:
|
||||
raise ValueError("history_xy must have shape [history_steps, 2]")
|
||||
if future.ndim != 2 or future.shape[-1] != 2 or len(future) < 1:
|
||||
raise ValueError("future_xy must have shape [action_horizon, 2]")
|
||||
previous = torch.cat((history[-1:], future[:-1]), dim=0)
|
||||
return future - previous
|
||||
|
||||
@classmethod
|
||||
def fit(
|
||||
cls, histories: Sequence[torch.Tensor], futures: Sequence[torch.Tensor]
|
||||
) -> ActionNormalizer:
|
||||
if not histories or len(histories) != len(futures):
|
||||
raise ValueError("histories and futures must be non-empty and equally sized")
|
||||
flat = torch.cat(
|
||||
[
|
||||
cls.to_deltas(history, future).reshape(-1, 2)
|
||||
for history, future in zip(histories, futures, strict=True)
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
return cls(flat.mean(dim=0), flat.std(dim=0, unbiased=False))
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, values: dict[str, Sequence[float]]) -> ActionNormalizer:
|
||||
return cls(torch.tensor(values["mean"]), torch.tensor(values["scale"]))
|
||||
|
||||
@property
|
||||
def statistics(self) -> ActionStatistics:
|
||||
return ActionStatistics(self.mean.clone(), self.scale.clone())
|
||||
|
||||
def encode(self, history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor:
|
||||
return (self.to_deltas(history_xy, future_xy) - self.mean) / self.scale
|
||||
|
||||
def decode(self, actions: torch.Tensor, history_xy: torch.Tensor) -> torch.Tensor:
|
||||
actions = torch.as_tensor(actions, dtype=torch.float32)
|
||||
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||
deltas = actions * self.scale.to(actions.device) + self.mean.to(actions.device)
|
||||
return history[-1:].to(actions.device) + torch.cumsum(deltas, dim=0)
|
||||
|
||||
def as_dict(self) -> dict[str, list[float]]:
|
||||
return self.statistics.as_dict()
|
||||
|
||||
def save(self, path: str | Path) -> Path:
|
||||
destination = Path(path).expanduser().resolve()
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
destination.write_text(json.dumps(self.as_dict(), indent=2) + "\n", encoding="utf-8")
|
||||
return destination
|
||||
|
||||
|
||||
def _resolve_data_path(manifest_path: Path, raw_path: str | Path) -> Path:
|
||||
path = Path(raw_path).expanduser()
|
||||
if path.is_absolute():
|
||||
return path.resolve()
|
||||
candidates = [manifest_path.parent / path]
|
||||
if len(manifest_path.parents) >= 3:
|
||||
candidates.insert(0, manifest_path.parents[2] / path)
|
||||
for candidate in candidates:
|
||||
if candidate.exists():
|
||||
return candidate.resolve()
|
||||
return candidates[0].resolve()
|
||||
|
||||
|
||||
def load_manifest(path: str | Path) -> list[SampleRecord]:
|
||||
manifest_path = Path(path).expanduser().resolve()
|
||||
records: list[SampleRecord] = []
|
||||
for line in manifest_path.read_text(encoding="utf-8").splitlines():
|
||||
if not line.strip():
|
||||
continue
|
||||
row = json.loads(line)
|
||||
records.append(
|
||||
SampleRecord(
|
||||
clip_id=str(row["clip_id"]),
|
||||
t0_us=int(row["t0_us"]),
|
||||
image_path=_resolve_data_path(manifest_path, row["image_path"]),
|
||||
trajectory_path=_resolve_data_path(manifest_path, row["trajectory_path"]),
|
||||
reasoning=str(row.get("reasoning", "")),
|
||||
)
|
||||
)
|
||||
if not records:
|
||||
raise ValueError(f"no samples found in {manifest_path}")
|
||||
return records
|
||||
|
||||
|
||||
def materialize_subset(
|
||||
source_manifest: str | Path = DEFAULT_SOURCE_MANIFEST,
|
||||
output_dir: str | Path = "data/overfit",
|
||||
*,
|
||||
max_samples: int = 1,
|
||||
) -> Path:
|
||||
"""Copy a self-contained PAI seed demonstration for synthetic experience."""
|
||||
|
||||
if max_samples < 1:
|
||||
raise ValueError("max_samples must be positive")
|
||||
destination = Path(output_dir).expanduser().resolve()
|
||||
image_dir = destination / "images"
|
||||
trajectory_dir = destination / "trajectories"
|
||||
image_dir.mkdir(parents=True, exist_ok=True)
|
||||
trajectory_dir.mkdir(parents=True, exist_ok=True)
|
||||
rows: list[dict[str, Any]] = []
|
||||
for index, record in enumerate(load_manifest(source_manifest)[:max_samples]):
|
||||
stem = f"sample_{index:02d}_{record.sample_id}"
|
||||
image_path = image_dir / f"{stem}{record.image_path.suffix.lower()}"
|
||||
trajectory_path = trajectory_dir / f"{stem}.npz"
|
||||
shutil.copy2(record.image_path, image_path)
|
||||
shutil.copy2(record.trajectory_path, trajectory_path)
|
||||
rows.append(
|
||||
{
|
||||
"clip_id": record.clip_id,
|
||||
"t0_us": record.t0_us,
|
||||
"reasoning": record.reasoning,
|
||||
"image_path": str(image_path.relative_to(destination)),
|
||||
"trajectory_path": str(trajectory_path.relative_to(destination)),
|
||||
}
|
||||
)
|
||||
manifest = destination / "overfit_manifest.jsonl"
|
||||
manifest.write_text(
|
||||
"".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return manifest
|
||||
|
||||
|
||||
def _resample_xy(xy: torch.Tensor, steps: int) -> torch.Tensor:
|
||||
xy = torch.as_tensor(xy, dtype=torch.float32)
|
||||
if xy.ndim != 2 or xy.shape[-1] != 2 or len(xy) < 2:
|
||||
raise ValueError("trajectory must contain at least two XY waypoints")
|
||||
if len(xy) == steps:
|
||||
return xy
|
||||
return F.interpolate(
|
||||
xy.T.unsqueeze(0), size=steps, mode="linear", align_corners=True
|
||||
)[0].T
|
||||
|
||||
|
||||
class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
|
||||
"""Four-source ReCAP overfit buffer synthesized from each PAI demonstration."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path: str | Path,
|
||||
*,
|
||||
image_size: int = 64,
|
||||
augmentation_factor: int = 1,
|
||||
history_steps: int = 8,
|
||||
action_horizon: int = 16,
|
||||
n_step: int = 4,
|
||||
gamma: float = 0.99,
|
||||
positive_quantile: float = 0.7,
|
||||
seed: int = 7,
|
||||
action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None,
|
||||
) -> None:
|
||||
if image_size < 1 or augmentation_factor < 1:
|
||||
raise ValueError("image_size and augmentation_factor must be positive")
|
||||
self.records = load_manifest(manifest_path)
|
||||
self.image_size = int(image_size)
|
||||
self.augmentation_factor = int(augmentation_factor)
|
||||
self.history_steps = int(history_steps)
|
||||
self.action_horizon = int(action_horizon)
|
||||
self.seed = int(seed)
|
||||
self.state_dim = self.history_steps * 2
|
||||
self.action_dim = 2
|
||||
|
||||
episodes: list[ExperienceEpisode] = []
|
||||
episode_record_indices: list[int] = []
|
||||
for record_index, record in enumerate(self.records):
|
||||
with np.load(record.trajectory_path) as arrays:
|
||||
history = _resample_xy(
|
||||
torch.from_numpy(np.asarray(arrays["history_xyz"], dtype=np.float32))[:, :2],
|
||||
self.history_steps,
|
||||
)
|
||||
future = _resample_xy(
|
||||
torch.from_numpy(np.asarray(arrays["future_xyz"], dtype=np.float32))[:, :2],
|
||||
self.action_horizon,
|
||||
)
|
||||
generated = build_synthetic_experience(
|
||||
history, future, episode_prefix=record.sample_id
|
||||
)
|
||||
episodes.extend(generated)
|
||||
episode_record_indices.extend([record_index] * len(generated))
|
||||
self.episodes = episodes
|
||||
self._episode_record_indices = episode_record_indices
|
||||
|
||||
raw_rewards = stack_episode_field(episodes, "rewards")
|
||||
raw_returns = discounted_returns(raw_rewards, gamma=gamma)
|
||||
self.normalized_returns = normalize_returns_to_unit_interval(raw_returns)
|
||||
baseline = raw_returns.mean(dim=0, keepdim=True).expand_as(raw_returns)
|
||||
dones = torch.zeros_like(raw_rewards, dtype=torch.bool)
|
||||
dones[:, -1] = True
|
||||
self.advantages = n_step_advantages(
|
||||
raw_rewards, baseline, n_step=n_step, gamma=gamma, dones=dones
|
||||
)
|
||||
intervention_masks = stack_episode_field(episodes, "intervention_mask").bool()
|
||||
self.positive_steps, self.advantage_threshold = quantile_advantage_labels(
|
||||
self.advantages,
|
||||
quantile=positive_quantile,
|
||||
intervention_mask=intervention_masks,
|
||||
)
|
||||
episode_advantages = self.advantages.mean(dim=-1)
|
||||
episode_interventions = intervention_masks.any(dim=-1)
|
||||
self.positive_episodes, self.episode_advantage_threshold = (
|
||||
quantile_advantage_labels(
|
||||
episode_advantages,
|
||||
quantile=positive_quantile,
|
||||
intervention_mask=episode_interventions,
|
||||
)
|
||||
)
|
||||
|
||||
if action_statistics is None:
|
||||
self.action_normalizer = ActionNormalizer.fit(
|
||||
[episode.history_xy for episode in episodes],
|
||||
[episode.future_xy for episode in episodes],
|
||||
)
|
||||
elif isinstance(action_statistics, ActionStatistics):
|
||||
self.action_normalizer = ActionNormalizer(
|
||||
action_statistics.mean, action_statistics.scale
|
||||
)
|
||||
else:
|
||||
self.action_normalizer = ActionNormalizer.from_dict(action_statistics)
|
||||
|
||||
@property
|
||||
def action_statistics(self) -> ActionStatistics:
|
||||
return self.action_normalizer.statistics
|
||||
|
||||
def save_action_statistics(self, path: str | Path) -> Path:
|
||||
return self.action_normalizer.save(path)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.episodes) * self.augmentation_factor
|
||||
|
||||
def _indices(self, index: int) -> tuple[int, int]:
|
||||
if index < 0:
|
||||
index += len(self)
|
||||
if index < 0 or index >= len(self):
|
||||
raise IndexError(index)
|
||||
return divmod(index, self.augmentation_factor)
|
||||
|
||||
def _load_image(self, record_index: int, variant: int) -> torch.Tensor:
|
||||
with Image.open(self.records[record_index].image_path) as source:
|
||||
image = source.convert("RGB")
|
||||
if variant:
|
||||
rng = np.random.default_rng(
|
||||
np.random.SeedSequence([self.seed, record_index, variant])
|
||||
)
|
||||
image = ImageEnhance.Brightness(image).enhance(float(rng.uniform(0.96, 1.04)))
|
||||
image = ImageEnhance.Contrast(image).enhance(float(rng.uniform(0.97, 1.03)))
|
||||
image = image.resize(
|
||||
(self.image_size, self.image_size), resample=Image.Resampling.BILINEAR
|
||||
)
|
||||
array = np.asarray(image, dtype=np.float32).copy() / 255.0
|
||||
return torch.from_numpy(array).permute(2, 0, 1).contiguous()
|
||||
|
||||
def __getitem__(self, index: int) -> dict[str, Any]:
|
||||
episode_index, variant = self._indices(index)
|
||||
episode = self.episodes[episode_index]
|
||||
record_index = self._episode_record_indices[episode_index]
|
||||
actions = self.action_normalizer.encode(
|
||||
episode.history_xy, episode.future_xy
|
||||
)
|
||||
return {
|
||||
"image": self._load_image(record_index, variant),
|
||||
"state": episode.history_xy.reshape(-1),
|
||||
"actions": actions,
|
||||
"history_xy": episode.history_xy.clone(),
|
||||
"future_xy": episode.future_xy.clone(),
|
||||
"rewards": episode.rewards.clone(),
|
||||
"returns": self.normalized_returns[episode_index].clone(),
|
||||
# The critic predicts the beginning-of-episode return for this
|
||||
# minimal trajectory-level buffer; the full vector remains
|
||||
# available for diagnostics and advantage recomputation.
|
||||
"return_target": self.normalized_returns[episode_index, 0].clone(),
|
||||
"value_target": self.normalized_returns[episode_index, 0].clone(),
|
||||
"advantages": self.advantages[episode_index].clone(),
|
||||
"positive_steps": self.positive_steps[episode_index].clone(),
|
||||
"advantage": self.advantages[episode_index].mean().clone(),
|
||||
"positive_condition": self.positive_episodes[episode_index].clone(),
|
||||
"advantage_label": (
|
||||
"positive" if bool(self.positive_episodes[episode_index]) else "negative"
|
||||
),
|
||||
"intervention_mask": episode.intervention_mask.clone(),
|
||||
"is_intervention": bool(episode.intervention_mask.any()),
|
||||
"source_type": episode.source_type,
|
||||
"reasoning": self.records[record_index].reasoning,
|
||||
"sample_id": f"{episode.episode_id}:aug{variant}",
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Literal, Sequence
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
ExperienceSource = Literal[
|
||||
"demo_success",
|
||||
"autonomous_understeer",
|
||||
"autonomous_stop",
|
||||
"intervention_correction",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExperienceEpisode:
|
||||
"""A tiny trajectory-level stand-in for the heterogeneous pi0.6 experience buffer."""
|
||||
|
||||
episode_id: str
|
||||
source_type: ExperienceSource
|
||||
history_xy: torch.Tensor
|
||||
future_xy: torch.Tensor
|
||||
rewards: torch.Tensor
|
||||
intervention_mask: torch.Tensor
|
||||
metadata: dict[str, float | str] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
history = torch.as_tensor(self.history_xy, dtype=torch.float32)
|
||||
future = torch.as_tensor(self.future_xy, dtype=torch.float32)
|
||||
rewards = torch.as_tensor(self.rewards, dtype=torch.float32)
|
||||
intervention = torch.as_tensor(self.intervention_mask, dtype=torch.bool)
|
||||
if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 1:
|
||||
raise ValueError("history_xy must have shape [history_steps, 2]")
|
||||
if future.ndim != 2 or future.shape[-1] != 2 or len(future) < 2:
|
||||
raise ValueError("future_xy must have shape [horizon, 2]")
|
||||
if rewards.shape != (len(future),):
|
||||
raise ValueError("rewards must have one value per future waypoint")
|
||||
if intervention.shape != rewards.shape:
|
||||
raise ValueError("intervention_mask must match rewards")
|
||||
object.__setattr__(self, "history_xy", history)
|
||||
object.__setattr__(self, "future_xy", future)
|
||||
object.__setattr__(self, "rewards", rewards)
|
||||
object.__setattr__(self, "intervention_mask", intervention)
|
||||
|
||||
|
||||
def discounted_returns(rewards: torch.Tensor, gamma: float = 0.99) -> torch.Tensor:
|
||||
"""Compute reward-to-go while preserving arbitrary leading batch dimensions."""
|
||||
|
||||
if not 0.0 <= gamma <= 1.0:
|
||||
raise ValueError("gamma must lie in [0, 1]")
|
||||
rewards = torch.as_tensor(rewards, dtype=torch.float32)
|
||||
if rewards.ndim < 1 or rewards.shape[-1] < 1:
|
||||
raise ValueError("rewards must have a non-empty time dimension")
|
||||
returns = torch.empty_like(rewards)
|
||||
running = torch.zeros_like(rewards[..., 0])
|
||||
for index in range(rewards.shape[-1] - 1, -1, -1):
|
||||
running = rewards[..., index] + gamma * running
|
||||
returns[..., index] = running
|
||||
return returns
|
||||
|
||||
|
||||
def normalize_returns_to_unit_interval(returns: torch.Tensor) -> torch.Tensor:
|
||||
"""Map the worst observed return to -1 and the best to 0.
|
||||
|
||||
pi0.6 uses non-positive reward conditioning. A constant collection is
|
||||
treated as equally optimal and therefore maps to zero.
|
||||
"""
|
||||
|
||||
returns = torch.as_tensor(returns, dtype=torch.float32)
|
||||
if returns.numel() == 0:
|
||||
raise ValueError("returns must be non-empty")
|
||||
minimum = returns.amin()
|
||||
maximum = returns.amax()
|
||||
span = maximum - minimum
|
||||
if float(span) <= torch.finfo(returns.dtype).eps:
|
||||
return torch.zeros_like(returns)
|
||||
return ((returns - maximum) / span).clamp(-1.0, 0.0)
|
||||
|
||||
|
||||
def n_step_advantages(
|
||||
rewards: torch.Tensor,
|
||||
values: torch.Tensor,
|
||||
*,
|
||||
n_step: int = 4,
|
||||
gamma: float = 0.99,
|
||||
dones: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Estimate A(s_t,a_t) from an n-step bootstrapped target.
|
||||
|
||||
``values`` may contain T values (the final bootstrap is zero) or T+1
|
||||
values. ``dones`` stops both reward accumulation and bootstrapping.
|
||||
"""
|
||||
|
||||
rewards = torch.as_tensor(rewards, dtype=torch.float32)
|
||||
values = torch.as_tensor(values, dtype=torch.float32, device=rewards.device)
|
||||
if rewards.ndim < 1 or rewards.shape[-1] < 1:
|
||||
raise ValueError("rewards must have a non-empty time dimension")
|
||||
if n_step < 1:
|
||||
raise ValueError("n_step must be positive")
|
||||
if not 0.0 <= gamma <= 1.0:
|
||||
raise ValueError("gamma must lie in [0, 1]")
|
||||
horizon = rewards.shape[-1]
|
||||
if values.shape[:-1] != rewards.shape[:-1] or values.shape[-1] not in {
|
||||
horizon,
|
||||
horizon + 1,
|
||||
}:
|
||||
raise ValueError("values must have shape [..., T] or [..., T+1]")
|
||||
if values.shape[-1] == horizon:
|
||||
values = torch.cat((values, torch.zeros_like(values[..., :1])), dim=-1)
|
||||
if dones is None:
|
||||
dones = torch.zeros_like(rewards, dtype=torch.bool)
|
||||
dones[..., -1] = True
|
||||
else:
|
||||
dones = torch.as_tensor(dones, dtype=torch.bool, device=rewards.device)
|
||||
if dones.shape != rewards.shape:
|
||||
raise ValueError("dones must match rewards")
|
||||
|
||||
advantages = torch.empty_like(rewards)
|
||||
for start in range(horizon):
|
||||
target = torch.zeros_like(rewards[..., 0])
|
||||
alive = torch.ones_like(target)
|
||||
discount = 1.0
|
||||
stop = min(start + n_step, horizon)
|
||||
for index in range(start, stop):
|
||||
target = target + discount * alive * rewards[..., index]
|
||||
alive = alive * (~dones[..., index]).to(rewards.dtype)
|
||||
discount *= gamma
|
||||
# These inputs are complete episodes rather than replay fragments, so
|
||||
# reaching the final recorded step is terminal even if callers omit a
|
||||
# final done flag. Bootstrap only when the n-step window ends earlier.
|
||||
if stop < horizon:
|
||||
target = target + discount * alive * values[..., stop]
|
||||
advantages[..., start] = target - values[..., start]
|
||||
return advantages
|
||||
|
||||
|
||||
def quantile_advantage_labels(
|
||||
advantages: torch.Tensor,
|
||||
*,
|
||||
quantile: float = 0.7,
|
||||
intervention_mask: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Convert advantages into ReCAP positive/negative conditions.
|
||||
|
||||
Intervention actions are always positive, even when their estimated
|
||||
numerical advantage falls below the empirical threshold.
|
||||
"""
|
||||
|
||||
advantages = torch.as_tensor(advantages, dtype=torch.float32)
|
||||
if advantages.numel() == 0:
|
||||
raise ValueError("advantages must be non-empty")
|
||||
if not 0.0 <= quantile <= 1.0:
|
||||
raise ValueError("quantile must lie in [0, 1]")
|
||||
threshold = torch.quantile(advantages.reshape(-1), quantile)
|
||||
labels = advantages >= threshold
|
||||
if intervention_mask is not None:
|
||||
intervention_mask = torch.as_tensor(
|
||||
intervention_mask, dtype=torch.bool, device=advantages.device
|
||||
)
|
||||
if intervention_mask.shape != advantages.shape:
|
||||
raise ValueError("intervention_mask must match advantages")
|
||||
labels = labels | intervention_mask
|
||||
return labels, threshold
|
||||
|
||||
|
||||
def _trajectory_rewards(candidate: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
error = torch.linalg.vector_norm(candidate - target, dim=-1)
|
||||
scale = torch.linalg.vector_norm(target - target[:1], dim=-1).amax().clamp_min(1.0)
|
||||
return -(error / scale).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def build_synthetic_experience(
|
||||
history_xy: torch.Tensor,
|
||||
expert_future_xy: torch.Tensor,
|
||||
*,
|
||||
episode_prefix: str = "sample",
|
||||
) -> list[ExperienceEpisode]:
|
||||
"""Build four deterministic experience sources from one PAI demonstration.
|
||||
|
||||
The small buffer deliberately includes successful demonstrations, two
|
||||
autonomous failure modes, and a human correction. It is only intended to
|
||||
validate the pi0.6/ReCAP data path and overfit behavior.
|
||||
"""
|
||||
|
||||
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||
expert = torch.as_tensor(expert_future_xy, dtype=torch.float32)
|
||||
if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 1:
|
||||
raise ValueError("history_xy must have shape [history_steps, 2]")
|
||||
if expert.ndim != 2 or expert.shape[-1] != 2 or len(expert) < 4:
|
||||
raise ValueError("expert_future_xy must have shape [horizon>=4, 2]")
|
||||
horizon = len(expert)
|
||||
anchor = history[-1]
|
||||
displacement = expert - anchor
|
||||
|
||||
understeer = anchor + displacement * expert.new_tensor([1.0, 0.2])
|
||||
stop_index = max(2, horizon // 3)
|
||||
stopped = expert.clone()
|
||||
stopped[stop_index:] = stopped[stop_index - 1]
|
||||
|
||||
progress = torch.linspace(0.0, 1.0, horizon, dtype=expert.dtype).unsqueeze(-1)
|
||||
correction_weight = ((progress - 0.2) / 0.55).clamp(0.0, 1.0)
|
||||
correction = understeer * (1.0 - correction_weight) + expert * correction_weight
|
||||
intervention = torch.arange(horizon) >= max(1, horizon // 4)
|
||||
|
||||
zero_rewards = torch.zeros(horizon, dtype=expert.dtype)
|
||||
return [
|
||||
ExperienceEpisode(
|
||||
f"{episode_prefix}:demo",
|
||||
"demo_success",
|
||||
history,
|
||||
expert,
|
||||
zero_rewards,
|
||||
torch.zeros(horizon, dtype=torch.bool),
|
||||
),
|
||||
ExperienceEpisode(
|
||||
f"{episode_prefix}:understeer",
|
||||
"autonomous_understeer",
|
||||
history,
|
||||
understeer,
|
||||
_trajectory_rewards(understeer, expert),
|
||||
torch.zeros(horizon, dtype=torch.bool),
|
||||
{"failure": "understeer"},
|
||||
),
|
||||
ExperienceEpisode(
|
||||
f"{episode_prefix}:stop",
|
||||
"autonomous_stop",
|
||||
history,
|
||||
stopped,
|
||||
_trajectory_rewards(stopped, expert),
|
||||
torch.zeros(horizon, dtype=torch.bool),
|
||||
{"failure": "premature_stop"},
|
||||
),
|
||||
ExperienceEpisode(
|
||||
f"{episode_prefix}:intervention",
|
||||
"intervention_correction",
|
||||
history,
|
||||
correction,
|
||||
_trajectory_rewards(correction, expert),
|
||||
intervention,
|
||||
{"failure": "understeer_then_corrected"},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def stack_episode_field(
|
||||
episodes: Sequence[ExperienceEpisode], field_name: str
|
||||
) -> torch.Tensor:
|
||||
if not episodes:
|
||||
raise ValueError("episodes must be non-empty")
|
||||
return torch.stack([torch.as_tensor(getattr(item, field_name)) for item in episodes])
|
||||
@@ -0,0 +1,44 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class FlowMatching:
|
||||
"""Low-timestep-biased rectified flow with pi*0.6 noise weighting."""
|
||||
|
||||
def __init__(self, inference_steps: int = 5, cutoff: float = 0.999) -> None:
|
||||
self.inference_steps = inference_steps
|
||||
self.cutoff = cutoff
|
||||
|
||||
def training_batch(self, target: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
|
||||
batch = target.shape[0]
|
||||
uniform = torch.rand(batch, device=target.device, dtype=target.dtype)
|
||||
timestep = self.cutoff * (1.0 - uniform.pow(1.0 / 1.5))
|
||||
noise = torch.randn_like(target)
|
||||
noisy = timestep[:, None, None] * target + (1 - timestep[:, None, None]) * noise
|
||||
velocity = target - noise
|
||||
weight = torch.exp(-timestep / 2.0)
|
||||
return noisy, timestep, velocity, weight
|
||||
|
||||
@staticmethod
|
||||
def weighted_loss(prediction: Tensor, target: Tensor, weight: Tensor) -> Tensor:
|
||||
per_example = (prediction.float() - target.float()).square().mean(dim=(1, 2))
|
||||
return (per_example * weight.float()).mean()
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(
|
||||
self,
|
||||
model_fn: Callable[[Tensor, Tensor], Tensor],
|
||||
shape: tuple[int, ...],
|
||||
device: torch.device | str,
|
||||
steps: int | None = None,
|
||||
) -> Tensor:
|
||||
count = steps or self.inference_steps
|
||||
x = torch.randn(shape, device=device)
|
||||
grid = torch.linspace(0.0, 1.0, count + 1, device=device)
|
||||
for index in range(count):
|
||||
t = grid[index].expand(shape[0])
|
||||
x = x + (grid[index + 1] - grid[index]) * model_fn(x, t)
|
||||
return x
|
||||
@@ -0,0 +1,98 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .data import Pi06ExperienceDataset
|
||||
from .flow import FlowMatching
|
||||
from .model import Pi06Config, Pi06Mini
|
||||
from .recap import classifier_free_guidance
|
||||
from .runtime import choose_device, write_json
|
||||
from .tokenizer import MiniTokenizer
|
||||
from .value import ValueConfig, ValueModel
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
|
||||
checkpoint=torch.load(checkpoint_path,map_location="cpu",weights_only=False)
|
||||
cfg=checkpoint["config"]; root=Path(__file__).resolve().parents[2]
|
||||
device=choose_device(cfg.get("device","auto")); tokenizer=MiniTokenizer()
|
||||
dataset=Pi06ExperienceDataset(
|
||||
root/"data/overfit/overfit_manifest.jsonl",image_size=cfg["data"]["image_size"],
|
||||
augmentation_factor=cfg["experience"]["augmentation_factor"],
|
||||
action_horizon=cfg["data"]["action_horizon"],n_step=cfg["experience"]["n_step"],
|
||||
gamma=cfg["experience"]["gamma"],positive_quantile=cfg["experience"]["positive_quantile"],
|
||||
action_statistics=checkpoint["action_statistics"],
|
||||
)
|
||||
mixed=Pi06Mini(Pi06Config(**checkpoint["policy_config"]));mixed.load_state_dict(checkpoint["mixed"]);mixed.to(device).eval()
|
||||
recap=Pi06Mini(Pi06Config(**checkpoint["policy_config"]));recap.load_state_dict(checkpoint["recap"]);recap.to(device).eval()
|
||||
value=ValueModel(ValueConfig(**checkpoint["value_config"]));value.load_state_dict(checkpoint["value"]);value.to(device).eval()
|
||||
flow=FlowMatching(cfg["inference"]["flow_steps"])
|
||||
demo=next(dataset[i] for i in range(len(dataset)) if dataset[i]["source_type"]=="demo_success")
|
||||
failure=next(dataset[i] for i in range(len(dataset)) if "stop" in dataset[i]["source_type"])
|
||||
image=demo["image"][None].to(device)
|
||||
|
||||
def condition(model:Pi06Mini, advantage:str|None, x:torch.Tensor,t:torch.Tensor)->torch.Tensor:
|
||||
seq=tokenizer.build_policy_sequence(
|
||||
demo["history_xy"],demo["reasoning"],demo["actions"],advantage,
|
||||
metadata=None,
|
||||
)
|
||||
ids=seq.input_ids[:seq.action_start_index][None].to(device)
|
||||
return model(image,ids,x,t,generated_token_start=seq.generated_token_start).vector_field
|
||||
|
||||
def sample_field(fn,seed:int=2):
|
||||
torch.manual_seed(seed)
|
||||
return flow.sample(fn,(1,dataset.action_horizon,2),device,cfg["inference"]["flow_steps"])[0].cpu()
|
||||
|
||||
mixed_actions=sample_field(lambda x,t:condition(mixed,None,x,t))
|
||||
uncond_actions=sample_field(lambda x,t:condition(recap,None,x,t))
|
||||
positive_actions=sample_field(lambda x,t:condition(recap,"positive",x,t))
|
||||
negative_actions=sample_field(lambda x,t:condition(recap,"negative",x,t))
|
||||
def cfg_field(x,t):
|
||||
return classifier_free_guidance(
|
||||
condition(recap,"positive",x,t),condition(recap,None,x,t),
|
||||
guidance_scale=cfg["inference"]["cfg_beta"],
|
||||
)
|
||||
cfg_actions=sample_field(cfg_field)
|
||||
decoded={
|
||||
"mixed":dataset.action_normalizer.decode(mixed_actions,demo["history_xy"]).numpy(),
|
||||
"unconditional":dataset.action_normalizer.decode(uncond_actions,demo["history_xy"]).numpy(),
|
||||
"positive":dataset.action_normalizer.decode(positive_actions,demo["history_xy"]).numpy(),
|
||||
"negative":dataset.action_normalizer.decode(negative_actions,demo["history_xy"]).numpy(),
|
||||
"cfg":dataset.action_normalizer.decode(cfg_actions,demo["history_xy"]).numpy(),
|
||||
}
|
||||
target=demo["future_xy"].numpy(); failure_xy=failure["future_xy"].numpy()
|
||||
metrics={name:float(np.mean((xy-target)**2)) for name,xy in decoded.items()}
|
||||
metrics["negative_to_failure_mse"]=float(np.mean((decoded["negative"]-failure_xy)**2))
|
||||
fig,ax=plt.subplots(figsize=(9,6));ax.plot(target[:,0],target[:,1],"o-",label="success target",linewidth=2)
|
||||
ax.plot(failure_xy[:,0],failure_xy[:,1],":",label="failure experience")
|
||||
for name,color in [("mixed","gray"),("unconditional","purple"),("positive","green"),("negative","red"),("cfg","orange")]:
|
||||
xy=decoded[name];ax.plot(xy[:,0],xy[:,1],"x--",label=f"{name} ({metrics[name]:.3f})",color=color)
|
||||
ax.scatter([0],[0],c="black",marker="s",label="ego");ax.axis("equal");ax.grid(True,alpha=.3);ax.legend(fontsize=8)
|
||||
ax.set_xlabel("forward x (m)");ax.set_ylabel("left y (m)");ax.set_title("π*0.6 Recap advantage-conditioned trajectories")
|
||||
outdir=Path(checkpoint_path).parent;trajectory_path=outdir/"recap_trajectory_comparison.png";fig.tight_layout();fig.savefig(trajectory_path,dpi=170);plt.close(fig)
|
||||
|
||||
value_rows=[]
|
||||
for i in range(len(dataset)):
|
||||
s=dataset[i];ids=tokenizer.build_value_input(s["history_xy"],s["source_type"])[None].to(device)
|
||||
prediction=float(value(s["image"][None].to(device),ids).expectation.cpu())
|
||||
value_rows.append({"source_type":s["source_type"],"target":float(s["value_target"]),"prediction":prediction})
|
||||
fig,ax=plt.subplots(figsize=(8,4));x=np.arange(len(value_rows));w=.35
|
||||
ax.bar(x-w/2,[r["target"] for r in value_rows],w,label="return target")
|
||||
ax.bar(x+w/2,[r["prediction"] for r in value_rows],w,label="value prediction")
|
||||
ax.set_xticks(x,[r["source_type"] for r in value_rows],rotation=20,ha="right");ax.set_ylim(-1.05,.05);ax.grid(True,axis="y",alpha=.3);ax.legend();ax.set_title("201-bin distributional value overfit")
|
||||
value_path=outdir/"value_comparison.png";fig.tight_layout();fig.savefig(value_path,dpi=170);plt.close(fig)
|
||||
write_json(outdir/"evaluation.json",{"trajectory_metrics":metrics,"values":value_rows})
|
||||
print({"trajectory_metrics":metrics,"values":value_rows})
|
||||
return trajectory_path,value_path
|
||||
|
||||
|
||||
def main():
|
||||
parser=argparse.ArgumentParser();parser.add_argument("--checkpoint",default="outputs/overfit/checkpoint.pt")
|
||||
evaluate(parser.parse_args().checkpoint)
|
||||
|
||||
|
||||
if __name__=="__main__":main()
|
||||
@@ -0,0 +1,444 @@
|
||||
r"""Minimal :math:`\pi^*_{0.6}` policy with Knowledge Insulation.
|
||||
|
||||
This is intentionally a small, pure-PyTorch architecture test rather than a
|
||||
scale reproduction. It preserves the paper's important information paths:
|
||||
|
||||
* a VLM stream predicts subtask and FAST-lite action tokens;
|
||||
* an optional positive/negative advantage token sits after the subtask and
|
||||
before both action representations;
|
||||
* a separately parameterized flow expert reads observation/subtask/advantage;
|
||||
* flow queries never read teacher-forced FAST-lite tokens; and
|
||||
* VLM keys and values are stop-gradient inputs to the flow expert (Knowledge
|
||||
Insulation), so continuous-action loss cannot update the VLM.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterator, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
@dataclass
|
||||
class Pi06Config:
|
||||
vocab_size: int = 512
|
||||
action_dim: int = 2
|
||||
image_channels: int = 3
|
||||
image_size: int = 64
|
||||
patch_size: int = 16
|
||||
d_model: int = 128
|
||||
n_heads: int = 4
|
||||
n_layers: int = 4
|
||||
mlp_ratio: int = 4
|
||||
max_text_tokens: int = 96
|
||||
max_action_tokens: int = 32
|
||||
action_start_token_id: int = 269
|
||||
dropout: float = 0.0
|
||||
rms_eps: float = 1e-6
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.d_model % self.n_heads:
|
||||
raise ValueError("d_model must be divisible by n_heads")
|
||||
if self.image_size % self.patch_size:
|
||||
raise ValueError("image_size must be divisible by patch_size")
|
||||
if not 0 <= self.action_start_token_id < self.vocab_size:
|
||||
raise ValueError("action_start_token_id must be inside the vocabulary")
|
||||
|
||||
@property
|
||||
def max_image_tokens(self) -> int:
|
||||
return (self.image_size // self.patch_size) ** 2
|
||||
|
||||
@property
|
||||
def max_prefix_tokens(self) -> int:
|
||||
return self.max_image_tokens + self.max_text_tokens
|
||||
|
||||
|
||||
@dataclass
|
||||
class Pi06Output:
|
||||
text_logits: Tensor
|
||||
vector_field: Tensor
|
||||
prefix_hidden: Tensor
|
||||
action_hidden: Tensor
|
||||
text_token_offset: int
|
||||
loss: Optional[Tensor] = None
|
||||
text_loss: Optional[Tensor] = None
|
||||
flow_loss: Optional[Tensor] = None
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
rms = x.float().pow(2).mean(dim=-1, keepdim=True)
|
||||
return x * torch.rsqrt(rms.to(x.dtype) + self.eps) * self.weight
|
||||
|
||||
|
||||
class AdaRMSNorm(nn.Module):
|
||||
"""RMSNorm modulated at every flow-expert layer by the flow timestep."""
|
||||
|
||||
def __init__(self, dim: int, cond_dim: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.norm = RMSNorm(dim, eps)
|
||||
self.modulation = nn.Linear(cond_dim, 2 * dim)
|
||||
nn.init.zeros_(self.modulation.weight)
|
||||
nn.init.zeros_(self.modulation.bias)
|
||||
|
||||
def forward(self, x: Tensor, condition: Tensor) -> Tensor:
|
||||
shift, scale = self.modulation(condition).chunk(2, dim=-1)
|
||||
return self.norm(x) * (1 + scale[:, None]) + shift[:, None]
|
||||
|
||||
|
||||
class TimeEmbedding(nn.Module):
|
||||
def __init__(self, dim: int) -> None:
|
||||
super().__init__()
|
||||
fourier_dim = max(8, dim // 2)
|
||||
self.fourier_dim = fourier_dim + fourier_dim % 2
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(self.fourier_dim, dim), nn.SiLU(), nn.Linear(dim, dim)
|
||||
)
|
||||
|
||||
def forward(self, timestep: Tensor) -> Tensor:
|
||||
timestep = timestep.reshape(-1).float()
|
||||
half = self.fourier_dim // 2
|
||||
frequencies = torch.exp(
|
||||
torch.linspace(
|
||||
math.log(1.0), math.log(1000.0), half,
|
||||
device=timestep.device, dtype=timestep.dtype,
|
||||
)
|
||||
)
|
||||
angles = 2 * math.pi * timestep[:, None] * frequencies[None]
|
||||
return self.mlp(torch.cat((angles.sin(), angles.cos()), dim=-1))
|
||||
|
||||
|
||||
class AttentionWeights(nn.Module):
|
||||
"""One expert's independent Q/K/V/output projections."""
|
||||
|
||||
def __init__(self, dim: int) -> None:
|
||||
super().__init__()
|
||||
self.q = nn.Linear(dim, dim, bias=False)
|
||||
self.k = nn.Linear(dim, dim, bias=False)
|
||||
self.v = nn.Linear(dim, dim, bias=False)
|
||||
self.out = nn.Linear(dim, dim, bias=False)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim: int, ratio: int, dropout: float) -> None:
|
||||
super().__init__()
|
||||
hidden = dim * ratio
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(dim, hidden),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden, dim),
|
||||
nn.Dropout(dropout),
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.net(x)
|
||||
|
||||
|
||||
def attention(
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
n_heads: int,
|
||||
allowed: Optional[Tensor],
|
||||
dropout: float,
|
||||
training: bool,
|
||||
) -> Tensor:
|
||||
batch, q_len, width = query.shape
|
||||
head_dim = width // n_heads
|
||||
|
||||
def split(x: Tensor) -> Tensor:
|
||||
return x.reshape(batch, -1, n_heads, head_dim).transpose(1, 2)
|
||||
|
||||
q, k, v = split(query), split(key), split(value)
|
||||
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim)
|
||||
if allowed is not None:
|
||||
scores = scores.masked_fill(
|
||||
~allowed[:, None], torch.finfo(scores.dtype).min
|
||||
)
|
||||
probabilities = scores.float().softmax(dim=-1).to(scores.dtype)
|
||||
probabilities = F.dropout(probabilities, p=dropout, training=training)
|
||||
output = probabilities @ v
|
||||
return output.transpose(1, 2).contiguous().reshape(batch, q_len, width)
|
||||
|
||||
|
||||
class KnowledgeInsulatedBlock(nn.Module):
|
||||
"""A dual-expert block whose VLM-to-flow connection is stop-gradient."""
|
||||
|
||||
def __init__(self, config: Pi06Config) -> None:
|
||||
super().__init__()
|
||||
d = config.d_model
|
||||
self.n_heads = config.n_heads
|
||||
self.dropout = config.dropout
|
||||
|
||||
self.vlm_norm1 = RMSNorm(d, config.rms_eps)
|
||||
self.vlm_attention = AttentionWeights(d)
|
||||
self.vlm_norm2 = RMSNorm(d, config.rms_eps)
|
||||
self.vlm_mlp = MLP(d, config.mlp_ratio, config.dropout)
|
||||
|
||||
self.action_norm1 = AdaRMSNorm(d, d, config.rms_eps)
|
||||
self.action_attention = AttentionWeights(d)
|
||||
self.action_norm2 = AdaRMSNorm(d, d, config.rms_eps)
|
||||
self.action_mlp = MLP(d, config.mlp_ratio, config.dropout)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
prefix: Tensor,
|
||||
action: Tensor,
|
||||
time_condition: Tensor,
|
||||
prefix_edges: Tensor,
|
||||
action_edges: Tensor,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
prefix_input = self.vlm_norm1(prefix)
|
||||
p = self.vlm_attention
|
||||
prefix_update = attention(
|
||||
p.q(prefix_input), p.k(prefix_input), p.v(prefix_input),
|
||||
self.n_heads, prefix_edges, self.dropout, self.training,
|
||||
)
|
||||
prefix = prefix + p.out(prefix_update)
|
||||
prefix = prefix + self.vlm_mlp(self.vlm_norm2(prefix))
|
||||
|
||||
action_input = self.action_norm1(action, time_condition)
|
||||
a = self.action_attention
|
||||
# Knowledge Insulation: the flow expert may consume VLM activations,
|
||||
# but the flow loss cannot alter either those activations or the VLM's
|
||||
# K/V projection weights. Text loss still trains the VLM normally.
|
||||
prefix_key = p.k(prefix_input).detach()
|
||||
prefix_value = p.v(prefix_input).detach()
|
||||
key = torch.cat((prefix_key, a.k(action_input)), dim=1)
|
||||
value = torch.cat((prefix_value, a.v(action_input)), dim=1)
|
||||
action_update = attention(
|
||||
a.q(action_input), key, value, self.n_heads,
|
||||
action_edges, self.dropout, self.training,
|
||||
)
|
||||
action = action + a.out(action_update)
|
||||
action = action + self.action_mlp(self.action_norm2(action, time_condition))
|
||||
return prefix, action
|
||||
|
||||
|
||||
class Pi06Mini(nn.Module):
|
||||
r"""Small policy retaining the structural contracts of :math:`\pi^*_{0.6}`."""
|
||||
|
||||
def __init__(self, config: Optional[Pi06Config] = None) -> None:
|
||||
super().__init__()
|
||||
self.config = config or Pi06Config()
|
||||
c = self.config
|
||||
self.patch_encoder = nn.Conv2d(
|
||||
c.image_channels, c.d_model, c.patch_size, stride=c.patch_size
|
||||
)
|
||||
self.token_embedding = nn.Embedding(c.vocab_size, c.d_model)
|
||||
self.action_projection = nn.Linear(c.action_dim, c.d_model)
|
||||
self.prefix_position = nn.Parameter(
|
||||
torch.randn(1, c.max_prefix_tokens, c.d_model) * 0.01
|
||||
)
|
||||
self.action_position = nn.Parameter(
|
||||
torch.randn(1, c.max_action_tokens, c.d_model) * 0.01
|
||||
)
|
||||
self.time_embedding = TimeEmbedding(c.d_model)
|
||||
self.layers = nn.ModuleList(
|
||||
KnowledgeInsulatedBlock(c) for _ in range(c.n_layers)
|
||||
)
|
||||
self.prefix_norm = RMSNorm(c.d_model, c.rms_eps)
|
||||
self.action_norm = AdaRMSNorm(c.d_model, c.d_model, c.rms_eps)
|
||||
self.text_head = nn.Linear(c.d_model, c.vocab_size, bias=False)
|
||||
self.vector_field_head = nn.Linear(c.d_model, c.action_dim)
|
||||
self.text_head.weight = self.token_embedding.weight
|
||||
|
||||
def encode_prefix(self, pixel_values: Tensor, input_ids: Tensor) -> tuple[Tensor, int]:
|
||||
patches = self.patch_encoder(pixel_values).flatten(2).transpose(1, 2)
|
||||
prefix = torch.cat((patches, self.token_embedding(input_ids)), dim=1)
|
||||
if prefix.shape[1] > self.config.max_prefix_tokens:
|
||||
raise ValueError("prefix exceeds configured maximum length")
|
||||
return prefix + self.prefix_position[:, : prefix.shape[1]], patches.shape[1]
|
||||
|
||||
def _derive_action_prefix_mask(self, input_ids: Tensor) -> Tensor:
|
||||
"""Allow flow to read every token strictly before the first ACTION_START."""
|
||||
|
||||
is_action = input_ids.eq(self.config.action_start_token_id)
|
||||
# If there is no teacher-forced FAST stream (normal flow inference),
|
||||
# every supplied observation/subtask/advantage token is readable.
|
||||
return is_action.long().cumsum(dim=1).eq(0)
|
||||
|
||||
def _masks(
|
||||
self,
|
||||
batch: int,
|
||||
prefix_length: int,
|
||||
action_length: int,
|
||||
text_offset: int,
|
||||
attention_mask: Optional[Tensor],
|
||||
action_prefix_mask: Tensor,
|
||||
generated_token_start: Optional[int],
|
||||
device: torch.device,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
valid = torch.ones(batch, prefix_length, dtype=torch.bool, device=device)
|
||||
if attention_mask is not None:
|
||||
valid[:, text_offset:] = attention_mask.bool()
|
||||
|
||||
generated_start = (
|
||||
prefix_length
|
||||
if generated_token_start is None
|
||||
else text_offset + generated_token_start
|
||||
)
|
||||
if not text_offset <= generated_start <= prefix_length:
|
||||
raise ValueError("generated_token_start must index input_ids")
|
||||
q = torch.arange(prefix_length, device=device)[:, None]
|
||||
k = torch.arange(prefix_length, device=device)[None]
|
||||
structural = torch.where(q < generated_start, k < generated_start, k <= q)
|
||||
prefix_edges = structural[None] & valid[:, None]
|
||||
identity = torch.eye(prefix_length, dtype=torch.bool, device=device)[None]
|
||||
prefix_edges = torch.where(valid[:, :, None], prefix_edges, identity)
|
||||
|
||||
supplied = action_prefix_mask.bool()
|
||||
if supplied.shape[1] == prefix_length - text_offset:
|
||||
supplied = torch.cat(
|
||||
(
|
||||
torch.ones(batch, text_offset, dtype=torch.bool, device=device),
|
||||
supplied,
|
||||
),
|
||||
dim=1,
|
||||
)
|
||||
if supplied.shape != valid.shape:
|
||||
raise ValueError("action_prefix_mask must cover text tokens or the full prefix")
|
||||
readable = valid & supplied
|
||||
prefix_keys = readable[:, None].expand(batch, action_length, prefix_length)
|
||||
action_keys = torch.ones(
|
||||
batch, action_length, action_length, dtype=torch.bool, device=device
|
||||
)
|
||||
return prefix_edges, torch.cat((prefix_keys, action_keys), dim=-1)
|
||||
|
||||
def vlm_parameters(self) -> Iterator[nn.Parameter]:
|
||||
"""Iterate unique parameters protected from flow-loss gradients by KI."""
|
||||
|
||||
modules: tuple[nn.Module, ...] = (
|
||||
self.patch_encoder,
|
||||
self.token_embedding,
|
||||
self.prefix_norm,
|
||||
self.text_head,
|
||||
)
|
||||
seen: set[int] = set()
|
||||
for parameter in (self.prefix_position,):
|
||||
seen.add(id(parameter))
|
||||
yield parameter
|
||||
for layer in self.layers:
|
||||
modules += (
|
||||
layer.vlm_norm1,
|
||||
layer.vlm_attention,
|
||||
layer.vlm_norm2,
|
||||
layer.vlm_mlp,
|
||||
)
|
||||
for module in modules:
|
||||
for parameter in module.parameters():
|
||||
if id(parameter) not in seen:
|
||||
seen.add(id(parameter))
|
||||
yield parameter
|
||||
|
||||
def action_parameters(self) -> Iterator[nn.Parameter]:
|
||||
"""Iterate parameters trained by the continuous flow loss."""
|
||||
|
||||
modules: list[nn.Module] = [
|
||||
self.action_projection,
|
||||
self.time_embedding,
|
||||
self.action_norm,
|
||||
self.vector_field_head,
|
||||
]
|
||||
for layer in self.layers:
|
||||
modules.extend(
|
||||
(
|
||||
layer.action_norm1,
|
||||
layer.action_attention,
|
||||
layer.action_norm2,
|
||||
layer.action_mlp,
|
||||
)
|
||||
)
|
||||
yield self.action_position
|
||||
for module in modules:
|
||||
yield from module.parameters()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pixel_values: Tensor,
|
||||
input_ids: Tensor,
|
||||
noisy_actions: Tensor,
|
||||
timesteps: Tensor,
|
||||
attention_mask: Optional[Tensor] = None,
|
||||
action_prefix_mask: Optional[Tensor] = None,
|
||||
generated_token_start: Optional[int] = None,
|
||||
text_labels: Optional[Tensor] = None,
|
||||
target_velocity: Optional[Tensor] = None,
|
||||
text_loss_weight: float = 1.0,
|
||||
flow_loss_weight: float = 1.0,
|
||||
) -> Pi06Output:
|
||||
if noisy_actions.ndim != 3:
|
||||
raise ValueError("noisy_actions must have shape [batch, horizon, action_dim]")
|
||||
if noisy_actions.shape[-1] != self.config.action_dim:
|
||||
raise ValueError("noisy_actions has the wrong action dimension")
|
||||
if noisy_actions.shape[1] > self.config.max_action_tokens:
|
||||
raise ValueError("action horizon exceeds configured maximum")
|
||||
|
||||
prefix, text_offset = self.encode_prefix(pixel_values, input_ids)
|
||||
horizon = noisy_actions.shape[1]
|
||||
action = self.action_projection(noisy_actions)
|
||||
action = action + self.action_position[:, :horizon]
|
||||
condition = self.time_embedding(timesteps).to(action.dtype)
|
||||
if action_prefix_mask is None:
|
||||
action_prefix_mask = self._derive_action_prefix_mask(input_ids)
|
||||
prefix_edges, action_edges = self._masks(
|
||||
prefix.shape[0], prefix.shape[1], horizon, text_offset,
|
||||
attention_mask, action_prefix_mask, generated_token_start, prefix.device,
|
||||
)
|
||||
for layer in self.layers:
|
||||
prefix, action = layer(
|
||||
prefix, action, condition, prefix_edges, action_edges
|
||||
)
|
||||
|
||||
prefix_hidden = self.prefix_norm(prefix)
|
||||
action_hidden = self.action_norm(action, condition)
|
||||
text_logits = self.text_head(prefix_hidden[:, text_offset:])
|
||||
vector_field = self.vector_field_head(action_hidden)
|
||||
|
||||
text_loss = None
|
||||
if text_labels is not None:
|
||||
text_loss = F.cross_entropy(
|
||||
text_logits.reshape(-1, text_logits.shape[-1]),
|
||||
text_labels.reshape(-1), ignore_index=-100,
|
||||
)
|
||||
flow_loss = None
|
||||
if target_velocity is not None:
|
||||
flow_loss = F.mse_loss(vector_field, target_velocity)
|
||||
terms = [
|
||||
term for term in (
|
||||
None if text_loss is None else text_loss_weight * text_loss,
|
||||
None if flow_loss is None else flow_loss_weight * flow_loss,
|
||||
) if term is not None
|
||||
]
|
||||
loss = sum(terms) if terms else None
|
||||
return Pi06Output(
|
||||
text_logits=text_logits,
|
||||
vector_field=vector_field,
|
||||
prefix_hidden=prefix_hidden,
|
||||
action_hidden=action_hidden,
|
||||
text_token_offset=text_offset,
|
||||
loss=loss,
|
||||
text_loss=text_loss,
|
||||
flow_loss=flow_loss,
|
||||
)
|
||||
|
||||
|
||||
def flow_matching_sample(actions: Tensor) -> tuple[Tensor, Tensor, Tensor]:
|
||||
"""Construct the linear conditional-flow target used by the policy."""
|
||||
|
||||
batch = actions.shape[0]
|
||||
timestep = torch.rand(batch, device=actions.device, dtype=actions.dtype)
|
||||
noise = torch.randn_like(actions)
|
||||
broadcast_t = timestep.reshape(batch, *([1] * (actions.ndim - 1)))
|
||||
noisy_actions = (1 - broadcast_t) * noise + broadcast_t * actions
|
||||
return noisy_actions, timestep, actions - noise
|
||||
@@ -0,0 +1,78 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def condition_dropout_mask(
|
||||
batch_size: int,
|
||||
*,
|
||||
drop_probability: float = 0.3,
|
||||
device: torch.device | str | None = None,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Sample the 30% advantage-condition dropout used for CFG training."""
|
||||
|
||||
if batch_size < 1:
|
||||
raise ValueError("batch_size must be positive")
|
||||
if not 0.0 <= drop_probability <= 1.0:
|
||||
raise ValueError("drop_probability must lie in [0, 1]")
|
||||
return torch.rand(batch_size, device=device, generator=generator) < drop_probability
|
||||
|
||||
|
||||
def apply_condition_dropout(
|
||||
condition: torch.Tensor,
|
||||
null_condition: torch.Tensor | float | int = 0,
|
||||
*,
|
||||
drop_probability: float = 0.3,
|
||||
mask: torch.Tensor | None = None,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Replace complete per-example ReCAP conditions with a null condition."""
|
||||
|
||||
condition = torch.as_tensor(condition)
|
||||
if condition.ndim < 1 or condition.shape[0] < 1:
|
||||
raise ValueError("condition must have a non-empty batch dimension")
|
||||
if mask is None:
|
||||
mask = condition_dropout_mask(
|
||||
condition.shape[0],
|
||||
drop_probability=drop_probability,
|
||||
device=condition.device,
|
||||
generator=generator,
|
||||
)
|
||||
else:
|
||||
mask = torch.as_tensor(mask, dtype=torch.bool, device=condition.device)
|
||||
if mask.shape != (condition.shape[0],):
|
||||
raise ValueError("mask must have shape [batch_size]")
|
||||
null = torch.as_tensor(null_condition, dtype=condition.dtype, device=condition.device)
|
||||
batch_mask = mask
|
||||
while mask.ndim < condition.ndim:
|
||||
mask = mask.unsqueeze(-1)
|
||||
dropped = torch.where(mask, null, condition)
|
||||
return dropped, batch_mask
|
||||
|
||||
|
||||
def classifier_free_guidance(
|
||||
conditional_prediction: torch.Tensor,
|
||||
unconditional_prediction: torch.Tensor,
|
||||
*,
|
||||
guidance_scale: float = 1.5,
|
||||
) -> torch.Tensor:
|
||||
"""Combine conditional and null-condition flow fields.
|
||||
|
||||
A scale of zero returns the unconditional prediction and a scale of one
|
||||
returns the ordinary positive-condition prediction.
|
||||
"""
|
||||
|
||||
conditional_prediction = torch.as_tensor(conditional_prediction)
|
||||
unconditional_prediction = torch.as_tensor(
|
||||
unconditional_prediction,
|
||||
dtype=conditional_prediction.dtype,
|
||||
device=conditional_prediction.device,
|
||||
)
|
||||
if conditional_prediction.shape != unconditional_prediction.shape:
|
||||
raise ValueError("conditional and unconditional predictions must have equal shape")
|
||||
if guidance_scale < 0.0:
|
||||
raise ValueError("guidance_scale must be non-negative")
|
||||
return unconditional_prediction + guidance_scale * (
|
||||
conditional_prediction - unconditional_prediction
|
||||
)
|
||||
@@ -0,0 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
|
||||
def load_config(path: str | Path) -> dict:
|
||||
with Path(path).open("r", encoding="utf-8") as handle:
|
||||
return yaml.safe_load(handle)
|
||||
|
||||
|
||||
def seed_everything(seed: int) -> None:
|
||||
random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
|
||||
|
||||
|
||||
def choose_device(name: str = "auto") -> torch.device:
|
||||
if name != "auto": return torch.device(name)
|
||||
if torch.cuda.is_available(): return torch.device("cuda")
|
||||
if torch.backends.mps.is_available(): return torch.device("mps")
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
def write_json(path: str | Path, value: object) -> None:
|
||||
target = Path(path); target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(json.dumps(value, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
@@ -0,0 +1,330 @@
|
||||
r"""Tiny discrete state, subtask, advantage, and FAST-lite tokenization.
|
||||
|
||||
The layout mirrors the ordering used by :math:`\pi^*_{0.6}`: the model first
|
||||
predicts a subtask, then receives an optional advantage indicator, and only
|
||||
then predicts the discrete action representation. ``omitted`` is represented
|
||||
by the physical absence of an indicator token, as required for unconditional
|
||||
classifier-free guidance training.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, Literal
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
Advantage = Literal["positive", "negative", "omitted"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CoordinateRange:
|
||||
minimum: float
|
||||
maximum: float
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.maximum <= self.minimum:
|
||||
raise ValueError("coordinate maximum must be greater than minimum")
|
||||
|
||||
|
||||
class CoordinateTokenizer:
|
||||
"""Uniformly discretize XY state values into literal byte tokens 0..255."""
|
||||
|
||||
bins = 256
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
x_range: tuple[float, float] = (-10.0, 10.0),
|
||||
y_range: tuple[float, float] = (-5.0, 5.0),
|
||||
) -> None:
|
||||
self.ranges = (CoordinateRange(*x_range), CoordinateRange(*y_range))
|
||||
|
||||
def encode(self, coordinates: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||
values = torch.as_tensor(coordinates, dtype=torch.float32)
|
||||
if values.shape[-1] != 2:
|
||||
raise ValueError("coordinates must end in an XY dimension of size 2")
|
||||
low = values.new_tensor([axis.minimum for axis in self.ranges])
|
||||
high = values.new_tensor([axis.maximum for axis in self.ranges])
|
||||
unit = ((values - low) / (high - low)).clamp(0.0, 1.0)
|
||||
return torch.round(unit * (self.bins - 1)).long()
|
||||
|
||||
def decode(self, tokens: Tensor | Iterable[Iterable[int]]) -> Tensor:
|
||||
encoded = torch.as_tensor(tokens, dtype=torch.long)
|
||||
if encoded.shape[-1] != 2:
|
||||
raise ValueError("coordinate tokens must end in an XY dimension of size 2")
|
||||
if encoded.numel() and (encoded.min() < 0 or encoded.max() >= self.bins):
|
||||
raise ValueError("coordinate tokens must be in the inclusive range 0..255")
|
||||
low = torch.tensor([axis.minimum for axis in self.ranges])
|
||||
high = torch.tensor([axis.maximum for axis in self.ranges])
|
||||
return low + encoded.float() / (self.bins - 1) * (high - low)
|
||||
|
||||
|
||||
class SubtaskVocabulary:
|
||||
labels = (
|
||||
"keep_lane",
|
||||
"steer_left",
|
||||
"steer_right",
|
||||
"slow_down",
|
||||
"stop",
|
||||
"accelerate",
|
||||
"unknown",
|
||||
)
|
||||
|
||||
def __init__(self, first_token_id: int = 260) -> None:
|
||||
self.first_token_id = int(first_token_id)
|
||||
self._label_to_token = {
|
||||
label: self.first_token_id + index for index, label in enumerate(self.labels)
|
||||
}
|
||||
self._token_to_label = {token: label for label, token in self._label_to_token.items()}
|
||||
|
||||
@property
|
||||
def token_ids(self) -> tuple[int, ...]:
|
||||
return tuple(self._label_to_token[label] for label in self.labels)
|
||||
|
||||
def classify(self, reasoning: str) -> str:
|
||||
text = reasoning.casefold()
|
||||
if any(term in text for term in ("stop", "halt", "stationary")):
|
||||
return "stop"
|
||||
if any(term in text for term in ("slow", "deceler", "reduce speed", "brake")):
|
||||
return "slow_down"
|
||||
if any(term in text for term in ("accelerat", "speed up")):
|
||||
return "accelerate"
|
||||
if any(term in text for term in ("left", "counterclockwise")):
|
||||
return "steer_left"
|
||||
if any(term in text for term in ("right", "clockwise")):
|
||||
return "steer_right"
|
||||
if any(term in text for term in ("lane", "straight", "follow", "continue", "proceed")):
|
||||
return "keep_lane"
|
||||
return "unknown"
|
||||
|
||||
def encode(self, label_or_reasoning: str) -> int:
|
||||
label = (
|
||||
label_or_reasoning
|
||||
if label_or_reasoning in self._label_to_token
|
||||
else self.classify(label_or_reasoning)
|
||||
)
|
||||
return self._label_to_token[label]
|
||||
|
||||
def decode(self, token_id: int | Tensor) -> str:
|
||||
try:
|
||||
return self._token_to_label[int(token_id)]
|
||||
except KeyError as error:
|
||||
raise ValueError(f"unknown subtask token: {int(token_id)}") from error
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PolicySequence:
|
||||
"""One policy token sequence plus masks needed by ``Pi06Mini``.
|
||||
|
||||
``labels`` is already shifted for next-token prediction. The advantage
|
||||
token itself is never a prediction target because it is externally
|
||||
supplied conditioning. ``action_prefix_mask`` allows the continuous
|
||||
expert to read observation, predicted subtask, and advantage (if present),
|
||||
while insulating it from teacher-forced FAST-lite tokens.
|
||||
"""
|
||||
|
||||
input_ids: Tensor
|
||||
labels: Tensor
|
||||
generated_token_start: int
|
||||
action_prefix_mask: Tensor
|
||||
subtask_index: int
|
||||
advantage_index: int | None
|
||||
action_start_index: int
|
||||
|
||||
|
||||
class MiniTokenizer:
|
||||
r"""Compose the minimal :math:`\pi^*_{0.6}` policy vocabulary."""
|
||||
|
||||
BOS = 256
|
||||
IMAGE = 257
|
||||
STATE = 258
|
||||
SUBTASK_QUERY = 259
|
||||
ADVANTAGE_POSITIVE = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
|
||||
ADVANTAGE_NEGATIVE = ADVANTAGE_POSITIVE + 1
|
||||
ACTION_START = ADVANTAGE_NEGATIVE + 1
|
||||
ACTION_END = ACTION_START + 1
|
||||
ACTION_TOKEN_START = ACTION_END + 1
|
||||
ACTION_BINS = 64
|
||||
ACTION_TOKEN_END = ACTION_TOKEN_START + ACTION_BINS - 1
|
||||
METADATA_DEMO = ACTION_TOKEN_END + 1
|
||||
METADATA_UNDERSTEER = METADATA_DEMO + 1
|
||||
METADATA_STOP = METADATA_DEMO + 2
|
||||
METADATA_INTERVENTION = METADATA_DEMO + 3
|
||||
ACTION_MINIMUM = -4.0
|
||||
ACTION_MAXIMUM = 4.0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
coordinate_tokenizer: CoordinateTokenizer | None = None,
|
||||
subtasks: SubtaskVocabulary | None = None,
|
||||
) -> None:
|
||||
self.coordinates = coordinate_tokenizer or CoordinateTokenizer()
|
||||
self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1)
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return self.METADATA_INTERVENTION + 1
|
||||
|
||||
def encode_metadata(self, source_type: str | None) -> Tensor | None:
|
||||
if source_type is None:
|
||||
return None
|
||||
if source_type == "demo_success": return torch.tensor(self.METADATA_DEMO)
|
||||
if "understeer" in source_type: return torch.tensor(self.METADATA_UNDERSTEER)
|
||||
if "stop" in source_type: return torch.tensor(self.METADATA_STOP)
|
||||
if "intervention" in source_type: return torch.tensor(self.METADATA_INTERVENTION)
|
||||
raise ValueError(f"unknown metadata source_type: {source_type}")
|
||||
|
||||
@property
|
||||
def action_token_ids(self) -> range:
|
||||
return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)
|
||||
|
||||
def build_observation(self, history_xy: Tensor, metadata: str | None = None) -> Tensor:
|
||||
coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1).cpu()
|
||||
specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
|
||||
query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
|
||||
parts = [specials, coordinate_tokens]
|
||||
metadata_token = self.encode_metadata(metadata)
|
||||
if metadata_token is not None: parts.append(metadata_token.reshape(1))
|
||||
parts.append(query)
|
||||
return torch.cat(parts)
|
||||
|
||||
def build_value_input(self, history_xy: Tensor, metadata: str | None = None) -> Tensor:
|
||||
"""Return image marker plus discretized state, without policy query tokens."""
|
||||
|
||||
coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1).cpu()
|
||||
specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
|
||||
parts = [specials, coordinate_tokens]
|
||||
metadata_token = self.encode_metadata(metadata)
|
||||
if metadata_token is not None: parts.append(metadata_token.reshape(1))
|
||||
return torch.cat(parts)
|
||||
|
||||
# Backward-compatible name used by the pi0.5 data utilities.
|
||||
build_prefix = build_observation
|
||||
|
||||
def encode_subtask(self, reasoning: str) -> Tensor:
|
||||
return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long)
|
||||
|
||||
def encode_advantage(self, advantage: Advantage | bool | None) -> Tensor | None:
|
||||
if advantage is True or advantage == "positive":
|
||||
return torch.tensor(self.ADVANTAGE_POSITIVE, dtype=torch.long)
|
||||
if advantage is False or advantage == "negative":
|
||||
return torch.tensor(self.ADVANTAGE_NEGATIVE, dtype=torch.long)
|
||||
if advantage is None or advantage == "omitted":
|
||||
return None
|
||||
raise ValueError("advantage must be positive, negative, omitted, bool, or None")
|
||||
|
||||
@staticmethod
|
||||
def _orthonormal_dct_matrix(
|
||||
horizon: int, *, device: torch.device, dtype: torch.dtype
|
||||
) -> Tensor:
|
||||
if horizon <= 0:
|
||||
raise ValueError("DCT horizon must be positive")
|
||||
time = torch.arange(horizon, device=device, dtype=dtype)
|
||||
frequency = torch.arange(horizon, device=device, dtype=dtype).unsqueeze(1)
|
||||
matrix = torch.cos(math.pi / horizon * (time + 0.5) * frequency)
|
||||
matrix[0] *= math.sqrt(1.0 / horizon)
|
||||
if horizon > 1:
|
||||
matrix[1:] *= math.sqrt(2.0 / horizon)
|
||||
return matrix
|
||||
|
||||
def dct_actions(self, actions: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||
values = torch.as_tensor(actions, dtype=torch.float32)
|
||||
if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
|
||||
raise ValueError("actions must have non-empty shape [horizon, 2]")
|
||||
transform = self._orthonormal_dct_matrix(
|
||||
values.shape[0], device=values.device, dtype=values.dtype
|
||||
)
|
||||
return transform @ values
|
||||
|
||||
def idct_actions(self, coefficients: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||
values = torch.as_tensor(coefficients, dtype=torch.float32)
|
||||
if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
|
||||
raise ValueError("DCT coefficients must have non-empty shape [horizon, 2]")
|
||||
transform = self._orthonormal_dct_matrix(
|
||||
values.shape[0], device=values.device, dtype=values.dtype
|
||||
)
|
||||
return transform.transpose(0, 1) @ values
|
||||
|
||||
def encode_actions(self, actions: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||
values = torch.as_tensor(actions, dtype=torch.float32)
|
||||
if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
|
||||
raise ValueError("actions must have non-empty shape [horizon, 2]")
|
||||
values = values.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
|
||||
coefficients = self.dct_actions(values) / math.sqrt(values.shape[0])
|
||||
coefficients = torch.where(
|
||||
coefficients.abs() < 1e-6, torch.zeros_like(coefficients), coefficients
|
||||
)
|
||||
unit = (
|
||||
coefficients.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
|
||||
- self.ACTION_MINIMUM
|
||||
) / (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
|
||||
payload = torch.round(unit * (self.ACTION_BINS - 1)).long().reshape(-1)
|
||||
payload = payload + self.ACTION_TOKEN_START
|
||||
boundaries = torch.tensor(
|
||||
[self.ACTION_START, self.ACTION_END], dtype=torch.long, device=payload.device
|
||||
)
|
||||
return torch.cat((boundaries[:1], payload, boundaries[1:]))
|
||||
|
||||
def decode_actions(self, tokens: Tensor | Iterable[int]) -> Tensor:
|
||||
encoded = torch.as_tensor(tokens, dtype=torch.long)
|
||||
if encoded.ndim != 1:
|
||||
raise ValueError("action tokens must be a one-dimensional sequence")
|
||||
if encoded.numel() < 4 or encoded[0] != self.ACTION_START or encoded[-1] != self.ACTION_END:
|
||||
raise ValueError("action tokens must be enclosed by ACTION_START and ACTION_END")
|
||||
payload = encoded[1:-1]
|
||||
if payload.numel() % 2:
|
||||
raise ValueError("action token payload must contain complete XY pairs")
|
||||
if payload.min() < self.ACTION_TOKEN_START or payload.max() > self.ACTION_TOKEN_END:
|
||||
raise ValueError("action token payload is outside the reserved action token interval")
|
||||
quantized = payload - self.ACTION_TOKEN_START
|
||||
coefficients = (
|
||||
quantized.float() / (self.ACTION_BINS - 1)
|
||||
* (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
|
||||
+ self.ACTION_MINIMUM
|
||||
)
|
||||
horizon = payload.numel() // 2
|
||||
return self.idct_actions(coefficients.reshape(horizon, 2) * math.sqrt(horizon))
|
||||
|
||||
def build_policy_sequence(
|
||||
self,
|
||||
history_xy: Tensor,
|
||||
subtask: str,
|
||||
actions: Tensor,
|
||||
advantage: Advantage | bool | None,
|
||||
metadata: str | None = None,
|
||||
) -> PolicySequence:
|
||||
observation = self.build_observation(history_xy, metadata)
|
||||
subtask_token = self.encode_subtask(subtask).reshape(1)
|
||||
advantage_token = self.encode_advantage(advantage)
|
||||
action_tokens = self.encode_actions(actions).cpu()
|
||||
pieces = [observation, subtask_token]
|
||||
if advantage_token is not None:
|
||||
pieces.append(advantage_token.reshape(1))
|
||||
pieces.append(action_tokens)
|
||||
input_ids = torch.cat(pieces)
|
||||
|
||||
subtask_index = observation.numel()
|
||||
advantage_index = subtask_index + 1 if advantage_token is not None else None
|
||||
action_start_index = subtask_index + 1 + int(advantage_token is not None)
|
||||
|
||||
labels = torch.full_like(input_ids, -100)
|
||||
# SUBTASK_QUERY predicts the subtask. The external advantage indicator
|
||||
# is deliberately skipped as a target. Advantage/subtask then predicts
|
||||
# ACTION_START and the remaining FAST-lite stream autoregressively.
|
||||
labels[subtask_index - 1] = input_ids[subtask_index]
|
||||
labels[action_start_index - 1 : -1] = input_ids[action_start_index:]
|
||||
|
||||
action_prefix_mask = torch.zeros_like(input_ids, dtype=torch.bool)
|
||||
action_prefix_mask[:action_start_index] = True
|
||||
return PolicySequence(
|
||||
input_ids=input_ids,
|
||||
labels=labels,
|
||||
generated_token_start=subtask_index,
|
||||
action_prefix_mask=action_prefix_mask,
|
||||
subtask_index=subtask_index,
|
||||
advantage_index=advantage_index,
|
||||
action_start_index=action_start_index,
|
||||
)
|
||||
@@ -0,0 +1,153 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
from pathlib import Path
|
||||
import torch
|
||||
|
||||
from .data import Pi06ExperienceDataset, materialize_subset
|
||||
from .flow import FlowMatching
|
||||
from .model import Pi06Config, Pi06Mini
|
||||
from .recap import apply_condition_dropout
|
||||
from .runtime import choose_device, load_config, seed_everything, write_json
|
||||
from .tokenizer import MiniTokenizer
|
||||
from .value import ValueConfig, ValueModel
|
||||
|
||||
|
||||
def build_models(cfg: dict, tokenizer: MiniTokenizer) -> tuple[Pi06Mini, ValueModel]:
|
||||
m = cfg["model"]
|
||||
policy = Pi06Mini(Pi06Config(
|
||||
vocab_size=tokenizer.vocab_size, action_dim=2,
|
||||
image_size=cfg["data"]["image_size"], patch_size=m["patch_size"],
|
||||
d_model=m["policy_d_model"], n_heads=m["n_heads"], n_layers=m["n_layers"],
|
||||
mlp_ratio=m["mlp_ratio"], max_text_tokens=m["max_text_tokens"],
|
||||
max_action_tokens=cfg["data"]["action_horizon"],
|
||||
action_start_token_id=tokenizer.ACTION_START, dropout=m["dropout"],
|
||||
))
|
||||
value = ValueModel(ValueConfig(
|
||||
vocab_size=tokenizer.vocab_size, image_size=cfg["data"]["image_size"],
|
||||
patch_size=m["patch_size"], d_model=m["value_d_model"], n_heads=m["n_heads"],
|
||||
n_layers=2, mlp_ratio=m["mlp_ratio"], max_text_tokens=m["max_text_tokens"],
|
||||
num_bins=m["value_bins"], value_min=m["value_min"], value_max=m["value_max"],
|
||||
))
|
||||
return policy, value
|
||||
|
||||
|
||||
def policy_step(
|
||||
model: Pi06Mini, sample: dict, tokenizer: MiniTokenizer, flow: FlowMatching,
|
||||
device: torch.device, advantage: str | None, cfg: dict,
|
||||
) -> tuple[torch.Tensor, dict[str, float]]:
|
||||
sequence = tokenizer.build_policy_sequence(
|
||||
sample["history_xy"], sample["reasoning"], sample["actions"], advantage,
|
||||
# Outcome/source metadata is reserved for the critic in this synthetic
|
||||
# proxy. The policy must distinguish behavior through advantage alone.
|
||||
metadata=None,
|
||||
)
|
||||
image = sample["image"][None].to(device)
|
||||
ids = sequence.input_ids[None].to(device)
|
||||
labels = sequence.labels[None].to(device)
|
||||
actions = sample["actions"][None].to(device)
|
||||
noisy, timestep, velocity, weight = flow.training_batch(actions)
|
||||
output = model(
|
||||
image, ids, noisy, timestep,
|
||||
action_prefix_mask=sequence.action_prefix_mask[None].to(device),
|
||||
generated_token_start=sequence.generated_token_start,
|
||||
text_labels=labels,
|
||||
)
|
||||
flow_loss = flow.weighted_loss(output.vector_field, velocity, weight)
|
||||
loss = cfg["train"]["token_loss_weight"] * output.text_loss + cfg["train"]["flow_loss_weight"] * flow_loss
|
||||
return loss, {
|
||||
"text_loss": float(output.text_loss.detach().cpu()),
|
||||
"flow_loss": float(flow_loss.detach().cpu()),
|
||||
}
|
||||
|
||||
|
||||
def train(config_path: str) -> Path:
|
||||
cfg = load_config(config_path)
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
seed_everything(cfg["seed"])
|
||||
device = choose_device(cfg.get("device", "auto"))
|
||||
manifest = materialize_subset(cfg["data"]["source_manifest"], root / "data/overfit")
|
||||
dataset = Pi06ExperienceDataset(
|
||||
manifest, image_size=cfg["data"]["image_size"],
|
||||
augmentation_factor=cfg["experience"]["augmentation_factor"],
|
||||
action_horizon=cfg["data"]["action_horizon"], n_step=cfg["experience"]["n_step"],
|
||||
gamma=cfg["experience"]["gamma"], positive_quantile=cfg["experience"]["positive_quantile"],
|
||||
)
|
||||
tokenizer = MiniTokenizer()
|
||||
base, value = build_models(cfg, tokenizer)
|
||||
base, value = base.to(device), value.to(device)
|
||||
flow = FlowMatching(cfg["inference"]["flow_steps"])
|
||||
outdir = root / cfg["output_dir"]; outdir.mkdir(parents=True, exist_ok=True)
|
||||
history: list[dict] = []
|
||||
|
||||
demos = [dataset[i] for i in range(len(dataset)) if dataset[i]["source_type"] == "demo_success"]
|
||||
optimizer = torch.optim.AdamW(base.parameters(), lr=cfg["train"]["learning_rate"])
|
||||
base.train()
|
||||
for step in range(1, cfg["train"]["base_steps"] + 1):
|
||||
sample = demos[(step - 1) % len(demos)]
|
||||
loss, metrics = policy_step(base, sample, tokenizer, flow, device, None, cfg)
|
||||
optimizer.zero_grad(set_to_none=True); loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(base.parameters(), cfg["train"]["grad_clip_norm"]); optimizer.step()
|
||||
if step == 1 or step % cfg["train"]["log_every"] == 0:
|
||||
row={"stage":"base","step":step,"loss":float(loss.detach().cpu()),**metrics};history.append(row);print(json.dumps(row))
|
||||
base_state = copy.deepcopy(base.state_dict())
|
||||
torch.save(base_state, outdir / "pi06_base.pt")
|
||||
|
||||
value_opt = torch.optim.AdamW(value.parameters(), lr=cfg["train"]["value_learning_rate"])
|
||||
value.train()
|
||||
for step in range(1, cfg["train"]["value_steps"] + 1):
|
||||
sample = dataset[(step - 1) % len(dataset)]
|
||||
ids = tokenizer.build_value_input(sample["history_xy"], sample["source_type"])[None].to(device)
|
||||
output = value(sample["image"][None].to(device), ids, target_values=sample["value_target"][None].to(device))
|
||||
value_opt.zero_grad(set_to_none=True); output.loss.backward(); value_opt.step()
|
||||
if step == 1 or step % cfg["train"]["log_every"] == 0:
|
||||
row={"stage":"value","step":step,"loss":float(output.loss.detach().cpu()),"prediction":float(output.expectation.detach().cpu()),"target":float(sample["value_target"])};history.append(row);print(json.dumps(row))
|
||||
|
||||
def fit_policy(stage: str, steps: int, conditioned: bool) -> Pi06Mini:
|
||||
model, _ = build_models(cfg, tokenizer); model.load_state_dict(base_state); model.to(device).train()
|
||||
opt = torch.optim.AdamW(model.parameters(), lr=cfg["train"]["learning_rate"])
|
||||
for step in range(1, steps + 1):
|
||||
sample = dataset[(step - 1) % len(dataset)]
|
||||
advantage = None
|
||||
if conditioned:
|
||||
dropped = bool(apply_condition_dropout(
|
||||
torch.tensor([sample["positive_condition"]]),
|
||||
drop_probability=cfg["recap"]["condition_dropout"],
|
||||
)[1][0])
|
||||
advantage = None if dropped else sample["advantage_label"]
|
||||
loss, metrics = policy_step(model, sample, tokenizer, flow, device, advantage, cfg)
|
||||
opt.zero_grad(set_to_none=True); loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"]); opt.step()
|
||||
if step == 1 or step % cfg["train"]["log_every"] == 0:
|
||||
row={"stage":stage,"step":step,"loss":float(loss.detach().cpu()),**metrics};history.append(row);print(json.dumps(row))
|
||||
return model
|
||||
|
||||
mixed = fit_policy("mixed", cfg["train"]["mixed_steps"], conditioned=False)
|
||||
recap = fit_policy("recap", cfg["train"]["recap_steps"], conditioned=True)
|
||||
torch.save(mixed.state_dict(), outdir / "mixed_sft.pt")
|
||||
torch.save(recap.state_dict(), outdir / "pi06star_recap.pt")
|
||||
checkpoint = {
|
||||
"base": base_state, "mixed": mixed.state_dict(), "recap": recap.state_dict(),
|
||||
"value": value.state_dict(), "policy_config": base.config.__dict__,
|
||||
"value_config": value.config.__dict__, "config": cfg,
|
||||
"action_statistics": dataset.action_normalizer.as_dict(),
|
||||
}
|
||||
torch.save(checkpoint, outdir / "checkpoint.pt")
|
||||
write_json(outdir / "training_history.json", history)
|
||||
write_json(outdir / "summary.json", {
|
||||
"device": str(device), "experience_samples": len(dataset),
|
||||
"policy_parameters": sum(p.numel() for p in base.parameters()),
|
||||
"value_parameters": sum(p.numel() for p in value.parameters()),
|
||||
"sources": [dataset[i]["source_type"] for i in range(len(dataset))],
|
||||
})
|
||||
return outdir / "checkpoint.pt"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser=argparse.ArgumentParser();parser.add_argument("--config",default="configs/overfit.yaml")
|
||||
train(parser.parse_args().config)
|
||||
|
||||
|
||||
if __name__ == "__main__": main()
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Independent 201-bin distributional value model for minimal Recap training."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .model import AttentionWeights, MLP, RMSNorm, attention
|
||||
|
||||
|
||||
@dataclass
|
||||
class ValueConfig:
|
||||
vocab_size: int = 512
|
||||
image_channels: int = 3
|
||||
image_size: int = 64
|
||||
patch_size: int = 16
|
||||
d_model: int = 96
|
||||
n_heads: int = 4
|
||||
n_layers: int = 3
|
||||
mlp_ratio: int = 4
|
||||
max_text_tokens: int = 64
|
||||
num_bins: int = 201
|
||||
value_min: float = -1.0
|
||||
value_max: float = 0.0
|
||||
dropout: float = 0.0
|
||||
rms_eps: float = 1e-6
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.d_model % self.n_heads:
|
||||
raise ValueError("d_model must be divisible by n_heads")
|
||||
if self.image_size % self.patch_size:
|
||||
raise ValueError("image_size must be divisible by patch_size")
|
||||
if self.num_bins < 2:
|
||||
raise ValueError("num_bins must be at least two")
|
||||
if self.value_max <= self.value_min:
|
||||
raise ValueError("value_max must be greater than value_min")
|
||||
|
||||
@property
|
||||
def max_image_tokens(self) -> int:
|
||||
return (self.image_size // self.patch_size) ** 2
|
||||
|
||||
@property
|
||||
def max_sequence_tokens(self) -> int:
|
||||
# One learned pooling token plus image and discrete-state tokens.
|
||||
return 1 + self.max_image_tokens + self.max_text_tokens
|
||||
|
||||
|
||||
@dataclass
|
||||
class ValueOutput:
|
||||
logits: Tensor
|
||||
probabilities: Tensor
|
||||
expectation: Tensor
|
||||
hidden: Tensor
|
||||
loss: Optional[Tensor] = None
|
||||
target_bins: Optional[Tensor] = None
|
||||
|
||||
|
||||
class ValueBlock(nn.Module):
|
||||
def __init__(self, config: ValueConfig) -> None:
|
||||
super().__init__()
|
||||
self.n_heads = config.n_heads
|
||||
self.dropout = config.dropout
|
||||
self.norm1 = RMSNorm(config.d_model, config.rms_eps)
|
||||
self.attention = AttentionWeights(config.d_model)
|
||||
self.norm2 = RMSNorm(config.d_model, config.rms_eps)
|
||||
self.mlp = MLP(config.d_model, config.mlp_ratio, config.dropout)
|
||||
|
||||
def forward(self, x: Tensor, edges: Tensor) -> Tensor:
|
||||
normalized = self.norm1(x)
|
||||
weights = self.attention
|
||||
update = attention(
|
||||
weights.q(normalized),
|
||||
weights.k(normalized),
|
||||
weights.v(normalized),
|
||||
self.n_heads,
|
||||
edges,
|
||||
self.dropout,
|
||||
self.training,
|
||||
)
|
||||
x = x + weights.out(update)
|
||||
return x + self.mlp(self.norm2(x))
|
||||
|
||||
|
||||
class DistributionalValueModel(nn.Module):
|
||||
"""Predict ``p(V | image, discrete state)`` over 201 return bins.
|
||||
|
||||
The value network is deliberately a separate model: it shares no
|
||||
parameters with ``Pi06Mini``. By default its support is the paper's
|
||||
task-normalized interval ``[-1, 0]``.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[ValueConfig] = None) -> None:
|
||||
super().__init__()
|
||||
self.config = config or ValueConfig()
|
||||
c = self.config
|
||||
self.patch_encoder = nn.Conv2d(
|
||||
c.image_channels, c.d_model, c.patch_size, stride=c.patch_size
|
||||
)
|
||||
self.token_embedding = nn.Embedding(c.vocab_size, c.d_model)
|
||||
self.value_token = nn.Parameter(torch.randn(1, 1, c.d_model) * 0.01)
|
||||
self.position = nn.Parameter(
|
||||
torch.randn(1, c.max_sequence_tokens, c.d_model) * 0.01
|
||||
)
|
||||
self.layers = nn.ModuleList(ValueBlock(c) for _ in range(c.n_layers))
|
||||
self.norm = RMSNorm(c.d_model, c.rms_eps)
|
||||
self.value_head = nn.Linear(c.d_model, c.num_bins)
|
||||
self.register_buffer(
|
||||
"support",
|
||||
torch.linspace(c.value_min, c.value_max, c.num_bins),
|
||||
persistent=True,
|
||||
)
|
||||
|
||||
def values_to_bins(self, values: Tensor) -> Tensor:
|
||||
"""Nearest-bin discretization of normalized Monte-Carlo returns."""
|
||||
|
||||
values = torch.as_tensor(values, dtype=self.support.dtype, device=self.support.device)
|
||||
unit = (values.clamp(self.config.value_min, self.config.value_max) - self.config.value_min)
|
||||
unit = unit / (self.config.value_max - self.config.value_min)
|
||||
return torch.round(unit * (self.config.num_bins - 1)).long()
|
||||
|
||||
def bins_to_values(self, bins: Tensor) -> Tensor:
|
||||
bins = torch.as_tensor(bins, dtype=torch.long, device=self.support.device)
|
||||
if bins.numel() and (bins.min() < 0 or bins.max() >= self.config.num_bins):
|
||||
raise ValueError("value bins are outside the configured support")
|
||||
return self.support[bins]
|
||||
|
||||
def expectation_from_logits(self, logits: Tensor) -> Tensor:
|
||||
return (logits.softmax(dim=-1) * self.support.to(logits.dtype)).sum(dim=-1)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pixel_values: Tensor,
|
||||
input_ids: Tensor,
|
||||
attention_mask: Optional[Tensor] = None,
|
||||
target_values: Optional[Tensor] = None,
|
||||
target_bins: Optional[Tensor] = None,
|
||||
) -> ValueOutput:
|
||||
if target_values is not None and target_bins is not None:
|
||||
raise ValueError("provide target_values or target_bins, not both")
|
||||
patches = self.patch_encoder(pixel_values).flatten(2).transpose(1, 2)
|
||||
text = self.token_embedding(input_ids)
|
||||
batch = input_ids.shape[0]
|
||||
pooling = self.value_token.expand(batch, -1, -1)
|
||||
hidden = torch.cat((pooling, patches, text), dim=1)
|
||||
if hidden.shape[1] > self.config.max_sequence_tokens:
|
||||
raise ValueError("value-model sequence exceeds configured maximum length")
|
||||
hidden = hidden + self.position[:, : hidden.shape[1]]
|
||||
|
||||
valid = torch.ones(batch, hidden.shape[1], dtype=torch.bool, device=hidden.device)
|
||||
if attention_mask is not None:
|
||||
if attention_mask.shape != input_ids.shape:
|
||||
raise ValueError("attention_mask must have the same shape as input_ids")
|
||||
valid[:, 1 + patches.shape[1] :] = attention_mask.bool()
|
||||
edges = valid[:, None, :].expand(batch, hidden.shape[1], hidden.shape[1])
|
||||
identity = torch.eye(hidden.shape[1], dtype=torch.bool, device=hidden.device)[None]
|
||||
edges = torch.where(valid[:, :, None], edges, identity)
|
||||
for layer in self.layers:
|
||||
hidden = layer(hidden, edges)
|
||||
hidden = self.norm(hidden)
|
||||
|
||||
logits = self.value_head(hidden[:, 0])
|
||||
probabilities = logits.softmax(dim=-1)
|
||||
expectation = (probabilities * self.support.to(logits.dtype)).sum(dim=-1)
|
||||
|
||||
bins = None
|
||||
if target_values is not None:
|
||||
bins = self.values_to_bins(target_values).to(logits.device)
|
||||
elif target_bins is not None:
|
||||
bins = torch.as_tensor(target_bins, dtype=torch.long, device=logits.device)
|
||||
if bins.numel() and (bins.min() < 0 or bins.max() >= self.config.num_bins):
|
||||
raise ValueError("target_bins are outside the configured support")
|
||||
loss = None if bins is None else F.cross_entropy(logits, bins.reshape(-1))
|
||||
return ValueOutput(
|
||||
logits=logits,
|
||||
probabilities=probabilities,
|
||||
expectation=expectation,
|
||||
hidden=hidden,
|
||||
loss=loss,
|
||||
target_bins=bins,
|
||||
)
|
||||
|
||||
|
||||
# Short name convenient in configs and training scripts.
|
||||
ValueModel = DistributionalValueModel
|
||||
@@ -0,0 +1,81 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from pi06mini.experience import (
|
||||
build_synthetic_experience,
|
||||
discounted_returns,
|
||||
n_step_advantages,
|
||||
normalize_returns_to_unit_interval,
|
||||
quantile_advantage_labels,
|
||||
)
|
||||
|
||||
|
||||
def _geometry() -> tuple[torch.Tensor, torch.Tensor]:
|
||||
history = torch.stack((torch.linspace(-7.0, 0.0, 8), torch.zeros(8)), dim=-1)
|
||||
x = torch.linspace(1.0, 16.0, 16)
|
||||
future = torch.stack((x, 0.025 * x.square()), dim=-1)
|
||||
return history, future
|
||||
|
||||
|
||||
def test_synthetic_buffer_has_all_experience_sources_and_failure_geometry() -> None:
|
||||
history, future = _geometry()
|
||||
episodes = build_synthetic_experience(history, future, episode_prefix="pai")
|
||||
|
||||
assert [episode.source_type for episode in episodes] == [
|
||||
"demo_success",
|
||||
"autonomous_understeer",
|
||||
"autonomous_stop",
|
||||
"intervention_correction",
|
||||
]
|
||||
demo, understeer, stopped, correction = episodes
|
||||
torch.testing.assert_close(demo.future_xy, future)
|
||||
assert understeer.future_xy[-1, 1] < future[-1, 1]
|
||||
assert torch.equal(stopped.future_xy[-1], stopped.future_xy[-2])
|
||||
assert correction.intervention_mask.any()
|
||||
assert not correction.intervention_mask[0]
|
||||
torch.testing.assert_close(correction.future_xy[-1], future[-1])
|
||||
assert torch.all(demo.rewards == 0)
|
||||
assert float(understeer.rewards.min()) < 0
|
||||
assert float(stopped.rewards.min()) < 0
|
||||
|
||||
|
||||
def test_returns_are_normalized_to_closed_negative_unit_interval() -> None:
|
||||
rewards = torch.tensor([[0.0, 0.0, 0.0], [-0.2, -0.4, -1.0]])
|
||||
returns = discounted_returns(rewards, gamma=1.0)
|
||||
normalized = normalize_returns_to_unit_interval(returns)
|
||||
|
||||
assert normalized.shape == rewards.shape
|
||||
assert float(normalized.min()) == -1.0
|
||||
assert float(normalized.max()) == 0.0
|
||||
assert torch.all(normalized <= 0.0)
|
||||
assert torch.all(normalized >= -1.0)
|
||||
assert torch.equal(
|
||||
normalize_returns_to_unit_interval(torch.ones(3)), torch.zeros(3)
|
||||
)
|
||||
|
||||
|
||||
def test_n_step_advantage_bootstraps_and_stops_at_terminal() -> None:
|
||||
rewards = torch.tensor([1.0, 2.0, 3.0])
|
||||
values = torch.tensor([0.5, 1.0, 1.5, 8.0])
|
||||
dones = torch.tensor([False, True, False])
|
||||
advantages = n_step_advantages(
|
||||
rewards, values, n_step=3, gamma=0.5, dones=dones
|
||||
)
|
||||
|
||||
# t=0 sees r0 + gamma*r1, then the terminal blocks the large bootstrap.
|
||||
assert torch.isclose(advantages[0], torch.tensor(1.5))
|
||||
assert torch.isclose(advantages[1], torch.tensor(1.0))
|
||||
# t=2 reaches the explicit end and therefore has no bootstrap.
|
||||
assert torch.isclose(advantages[2], torch.tensor(1.5))
|
||||
|
||||
|
||||
def test_quantile_labels_force_interventions_positive() -> None:
|
||||
advantages = torch.tensor([-3.0, -2.0, -1.0, 2.0])
|
||||
interventions = torch.tensor([True, False, False, False])
|
||||
labels, threshold = quantile_advantage_labels(
|
||||
advantages, quantile=0.75, intervention_mask=interventions
|
||||
)
|
||||
|
||||
assert labels.tolist() == [True, False, False, True]
|
||||
assert float(threshold) > -1.0
|
||||
@@ -0,0 +1,80 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from pi06mini.data import Pi06ExperienceDataset, load_manifest, materialize_subset
|
||||
|
||||
|
||||
def _make_source_project(root: Path) -> Path:
|
||||
processed = root / "data" / "processed"
|
||||
(processed / "images").mkdir(parents=True)
|
||||
(processed / "trajectories").mkdir(parents=True)
|
||||
pixels = np.zeros((40, 64, 3), dtype=np.uint8)
|
||||
pixels[..., 0] = 80
|
||||
pixels[..., 1] = np.arange(64, dtype=np.uint8)[None, :]
|
||||
Image.fromarray(pixels).save(processed / "images" / "sample.jpg")
|
||||
|
||||
hx = np.linspace(-7.0, 0.0, 8, dtype=np.float32)
|
||||
history = np.stack((hx, 0.01 * hx, np.zeros(8)), axis=-1)
|
||||
fx = np.linspace(1.0, 24.0, 16, dtype=np.float32)
|
||||
future = np.stack((fx, 0.015 * fx**2, np.zeros(16)), axis=-1)
|
||||
np.savez_compressed(
|
||||
processed / "trajectories" / "sample.npz",
|
||||
history_xyz=history,
|
||||
future_xyz=future,
|
||||
)
|
||||
manifest = processed / "overfit_manifest.jsonl"
|
||||
manifest.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"clip_id": "clip-a",
|
||||
"t0_us": 123,
|
||||
"reasoning": "Steer left.",
|
||||
"image_path": "data/processed/images/sample.jpg",
|
||||
"trajectory_path": "data/processed/trajectories/sample.npz",
|
||||
}
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return manifest
|
||||
|
||||
|
||||
def test_experience_dataset_contract_and_intervention_label(tmp_path: Path) -> None:
|
||||
source = _make_source_project(tmp_path / "source")
|
||||
manifest = materialize_subset(source, tmp_path / "target")
|
||||
assert load_manifest(manifest)[0].image_path.is_file()
|
||||
|
||||
dataset = Pi06ExperienceDataset(manifest, image_size=32)
|
||||
assert len(dataset) == 4
|
||||
assert {item.source_type for item in dataset.episodes} == {
|
||||
"demo_success",
|
||||
"autonomous_understeer",
|
||||
"autonomous_stop",
|
||||
"intervention_correction",
|
||||
}
|
||||
assert float(dataset.normalized_returns.min()) == -1.0
|
||||
assert float(dataset.normalized_returns.max()) == 0.0
|
||||
|
||||
items = [dataset[index] for index in range(len(dataset))]
|
||||
intervention = next(item for item in items if item["is_intervention"])
|
||||
assert bool(intervention["positive_condition"])
|
||||
assert torch.all(intervention["positive_steps"][intervention["intervention_mask"]])
|
||||
assert intervention["image"].shape == (3, 32, 32)
|
||||
assert intervention["state"].shape == (16,)
|
||||
assert intervention["actions"].shape == (16, 2)
|
||||
assert intervention["returns"].shape == (16,)
|
||||
assert intervention["return_target"].shape == ()
|
||||
assert intervention["value_target"].shape == ()
|
||||
assert intervention["advantages"].shape == (16,)
|
||||
assert intervention["advantage_label"] == "positive"
|
||||
assert intervention["reasoning"] == "Steer left."
|
||||
decoded = dataset.action_normalizer.decode(
|
||||
intervention["actions"], intervention["history_xy"]
|
||||
)
|
||||
torch.testing.assert_close(decoded, intervention["future_xy"], atol=1e-5, rtol=1e-5)
|
||||
@@ -0,0 +1,123 @@
|
||||
import torch
|
||||
|
||||
from pi06mini.model import Pi06Config, Pi06Mini
|
||||
from pi06mini.tokenizer import MiniTokenizer
|
||||
|
||||
|
||||
def tiny_model() -> Pi06Mini:
|
||||
tokenizer = MiniTokenizer()
|
||||
return Pi06Mini(
|
||||
Pi06Config(
|
||||
vocab_size=tokenizer.vocab_size,
|
||||
image_size=16,
|
||||
patch_size=16,
|
||||
d_model=32,
|
||||
n_heads=4,
|
||||
n_layers=1,
|
||||
max_text_tokens=32,
|
||||
max_action_tokens=4,
|
||||
action_start_token_id=tokenizer.ACTION_START,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_flow_reads_subtask_and_advantage_but_not_fast_targets() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
sequence = tokenizer.build_policy_sequence(
|
||||
torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
|
||||
)
|
||||
model = tiny_model()
|
||||
input_ids = sequence.input_ids[None]
|
||||
derived = model._derive_action_prefix_mask(input_ids)
|
||||
torch.testing.assert_close(derived[0], sequence.action_prefix_mask)
|
||||
|
||||
text_offset = 1
|
||||
prefix, flow = model._masks(
|
||||
batch=1,
|
||||
prefix_length=1 + input_ids.shape[1],
|
||||
action_length=2,
|
||||
text_offset=text_offset,
|
||||
attention_mask=torch.ones_like(input_ids, dtype=torch.bool),
|
||||
action_prefix_mask=derived,
|
||||
generated_token_start=sequence.generated_token_start,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
generated_start = text_offset + sequence.generated_token_start
|
||||
assert prefix[0, 0, :generated_start].all()
|
||||
assert not prefix[0, 0, generated_start:].any()
|
||||
# Flow sees image plus all text strictly before ACTION_START.
|
||||
readable_prefix = text_offset + sequence.action_start_index
|
||||
assert flow[0, :, :readable_prefix].all()
|
||||
assert not flow[0, :, readable_prefix : 1 + input_ids.shape[1]].any()
|
||||
assert flow[0, :, 1 + input_ids.shape[1] :].all()
|
||||
|
||||
|
||||
def test_flow_loss_is_knowledge_insulated_from_every_vlm_parameter() -> None:
|
||||
torch.manual_seed(4)
|
||||
tokenizer = MiniTokenizer()
|
||||
sequence = tokenizer.build_policy_sequence(
|
||||
torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
|
||||
)
|
||||
model = tiny_model()
|
||||
output = model(
|
||||
pixel_values=torch.randn(1, 3, 16, 16),
|
||||
input_ids=sequence.input_ids[None],
|
||||
noisy_actions=torch.randn(1, 2, 2),
|
||||
timesteps=torch.rand(1),
|
||||
action_prefix_mask=sequence.action_prefix_mask[None],
|
||||
generated_token_start=sequence.generated_token_start,
|
||||
target_velocity=torch.randn(1, 2, 2),
|
||||
)
|
||||
assert output.flow_loss is not None
|
||||
output.flow_loss.backward()
|
||||
|
||||
vlm_gradients = [parameter.grad for parameter in model.vlm_parameters()]
|
||||
assert vlm_gradients
|
||||
assert all(gradient is None or torch.count_nonzero(gradient) == 0 for gradient in vlm_gradients)
|
||||
action_gradients = [parameter.grad for parameter in model.action_parameters()]
|
||||
assert any(gradient is not None and torch.count_nonzero(gradient) > 0 for gradient in action_gradients)
|
||||
|
||||
|
||||
def test_knowledge_insulation_stops_gradients_not_conditioning() -> None:
|
||||
torch.manual_seed(8)
|
||||
tokenizer = MiniTokenizer()
|
||||
model = tiny_model().eval()
|
||||
positive = tokenizer.build_policy_sequence(
|
||||
torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "positive"
|
||||
)
|
||||
negative = tokenizer.build_policy_sequence(
|
||||
torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "negative"
|
||||
)
|
||||
image = torch.randn(1, 3, 16, 16)
|
||||
noise = torch.randn(1, 2, 2)
|
||||
timestep = torch.full((1,), 0.5)
|
||||
with torch.no_grad():
|
||||
pos = model(
|
||||
image, positive.input_ids[None], noise, timestep,
|
||||
generated_token_start=positive.generated_token_start,
|
||||
).vector_field
|
||||
neg = model(
|
||||
image, negative.input_ids[None], noise, timestep,
|
||||
generated_token_start=negative.generated_token_start,
|
||||
).vector_field
|
||||
assert not torch.allclose(pos, neg)
|
||||
|
||||
|
||||
def test_joint_loss_still_trains_vlm_through_token_objective() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
sequence = tokenizer.build_policy_sequence(
|
||||
torch.zeros(2, 2), "stop", torch.zeros(2, 2), "negative"
|
||||
)
|
||||
model = tiny_model()
|
||||
output = model(
|
||||
torch.randn(1, 3, 16, 16),
|
||||
sequence.input_ids[None],
|
||||
torch.randn(1, 2, 2),
|
||||
torch.rand(1),
|
||||
generated_token_start=sequence.generated_token_start,
|
||||
text_labels=sequence.labels[None],
|
||||
target_velocity=torch.randn(1, 2, 2),
|
||||
)
|
||||
output.loss.backward()
|
||||
assert model.token_embedding.weight.grad is not None
|
||||
assert torch.count_nonzero(model.token_embedding.weight.grad) > 0
|
||||
@@ -0,0 +1,49 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from pi06mini.recap import (
|
||||
apply_condition_dropout,
|
||||
classifier_free_guidance,
|
||||
condition_dropout_mask,
|
||||
)
|
||||
|
||||
|
||||
def test_default_condition_dropout_is_statistically_thirty_percent() -> None:
|
||||
generator = torch.Generator().manual_seed(11)
|
||||
mask = condition_dropout_mask(20_000, generator=generator)
|
||||
|
||||
assert mask.dtype == torch.bool
|
||||
assert abs(float(mask.float().mean()) - 0.3) < 0.015
|
||||
|
||||
|
||||
def test_condition_dropout_replaces_whole_examples() -> None:
|
||||
conditions = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
|
||||
dropped, mask = apply_condition_dropout(
|
||||
conditions,
|
||||
null_condition=-1.0,
|
||||
mask=torch.tensor([False, True, False]),
|
||||
)
|
||||
|
||||
assert mask.tolist() == [False, True, False]
|
||||
torch.testing.assert_close(dropped[0], conditions[0])
|
||||
torch.testing.assert_close(dropped[1], torch.tensor([-1.0, -1.0]))
|
||||
torch.testing.assert_close(dropped[2], conditions[2])
|
||||
|
||||
|
||||
def test_cfg_interpolates_and_extrapolates_flow_predictions() -> None:
|
||||
conditional = torch.tensor([[2.0, 4.0]])
|
||||
unconditional = torch.tensor([[0.0, 2.0]])
|
||||
|
||||
torch.testing.assert_close(
|
||||
classifier_free_guidance(conditional, unconditional, guidance_scale=0.0),
|
||||
unconditional,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
classifier_free_guidance(conditional, unconditional, guidance_scale=1.0),
|
||||
conditional,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
classifier_free_guidance(conditional, unconditional, guidance_scale=1.5),
|
||||
torch.tensor([[3.0, 5.0]]),
|
||||
)
|
||||
@@ -0,0 +1,65 @@
|
||||
import torch
|
||||
|
||||
from pi06mini.tokenizer import MiniTokenizer
|
||||
|
||||
|
||||
def test_advantage_condition_is_between_subtask_and_actions() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
history = torch.tensor([[0.0, 0.0], [1.0, 0.2]])
|
||||
actions = torch.tensor([[0.1, 0.0], [0.2, -0.1]])
|
||||
|
||||
positive = tokenizer.build_policy_sequence(
|
||||
history, "steer_left", actions, "positive"
|
||||
)
|
||||
assert positive.advantage_index == positive.subtask_index + 1
|
||||
assert positive.action_start_index == positive.advantage_index + 1
|
||||
assert positive.input_ids[positive.advantage_index] == tokenizer.ADVANTAGE_POSITIVE
|
||||
assert positive.input_ids[positive.action_start_index] == tokenizer.ACTION_START
|
||||
assert positive.action_prefix_mask[: positive.action_start_index].all()
|
||||
assert not positive.action_prefix_mask[positive.action_start_index :].any()
|
||||
|
||||
negative = tokenizer.build_policy_sequence(
|
||||
history, "steer_left", actions, "negative"
|
||||
)
|
||||
assert negative.input_ids[negative.advantage_index] == tokenizer.ADVANTAGE_NEGATIVE
|
||||
|
||||
|
||||
def test_omitted_advantage_is_absent_not_a_third_condition_token() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
sequence = tokenizer.build_policy_sequence(
|
||||
torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "omitted"
|
||||
)
|
||||
|
||||
assert sequence.advantage_index is None
|
||||
assert sequence.action_start_index == sequence.subtask_index + 1
|
||||
assert sequence.input_ids[sequence.action_start_index] == tokenizer.ACTION_START
|
||||
assert tokenizer.ADVANTAGE_POSITIVE not in sequence.input_ids
|
||||
assert tokenizer.ADVANTAGE_NEGATIVE not in sequence.input_ids
|
||||
|
||||
|
||||
def test_labels_predict_subtask_and_fast_but_not_external_advantage() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
positive = tokenizer.build_policy_sequence(
|
||||
torch.zeros(2, 2), "stop", torch.zeros(2, 2), True
|
||||
)
|
||||
assert positive.labels[positive.subtask_index - 1] == positive.input_ids[positive.subtask_index]
|
||||
assert positive.labels[positive.subtask_index] == -100
|
||||
assert positive.labels[positive.advantage_index] == tokenizer.ACTION_START
|
||||
|
||||
omitted = tokenizer.build_policy_sequence(
|
||||
torch.zeros(2, 2), "stop", torch.zeros(2, 2), None
|
||||
)
|
||||
assert omitted.labels[omitted.subtask_index] == tokenizer.ACTION_START
|
||||
|
||||
|
||||
def test_fast_lite_dct_round_trip_has_bounded_error() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
actions = torch.tensor(
|
||||
[[0.0, 0.0], [0.2, -0.1], [0.5, -0.2], [0.8, -0.1]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
exact = tokenizer.idct_actions(tokenizer.dct_actions(actions))
|
||||
torch.testing.assert_close(exact, actions, atol=1e-5, rtol=1e-5)
|
||||
|
||||
decoded = tokenizer.decode_actions(tokenizer.encode_actions(actions))
|
||||
assert torch.sqrt(torch.mean((decoded - actions) ** 2)) < 0.3
|
||||
@@ -0,0 +1,73 @@
|
||||
import torch
|
||||
|
||||
from pi06mini.model import Pi06Config, Pi06Mini
|
||||
from pi06mini.tokenizer import MiniTokenizer
|
||||
from pi06mini.value import DistributionalValueModel, ValueConfig
|
||||
|
||||
|
||||
def tiny_value_model() -> DistributionalValueModel:
|
||||
tokenizer = MiniTokenizer()
|
||||
return DistributionalValueModel(
|
||||
ValueConfig(
|
||||
vocab_size=tokenizer.vocab_size,
|
||||
image_size=16,
|
||||
patch_size=16,
|
||||
d_model=32,
|
||||
n_heads=4,
|
||||
n_layers=1,
|
||||
max_text_tokens=16,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_distributional_value_has_201_bins_and_expected_value() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
model = tiny_value_model()
|
||||
state = tokenizer.build_value_input(torch.zeros(2, 2))[None].repeat(2, 1)
|
||||
output = model(
|
||||
torch.randn(2, 3, 16, 16),
|
||||
state,
|
||||
target_values=torch.tensor([-1.0, 0.0]),
|
||||
)
|
||||
|
||||
assert output.logits.shape == (2, 201)
|
||||
assert output.probabilities.shape == (2, 201)
|
||||
torch.testing.assert_close(output.probabilities.sum(-1), torch.ones(2))
|
||||
assert torch.all((-1.0 <= output.expectation) & (output.expectation <= 0.0))
|
||||
assert output.target_bins.tolist() == [0, 200]
|
||||
assert output.loss is not None and output.loss.ndim == 0
|
||||
|
||||
|
||||
def test_value_discretization_and_expectation_use_configured_support() -> None:
|
||||
model = tiny_value_model()
|
||||
bins = model.values_to_bins(torch.tensor([-2.0, -0.5, 1.0]))
|
||||
assert bins.tolist() == [0, 100, 200]
|
||||
torch.testing.assert_close(
|
||||
model.bins_to_values(torch.tensor([0, 100, 200])),
|
||||
torch.tensor([-1.0, -0.5, 0.0]),
|
||||
)
|
||||
logits = torch.full((3, 201), -100.0)
|
||||
logits[0, 0] = logits[1, 100] = logits[2, 200] = 100.0
|
||||
torch.testing.assert_close(
|
||||
model.expectation_from_logits(logits),
|
||||
torch.tensor([-1.0, -0.5, 0.0]),
|
||||
)
|
||||
|
||||
|
||||
def test_value_model_and_policy_share_no_parameters() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
policy = Pi06Mini(
|
||||
Pi06Config(
|
||||
vocab_size=tokenizer.vocab_size,
|
||||
image_size=16,
|
||||
patch_size=16,
|
||||
d_model=32,
|
||||
n_heads=4,
|
||||
n_layers=1,
|
||||
action_start_token_id=tokenizer.ACTION_START,
|
||||
)
|
||||
)
|
||||
value = tiny_value_model()
|
||||
policy_ids = {id(parameter) for parameter in policy.parameters()}
|
||||
value_ids = {id(parameter) for parameter in value.parameters()}
|
||||
assert policy_ids.isdisjoint(value_ids)
|
||||
Reference in New Issue
Block a user