commit 53ef020b46d1c449b75df27f4b71df586699767c Author: seeker Date: Sat Jul 11 22:35:54 2026 +0800 Initial pi0.6 minimal validation diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ff92181 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.ruff_cache/ +*.egg-info/ +.venv/ +data/ +outputs/**/*.pt +.DS_Store diff --git a/README.md b/README.md new file mode 100644 index 0000000..d5ff69b --- /dev/null +++ b/README.md @@ -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) diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..fcd8290 --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -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): . +- Physical Intelligence, **π0.6 model card**: + . +- 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. diff --git a/configs/overfit.yaml b/configs/overfit.yaml new file mode 100644 index 0000000..92a0a1f --- /dev/null +++ b/configs/overfit.yaml @@ -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 diff --git a/outputs/overfit/evaluation.json b/outputs/overfit/evaluation.json new file mode 100644 index 0000000..7d7880f --- /dev/null +++ b/outputs/overfit/evaluation.json @@ -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 + } + ] +} \ No newline at end of file diff --git a/outputs/overfit/recap_trajectory_comparison.png b/outputs/overfit/recap_trajectory_comparison.png new file mode 100644 index 0000000..3b81dd6 Binary files /dev/null and b/outputs/overfit/recap_trajectory_comparison.png differ diff --git a/outputs/overfit/summary.json b/outputs/overfit/summary.json new file mode 100644 index 0000000..7cdb9ed --- /dev/null +++ b/outputs/overfit/summary.json @@ -0,0 +1,12 @@ +{ + "device": "mps", + "experience_samples": 4, + "policy_parameters": 817058, + "value_parameters": 171977, + "sources": [ + "demo_success", + "autonomous_understeer", + "autonomous_stop", + "intervention_correction" + ] +} \ No newline at end of file diff --git a/outputs/overfit/training_history.json b/outputs/overfit/training_history.json new file mode 100644 index 0000000..7cdf30f --- /dev/null +++ b/outputs/overfit/training_history.json @@ -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 + } +] \ No newline at end of file diff --git a/outputs/overfit/value_comparison.png b/outputs/overfit/value_comparison.png new file mode 100644 index 0000000..5b4143a Binary files /dev/null and b/outputs/overfit/value_comparison.png differ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..9643377 --- /dev/null +++ b/pyproject.toml @@ -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" diff --git a/src/pi06mini/__init__.py b/src/pi06mini/__init__.py new file mode 100644 index 0000000..8db6dd2 --- /dev/null +++ b/src/pi06mini/__init__.py @@ -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", +] diff --git a/src/pi06mini/data.py b/src/pi06mini/data.py new file mode 100644 index 0000000..ccb1bdc --- /dev/null +++ b/src/pi06mini/data.py @@ -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}", + } diff --git a/src/pi06mini/experience.py b/src/pi06mini/experience.py new file mode 100644 index 0000000..10f66e9 --- /dev/null +++ b/src/pi06mini/experience.py @@ -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]) diff --git a/src/pi06mini/flow.py b/src/pi06mini/flow.py new file mode 100644 index 0000000..f23ba89 --- /dev/null +++ b/src/pi06mini/flow.py @@ -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 diff --git a/src/pi06mini/infer.py b/src/pi06mini/infer.py new file mode 100644 index 0000000..a46a3fc --- /dev/null +++ b/src/pi06mini/infer.py @@ -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() diff --git a/src/pi06mini/model.py b/src/pi06mini/model.py new file mode 100644 index 0000000..712b7d3 --- /dev/null +++ b/src/pi06mini/model.py @@ -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 diff --git a/src/pi06mini/recap.py b/src/pi06mini/recap.py new file mode 100644 index 0000000..176d219 --- /dev/null +++ b/src/pi06mini/recap.py @@ -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 + ) diff --git a/src/pi06mini/runtime.py b/src/pi06mini/runtime.py new file mode 100644 index 0000000..118fb34 --- /dev/null +++ b/src/pi06mini/runtime.py @@ -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") diff --git a/src/pi06mini/tokenizer.py b/src/pi06mini/tokenizer.py new file mode 100644 index 0000000..fec7dbc --- /dev/null +++ b/src/pi06mini/tokenizer.py @@ -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, + ) diff --git a/src/pi06mini/train.py b/src/pi06mini/train.py new file mode 100644 index 0000000..f9d033d --- /dev/null +++ b/src/pi06mini/train.py @@ -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() diff --git a/src/pi06mini/value.py b/src/pi06mini/value.py new file mode 100644 index 0000000..fc0c0ce --- /dev/null +++ b/src/pi06mini/value.py @@ -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 diff --git a/tests/test_experience.py b/tests/test_experience.py new file mode 100644 index 0000000..39ac302 --- /dev/null +++ b/tests/test_experience.py @@ -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 diff --git a/tests/test_experience_data.py b/tests/test_experience_data.py new file mode 100644 index 0000000..a3d0a20 --- /dev/null +++ b/tests/test_experience_data.py @@ -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) diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..389e03f --- /dev/null +++ b/tests/test_model.py @@ -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 diff --git a/tests/test_recap.py b/tests/test_recap.py new file mode 100644 index 0000000..8defff4 --- /dev/null +++ b/tests/test_recap.py @@ -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]]), + ) diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py new file mode 100644 index 0000000..8bd686c --- /dev/null +++ b/tests/test_tokenizer.py @@ -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 diff --git a/tests/test_value.py b/tests/test_value.py new file mode 100644 index 0000000..13a2657 --- /dev/null +++ b/tests/test_value.py @@ -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)