commit b4c8db846b1a9fc52e3a4d9593c5e67b10e1353d Author: seeker Date: Sun Jul 12 00:52:00 2026 +0800 Initial pi0.7 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..046811e --- /dev/null +++ b/README.md @@ -0,0 +1,246 @@ +# π0.7 Minimal Validation + +这是一个面向**结构验证和小样本过拟合**的 π0.7-style 教学项目。它在前面的 +π0.6/RECAP 最小实验之上，重点验证一条可解释的 steerable policy 链路：利用 +短时历史形成 memory，在高层预测视觉 subgoal，通过可独立缺省的 policy +metadata 控制动作风格，并在存在控制延迟时用 RTC 对齐动作 chunk。 + +本项目不是 Physical Intelligence 官方实现，也不能用于真实车辆或机器人控制。 + +## 公开资料边界 + +截至本项目实现时，作者没有在 Physical Intelligence 官方网站、公开论文索引或 +`openpi` 仓库中找到可核验的 π0.7 原始论文、官方训练源码、checkpoint 或完整 +model card。因此这里不能声称“逐层复现论文”，也不会虚构官方模块尺寸、损失、 +数据配方或实验数字。 + +本仓库所说的 **paper-faithful toy** 有严格边界：若 π0.7 的公开描述强调 memory、 +goal generation、high-level control、steerable metadata、CFG 和 real-time +chunking，本项目保留这些模块的输入输出关系与必要消融；所有未公开细节都使用 +最小、明确标注且有测试的本地替代。`TinyMEMEncoder`、`TinyGoalGenerator`、 +`HighLevelPolicy` 和 `Pi07Mini` 都是本仓库命名，不代表官方内部类名。 + +若后续公开原始论文或源码，应重新审计本项目，而不是把当前实现当成事实来源。 + +## 模型阶段与输入输出 + +| 阶段 | 输入 | 模块 | 输出 | 训练或验证目标 | +|---|---|---|---|---| +| 1. 历史构造 | 当前图像、过去 XY；不接受 future | deterministic history proxy | `6` 帧 oldest-to-newest 图像、历史状态、共享当前状态的 ambiguous history pair | 检查无未来轨迹泄漏 | +| 2. Temporal memory | 历史图像与历史状态 | `TinyMEMEncoder` | memory tokens / pooled memory embedding | 区分当前状态相同、过去路径不同的样本 | +| 3. Goal generation | 当前 observation、memory、任务文本 | `TinyGoalGenerator` | endpoint heatmap / subgoal XY | 只监督 future endpoint，不把完整真值轨迹输入 policy | +| 4. High-level policy | observation、memory、任务与可见 metadata | `HighLevelPolicy` | high-level context、goal/metadata logits | goal loss 与 high-level classification loss | +| 5. Steerable flow policy | high-level context、预测 goal、metadata、noisy action、flow time | `Pi07Mini` low-level action expert | 连续 action vector field `[16,2]` | flow-matching loss | +| 6. RTC alignment | action chunk、随机 delay `d∈[0,12]`、已执行 action prefix | real-time chunking target alignment | 从 `action[d]` 开始的 target 与有效位置 mask | 仅在有效位置计算 masked flow loss | +| 7. CFG inference | 同一 observation 的 conditional 与 metadata-omitted 两次前向 | classifier-free guidance | guided action chunk | 与普通 conditional、unconditional 和真值同图比较 | +| 8. 轨迹恢复 | normalized action deltas、历史终点、训练统计 | denormalize + cumulative sum | 未来 XY 轨迹 | 轨迹 MSE 与可视化，不等同安全指标 | + +核心信息流： + +```text +history frames + history states + ↓ + TinyMEMEncoder ───────────────┐ + │ +current image + task ────────────────┼→ TinyGoalGenerator → predicted subgoal + │ +optional speed/quality/mistake/control metadata + ↓ + HighLevelPolicy + ↓ +noisy action + flow time + RTC prefix → Pi07Mini action expert + ↓ + continuous action chunk +``` + +## 重点结构 + +### 1. MEM：短时历史而非单帧 + +PAI 的极小 manifest 只有一张相机图像，本项目依据已观测的过去 XY 对该图像做 +确定性的微小平移和曝光变化，构造短时视觉 history proxy。它不是实际多帧视频， +只能验证 temporal memory 的接口。 + +`ambiguous_history_pair` 刻意构造两个“当前状态相同、历史路径不同”的序列。 +完整模型应比 `no_history` 消融更能区分它们。这里的 MEM 仅表示本仓库的 temporal +memory encoder；在官方资料缺失时，不推断该缩写的官方全称。 + +### 2. Goal generator：只用 endpoint 监督 + +训练目标是 future trajectory 的最后一个 XY endpoint，经高斯渲染得到 heatmap。 +API 只接受 endpoint，不接受整条 future trajectory，从而避免把低层动作答案作为 +视觉 subgoal 条件泄漏给 action expert。推理和最终评估必须使用**模型预测的 goal**； +真值 heatmap 只可用于训练监督和 oracle 诊断。 + +### 3. High-level 与 steerable metadata + +最小数据由一条 expert demonstration 确定性构造五种 policy mode： + +| Mode | speed | quality | mistake | control | +|---|---|---|---|---| +| expert | normal | expert | none | human | +| fast | fast | good | none | autonomous | +| slow | slow | good | none | autonomous | +| understeer | normal | poor | understeer | autonomous | +| stop | slow | poor | stop | autonomous | + +训练时 `speed`、`quality`、`mistake`、`control` 独立做 metadata dropout，任务文本 +始终保留。这既检验属性是否能单独控制行为，也提供 CFG 所需的 omitted condition。 +这些合成标签不是人类偏好数据，也不能证明真实世界可控性。 + +### 4. RTC：控制延迟下的 action chunk 对齐 + +若推理耗时等价于已经执行 `d` 个 action，模型不应再次预测旧 chunk 的开头： + +```text +original target: a0 a1 ... a15 +delay d=4: a4 a5 ... a15 PAD PAD PAD PAD +valid mask: 1 1 ... 1 0 0 0 0 +``` + +默认随机采样 `d=0..12`，返回已执行 prefix，并对 padding tail 屏蔽损失。RTC 在 +这里表示训练/评估时的离线对齐机制，不包含实际控制器、异步推理线程或硬实时保证。 + +### 5. CFG + +推理时在相同 observation、goal、memory 和初始噪声下计算 metadata conditional +与 metadata-omitted 预测： + +```text +v_cfg = v_uncond + beta * (v_cond - v_uncond) +``` + +默认 `beta=2.0`。必须同时报告 `beta=0,1,2`，因为高 guidance 可能放大过拟合、 +产生分布外动作，不能只展示最好看的曲线。 + +## 数据与防泄漏约束 + +- history proxy 和 ambiguous pair 只依赖当前图像与过去状态； +- subgoal condition 只使用 endpoint target，不能编码完整 action chunk； +- flow expert 训练时不能读取 teacher-forced 真值动作 token； +- 所有模型对比使用同一 initial noise、归一化统计和 RTC delay； +- 数据增强的随机数由 sample index 和 seed 决定，保证复现实验； +- 仓库不包含或重新分发原始 PAI 数据。 + +## 安装 + +```bash +cd pi07-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`。 + +## 测试 + +```bash +pytest -q +``` + +测试至少应覆盖：history/goal 无 future 泄漏、ambiguous history、metadata 独立 +dropout、五类策略构造、action normalization、flow matching、RTC shift/mask、 +MEM/goal/high-level/policy 前向形状、CFG 公式和 checkpoint reload。 + +## 三阶段过拟合训练 + +```bash +pi07mini-train --config configs/overfit.yaml +``` + +或： + +```bash +python -m pi07mini.train --config configs/overfit.yaml +``` + +| 阶段 | 默认步数 | 更新内容 | 目的 | +|---|---:|---|---| +| goal | 300 | MEM + goal generator | 先证明 history-conditioned endpoint heatmap 可被记忆 | +| high-level | 200 | MEM + high-level policy | 记住五类 metadata/goal 高层条件 | +| policy | 900 | 完整 `Pi07Mini` | 联合训练 metadata dropout、flow action 与 RTC masked target | + +默认预算面向极小样本过拟合，不代表官方训练 schedule。训练完成后应保存实际 +训练会生成 loss history、checkpoint、evaluation JSON 和消融效果图。 + +## 当前最小实验结果 + +固定相同 flow 初始噪声时，高层策略预测 `steer_left`。轨迹 MSE 如下： + +| 条件 | MSE | +|---|---:| +| Full context + oracle endpoint goal | 0.470 | +| Full context + generated goal | 0.491 | +| No metadata | 2.458 | +| No goal | 4.043 | +| No history | 2.159 | +| Metadata CFG beta=2 | 10.576 | + +完整上下文明显优于移除 metadata、goal 或 history 的消融。TinyGoalGenerator 的像素级结果仍较粗糙；生成 goal 在当前 policy 上可提供有效条件，但不能视为真实未来图像生成能力。较强 CFG 在此极小数据实验中造成退化。 + +## 推理、画图与必要消融 + +```bash +pi07mini-infer --checkpoint outputs/overfit/checkpoint.pt +``` + +或： + +```bash +python -m pi07mini.infer --checkpoint outputs/overfit/checkpoint.pt +``` + +最终展示不能只有 full model。相同样本、noise seed 和 RTC delay 下至少比较： + +| 实验 | 移除内容 | 要回答的问题 | +|---|---|---| +| `full` | 无 | 完整链路能否过拟合并重建对应 mode 的轨迹？ | +| `no_metadata` | speed/quality/mistake/control | policy mode 差异是否真的来自 steerable condition？ | +| `no_goal` | predicted goal | endpoint condition 是否改善几何方向和终点？ | +| `no_history` | MEM history，只保留当前 observation | temporal history 是否解决相同当前状态的歧义？ | + +还必须提供： + +- predicted goal heatmap 与 ground-truth endpoint heatmap； +- expert/fast/slow/understeer/stop 的真值与预测轨迹； +- CFG `beta=0,1,2`； +- RTC delay `0,4,8,12` 的 masked trajectory error； +- 每个 ablation 的轨迹 MSE，以及 goal endpoint error； +- full/no-history 在 ambiguous history pair 上的差异指标。 + +小样本过拟合中，某个消融偶然更好并不证明模块无效；应展示全部固定设置结果， +不调换 noise seed 或只挑选最佳样本。 + +## 安全限制 + +- 严禁连接车辆、机械臂、移动机器人或任何会造成物理影响的执行器。 +- 合成 fast/understeer/stop 轨迹只用于离线可控性探针，不是安全驾驶示范。 +- endpoint 命中、轨迹 MSE 下降、loss 收敛和 CFG 改善都不代表闭环稳定性。 +- history proxy 不是视频；RTC target shift 也不是硬实时系统验证。 +- metadata 可控性可能放大危险属性，尤其是 `fast`、`poor` 或高 CFG beta。 +- 真实部署至少需要真实时序传感器、闭环 rollout、OOD 与延迟压力测试、动力学 + 约束、碰撞检查、冗余控制器、故障降级、急停和人工监督。 +- PyTorch checkpoint 使用 pickle 机制，只加载可信来源的文件。 + +## 能得出的结论 + +成功实验最多能说明：本仓库定义的 memory、endpoint goal、high-level metadata、 +conditional flow、CFG 与 RTC 对齐可以在几个合成 PAI trajectory mode 上形成一条 +可训练、可消融的最小链路。它不能证明与未公开 π0.7 的结构或数值等价，不能 +证明 emergent capability，也不能支持任何真实世界安全或泛化主张。 + +## 可核验背景资料 + +- [Physical Intelligence](https://www.pi.website/) +- [Physical-Intelligence/openpi](https://github.com/Physical-Intelligence/openpi) + +在原始 π0.7 论文或官方源码可核验前，本 README 不提供推测性的论文链接、arXiv +编号或官方结果表。 diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..9689d12 --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,31 @@ +# Third-party notices + +This repository is an independent educational implementation. It does not +contain code, weights, datasets, or other assets copied from Physical +Intelligence. + +At the time this project was prepared, no original pi0.7 paper, official +pi0.7 training implementation, checkpoint, or complete technical model card +was available in the public Physical Intelligence/openpi materials inspected +by the authors. Consequently, names such as `TinyMEMEncoder`, +`TinyGoalGenerator`, `HighLevelPolicy`, and `Pi07Mini` describe this +repository's own teaching implementation. They must not be interpreted as +official class names or proof of the unpublished system's internals. + +The project uses the earlier public pi-family literature and public openpi +repository only as background for vision-language-action policies, action +chunking, and flow matching. No source file from openpi is vendored here. + +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 and terms of any +external images, trajectories, datasets, or checkpoints they supply. The +default configuration points to a local manifest and this repository does not +redistribute the source PAI data. diff --git a/configs/overfit.yaml b/configs/overfit.yaml new file mode 100644 index 0000000..5203fac --- /dev/null +++ b/configs/overfit.yaml @@ -0,0 +1,98 @@ +seed: 7 +device: auto +output_dir: outputs/overfit + +data: + source_manifest: ../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl + max_samples: 1 + augmentation_factor: 1 + image_size: 64 + history_frames: 6 + action_horizon: 16 + policy_variants: + - expert + - fast + - slow + - understeer + - stop + +model: + patch_size: 16 + d_model: 96 + mem_d_model: 64 + goal_d_model: 64 + high_level_d_model: 64 + n_heads: 4 + n_layers: 3 + mlp_ratio: 3 + dropout: 0.0 + max_text_tokens: 96 + action_dim: 2 + action_horizon: 16 + goal_heatmap_size: 32 + +memory: + enabled: true + history_frames: 6 + ambiguous_pair_weight: 0.25 + +goal: + enabled: true + heatmap_size: 32 + heatmap_sigma: 1.5 + xy_extent: [-8.0, 64.0, -16.0, 16.0] + endpoint_only_target: true + +high_level: + enabled: true + predict_goal: true + predict_metadata: true + metadata_fields: + - speed + - quality + - mistake + - control + +steering: + metadata_dropout: 0.25 + preserve_task_prompt: true + cfg_training_dropout: 0.25 + cfg_beta: 2.0 + +rtc: + enabled: true + max_delay_steps: 12 + include_executed_prefix: true + masked_loss: true + +flow: + timestep_beta_alpha: 1.5 + timestep_beta_beta: 1.0 + integration_steps: 10 + +train: + goal_steps: 300 + high_level_steps: 200 + policy_steps: 900 + batch_size: 5 + learning_rate: 0.001 + weight_decay: 0.0 + grad_clip_norm: 1.0 + goal_loss_weight: 1.0 + high_level_loss_weight: 1.0 + flow_loss_weight: 2.0 + rtc_loss_weight: 1.0 + log_every: 20 + +evaluation: + fixed_noise_seed: 1 + cfg_betas: [0.0, 1.0, 2.0] + ablations: + - full + - no_metadata + - no_goal + - no_history + compare_rtc_delays: [0, 4, 8, 12] + save_trajectory_plot: true + save_goal_plot: true + save_ablation_plot: true diff --git a/outputs/overfit/context_ablation.png b/outputs/overfit/context_ablation.png new file mode 100644 index 0000000..8d01cf1 Binary files /dev/null and b/outputs/overfit/context_ablation.png differ diff --git a/outputs/overfit/evaluation.json b/outputs/overfit/evaluation.json new file mode 100644 index 0000000..4b36128 --- /dev/null +++ b/outputs/overfit/evaluation.json @@ -0,0 +1,11 @@ +{ + "predicted_subtask": "steer_left", + "trajectory_metrics": { + "full_oracle_goal": 0.47031369805336, + "generated_goal": 0.4905388057231903, + "no_metadata": 2.457864284515381, + "no_goal": 4.042560577392578, + "no_history": 2.1593737602233887, + "cfg_beta2": 10.576497077941895 + } +} \ No newline at end of file diff --git a/outputs/overfit/goal_generation.png b/outputs/overfit/goal_generation.png new file mode 100644 index 0000000..a38255f Binary files /dev/null and b/outputs/overfit/goal_generation.png differ diff --git a/outputs/overfit/summary.json b/outputs/overfit/summary.json new file mode 100644 index 0000000..1f6e681 --- /dev/null +++ b/outputs/overfit/summary.json @@ -0,0 +1,7 @@ +{ + "device": "mps", + "samples": 5, + "policy_parameters": 1088930, + "goal_parameters": 75731, + "high_level_parameters": 652800 +} \ 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..01f5f9d --- /dev/null +++ b/outputs/overfit/training_history.json @@ -0,0 +1,459 @@ +[ + { + "stage": "goal", + "step": 1, + "loss": 1.0040572881698608 + }, + { + "stage": "goal", + "step": 20, + "loss": 0.34230467677116394 + }, + { + "stage": "goal", + "step": 40, + "loss": 0.37451064586639404 + }, + { + "stage": "goal", + "step": 60, + "loss": 0.26792070269584656 + }, + { + "stage": "goal", + "step": 80, + "loss": 0.10219617933034897 + }, + { + "stage": "goal", + "step": 100, + "loss": 0.2711445987224579 + }, + { + "stage": "goal", + "step": 120, + "loss": 0.2289590686559677 + }, + { + "stage": "goal", + "step": 140, + "loss": 0.36456483602523804 + }, + { + "stage": "goal", + "step": 160, + "loss": 0.0842418223619461 + }, + { + "stage": "goal", + "step": 180, + "loss": 0.9889341592788696 + }, + { + "stage": "goal", + "step": 200, + "loss": 0.5776903629302979 + }, + { + "stage": "goal", + "step": 220, + "loss": 0.44620272517204285 + }, + { + "stage": "goal", + "step": 240, + "loss": 0.05634145438671112 + }, + { + "stage": "goal", + "step": 260, + "loss": 0.05701519176363945 + }, + { + "stage": "goal", + "step": 280, + "loss": 0.06764699518680573 + }, + { + "stage": "goal", + "step": 300, + "loss": 0.10973378270864487 + }, + { + "stage": "high_level", + "step": 1, + "loss": 76.15934753417969 + }, + { + "stage": "high_level", + "step": 20, + "loss": -0.0 + }, + { + "stage": "high_level", + "step": 40, + "loss": -0.0 + }, + { + "stage": "high_level", + "step": 60, + "loss": -0.0 + }, + { + "stage": "high_level", + "step": 80, + "loss": -0.0 + }, + { + "stage": "high_level", + "step": 100, + "loss": -0.0 + }, + { + "stage": "high_level", + "step": 120, + "loss": -0.0 + }, + { + "stage": "high_level", + "step": 140, + "loss": -0.0 + }, + { + "stage": "high_level", + "step": 160, + "loss": -0.0 + }, + { + "stage": "high_level", + "step": 180, + "loss": -0.0 + }, + { + "stage": "high_level", + "step": 200, + "loss": -0.0 + }, + { + "stage": "policy", + "step": 1, + "loss": 51.04338836669922, + "text_loss": 44.21609115600586, + "flow_loss": 3.4136481285095215 + }, + { + "stage": "policy", + "step": 20, + "loss": 26.20608901977539, + "text_loss": 5.128565311431885, + "flow_loss": 10.538762092590332 + }, + { + "stage": "policy", + "step": 40, + "loss": 13.018914222717285, + "text_loss": 3.846059799194336, + "flow_loss": 4.586427211761475 + }, + { + "stage": "policy", + "step": 60, + "loss": 11.344329833984375, + "text_loss": 3.335771083831787, + "flow_loss": 4.004279136657715 + }, + { + "stage": "policy", + "step": 80, + "loss": 6.586862087249756, + "text_loss": 2.5455708503723145, + "flow_loss": 2.0206456184387207 + }, + { + "stage": "policy", + "step": 100, + "loss": 11.200895309448242, + "text_loss": 2.1739449501037598, + "flow_loss": 4.513474941253662 + }, + { + "stage": "policy", + "step": 120, + "loss": 17.273035049438477, + "text_loss": 2.1746408939361572, + "flow_loss": 7.549196720123291 + }, + { + "stage": "policy", + "step": 140, + "loss": 11.353687286376953, + "text_loss": 2.2218563556671143, + "flow_loss": 4.565915584564209 + }, + { + "stage": "policy", + "step": 160, + "loss": 14.110220909118652, + "text_loss": 2.3123977184295654, + "flow_loss": 5.898911476135254 + }, + { + "stage": "policy", + "step": 180, + "loss": 19.999370574951172, + "text_loss": 2.749119997024536, + "flow_loss": 8.62512493133545 + }, + { + "stage": "policy", + "step": 200, + "loss": 8.102465629577637, + "text_loss": 1.631801962852478, + "flow_loss": 3.2353317737579346 + }, + { + "stage": "policy", + "step": 220, + "loss": 21.903854370117188, + "text_loss": 1.6063504219055176, + "flow_loss": 10.148752212524414 + }, + { + "stage": "policy", + "step": 240, + "loss": 10.197904586791992, + "text_loss": 1.5640149116516113, + "flow_loss": 4.3169450759887695 + }, + { + "stage": "policy", + "step": 260, + "loss": 6.899936676025391, + "text_loss": 1.3081820011138916, + "flow_loss": 2.795877456665039 + }, + { + "stage": "policy", + "step": 280, + "loss": 10.051774978637695, + "text_loss": 1.2130903005599976, + "flow_loss": 4.419342517852783 + }, + { + "stage": "policy", + "step": 300, + "loss": 7.963553428649902, + "text_loss": 1.1586928367614746, + "flow_loss": 3.402430295944214 + }, + { + "stage": "policy", + "step": 320, + "loss": 5.177718639373779, + "text_loss": 1.2673476934432983, + "flow_loss": 1.9551855325698853 + }, + { + "stage": "policy", + "step": 340, + "loss": 4.487771034240723, + "text_loss": 1.0361082553863525, + "flow_loss": 1.7258312702178955 + }, + { + "stage": "policy", + "step": 360, + "loss": 4.356203556060791, + "text_loss": 1.3223515748977661, + "flow_loss": 1.5169259309768677 + }, + { + "stage": "policy", + "step": 380, + "loss": 5.289548397064209, + "text_loss": 1.4994667768478394, + "flow_loss": 1.89504075050354 + }, + { + "stage": "policy", + "step": 400, + "loss": 8.761286735534668, + "text_loss": 1.4391229152679443, + "flow_loss": 3.6610820293426514 + }, + { + "stage": "policy", + "step": 420, + "loss": 2.400188446044922, + "text_loss": 1.0106937885284424, + "flow_loss": 0.694747269153595 + }, + { + "stage": "policy", + "step": 440, + "loss": 9.683706283569336, + "text_loss": 1.044368863105774, + "flow_loss": 4.319668769836426 + }, + { + "stage": "policy", + "step": 460, + "loss": 10.443679809570312, + "text_loss": 0.7045990228652954, + "flow_loss": 4.869540214538574 + }, + { + "stage": "policy", + "step": 480, + "loss": 6.514384746551514, + "text_loss": 1.185524582862854, + "flow_loss": 2.6644301414489746 + }, + { + "stage": "policy", + "step": 500, + "loss": 2.255455255508423, + "text_loss": 1.0136079788208008, + "flow_loss": 0.620923638343811 + }, + { + "stage": "policy", + "step": 520, + "loss": 4.1700439453125, + "text_loss": 1.1829472780227661, + "flow_loss": 1.4935483932495117 + }, + { + "stage": "policy", + "step": 540, + "loss": 6.956455230712891, + "text_loss": 1.1096038818359375, + "flow_loss": 2.9234256744384766 + }, + { + "stage": "policy", + "step": 560, + "loss": 4.5965800285339355, + "text_loss": 1.5257400274276733, + "flow_loss": 1.5354200601577759 + }, + { + "stage": "policy", + "step": 580, + "loss": 1.7948781251907349, + "text_loss": 0.835124135017395, + "flow_loss": 0.4798769950866699 + }, + { + "stage": "policy", + "step": 600, + "loss": 11.411609649658203, + "text_loss": 0.7508105635643005, + "flow_loss": 5.330399513244629 + }, + { + "stage": "policy", + "step": 620, + "loss": 14.173189163208008, + "text_loss": 0.553065836429596, + "flow_loss": 6.810061454772949 + }, + { + "stage": "policy", + "step": 640, + "loss": 2.423476457595825, + "text_loss": 0.6566628813743591, + "flow_loss": 0.8834067583084106 + }, + { + "stage": "policy", + "step": 660, + "loss": 7.33753776550293, + "text_loss": 0.5569484829902649, + "flow_loss": 3.3902945518493652 + }, + { + "stage": "policy", + "step": 680, + "loss": 7.2152018547058105, + "text_loss": 0.6606324911117554, + "flow_loss": 3.277284622192383 + }, + { + "stage": "policy", + "step": 700, + "loss": 2.5858616828918457, + "text_loss": 0.471678227186203, + "flow_loss": 1.0570917129516602 + }, + { + "stage": "policy", + "step": 720, + "loss": 5.87769889831543, + "text_loss": 0.5711548924446106, + "flow_loss": 2.6532719135284424 + }, + { + "stage": "policy", + "step": 740, + "loss": 6.787323951721191, + "text_loss": 0.5064371228218079, + "flow_loss": 3.1404433250427246 + }, + { + "stage": "policy", + "step": 760, + "loss": 11.166234970092773, + "text_loss": 0.5802748203277588, + "flow_loss": 5.292980194091797 + }, + { + "stage": "policy", + "step": 780, + "loss": 3.563460350036621, + "text_loss": 0.6946841478347778, + "flow_loss": 1.4343880414962769 + }, + { + "stage": "policy", + "step": 800, + "loss": 3.2471487522125244, + "text_loss": 0.44556593894958496, + "flow_loss": 1.4007914066314697 + }, + { + "stage": "policy", + "step": 820, + "loss": 3.1103086471557617, + "text_loss": 0.7132699489593506, + "flow_loss": 1.1985193490982056 + }, + { + "stage": "policy", + "step": 840, + "loss": 4.753307342529297, + "text_loss": 0.48397523164749146, + "flow_loss": 2.1346659660339355 + }, + { + "stage": "policy", + "step": 860, + "loss": 3.5492496490478516, + "text_loss": 0.366057813167572, + "flow_loss": 1.5915958881378174 + }, + { + "stage": "policy", + "step": 880, + "loss": 4.0379414558410645, + "text_loss": 0.6511648893356323, + "flow_loss": 1.6933883428573608 + }, + { + "stage": "policy", + "step": 900, + "loss": 1.7170453071594238, + "text_loss": 0.2673894464969635, + "flow_loss": 0.7248279452323914 + } +] \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..ccc1f10 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,38 @@ +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "pi07-minimal-validation" +version = "0.1.0" +description = "A tiny, safety-scoped pi0.7-style steerable policy validation experiment" +readme = "README.md" +requires-python = ">=3.10" +license = { text = "MIT" } +authors = [{ name = "pi07 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] +pi07mini-train = "pi07mini.train:main" +pi07mini-infer = "pi07mini.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/pi07mini/__init__.py b/src/pi07mini/__init__.py new file mode 100644 index 0000000..7936b42 --- /dev/null +++ b/src/pi07mini/__init__.py @@ -0,0 +1,32 @@ +"""Minimal pi0.7 architecture validation package.""" + +from .memory import MEMConfig, MEMVideoEncoder +from .model import ( + GoalGeneratorOutput, + HighLevelOutput, + HighLevelPolicy, + Pi07Config, + Pi07Mini, + Pi07Output, + PrefixLayout, + TinyGoalGenerator, + flow_matching_sample, +) +from .tokenizer import EpisodeMetadata, MiniTokenizer, PolicySequence + +__all__ = [ + "EpisodeMetadata", + "GoalGeneratorOutput", + "HighLevelOutput", + "HighLevelPolicy", + "MEMConfig", + "MEMVideoEncoder", + "MiniTokenizer", + "Pi07Config", + "Pi07Mini", + "Pi07Output", + "PolicySequence", + "PrefixLayout", + "TinyGoalGenerator", + "flow_matching_sample", +] diff --git a/src/pi07mini/context.py b/src/pi07mini/context.py new file mode 100644 index 0000000..2e05286 --- /dev/null +++ b/src/pi07mini/context.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping, Sequence + +import torch +import torch.nn.functional as F + + +PROMPT_COMPONENT_ORDER = ("task", "speed", "quality", "mistake", "control") + + +@dataclass(frozen=True) +class PolicyMetadata: + """Steerable policy attributes used by the tiny pi0.7 experiment.""" + + speed: str + quality: str + mistake: str + control: str + speed_steps: int = 1_000 + quality_score: int = 3 + control_mode: str = "ee" + source_type: str = "autonomous" + + def __post_init__(self) -> None: + if self.speed_steps <= 0: + raise ValueError("speed_steps must be positive") + if not 1 <= self.quality_score <= 5: + raise ValueError("quality_score must lie in [1,5]") + if self.control_mode not in {"joint", "ee"}: + raise ValueError("control_mode must be joint or ee") + + def as_dict(self) -> dict[str, str]: + return { + "speed": self.speed, + "quality": self.quality, + "mistake": self.mistake, + "control": self.control, + } + + def as_token_metadata(self) -> dict[str, int | bool | str]: + """Return the scalar metadata contract consumed by MiniTokenizer.""" + + return { + "speed_steps": self.speed_steps, + "quality": self.quality_score, + "mistake": self.mistake != "none", + "control_mode": self.control_mode, + "source_type": self.source_type, + } + + +def prompt_components( + metadata: PolicyMetadata, + *, + task: str = "Follow the visible road and predict the next waypoint chunk.", +) -> dict[str, str]: + return {"task": task, **metadata.as_dict()} + + +def drop_prompt_components( + components: Mapping[str, str], + *, + drop_probability: float, + generator: torch.Generator | None = None, + preserve: Sequence[str] = ("task",), +) -> tuple[dict[str, str], torch.Tensor]: + """Apply classifier-free-style dropout to independent prompt attributes. + + The returned mask follows ``PROMPT_COMPONENT_ORDER`` and marks components + that remain visible. Task text is preserved by default, while metadata can + be independently omitted so the policy cannot require every attribute. + """ + + if not 0.0 <= drop_probability <= 1.0: + raise ValueError("drop_probability must lie in [0, 1]") + preserved = set(preserve) + visible: dict[str, str] = {} + mask = torch.zeros(len(PROMPT_COMPONENT_ORDER), dtype=torch.bool) + for index, name in enumerate(PROMPT_COMPONENT_ORDER): + if name not in components: + continue + keep = name in preserved or bool( + torch.rand((), generator=generator) >= drop_probability + ) + if keep: + visible[name] = str(components[name]) + mask[index] = True + return visible, mask + + +def format_policy_prompt(components: Mapping[str, str]) -> str: + """Format components without inventing values for dropped attributes.""" + + pieces: list[str] = [] + if task := components.get("task"): + pieces.append(str(task).strip()) + attributes = [ + f"{name}={components[name]}" + for name in PROMPT_COMPONENT_ORDER[1:] + if name in components + ] + if attributes: + pieces.append("Policy metadata: " + ", ".join(attributes) + ".") + return "\n".join(pieces) + + +def endpoint_subgoal_heatmap( + endpoint_xy: torch.Tensor, + *, + size: int = 32, + sigma: float = 1.5, + xy_extent: tuple[float, float, float, float] = (-8.0, 64.0, -16.0, 16.0), +) -> torch.Tensor: + """Render a normalized Gaussian subgoal heatmap from one endpoint only. + + It deliberately accepts no trajectory tensor. This API makes it difficult + to accidentally leak the entire ground-truth action chunk into the visual + subgoal condition. + """ + + endpoint = torch.as_tensor(endpoint_xy, dtype=torch.float32).reshape(2) + if size < 2 or sigma <= 0.0: + raise ValueError("size must be >=2 and sigma must be positive") + x_min, x_max, y_min, y_max = xy_extent + if not x_min < x_max or not y_min < y_max: + raise ValueError("xy_extent must contain increasing bounds") + px = (endpoint[0] - x_min) / (x_max - x_min) * (size - 1) + py = (y_max - endpoint[1]) / (y_max - y_min) * (size - 1) + yy, xx = torch.meshgrid( + torch.arange(size, dtype=torch.float32), + torch.arange(size, dtype=torch.float32), + indexing="ij", + ) + heatmap = torch.exp(-((xx - px) ** 2 + (yy - py) ** 2) / (2.0 * sigma**2)) + heatmap = heatmap / heatmap.amax().clamp_min(torch.finfo(heatmap.dtype).eps) + return heatmap.unsqueeze(0) + + +def make_history_proxy_frames( + image: torch.Tensor, + history_xy: torch.Tensor, + *, + frames: int = 6, +) -> torch.Tensor: + """Create oldest-to-newest history frames from one PAI image and past XY. + + PAI's tiny overfit manifest only exposes one camera frame. We therefore + synthesize a deterministic visual-history proxy by rolling that frame from + past ego displacement and applying a very small temporal exposure change. + No future waypoint is accepted by this function. + """ + + image = torch.as_tensor(image, dtype=torch.float32) + history = torch.as_tensor(history_xy, dtype=torch.float32) + if image.ndim != 3 or image.shape[0] not in {1, 3}: + raise ValueError("image must have shape [C,H,W]") + if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 2: + raise ValueError("history_xy must have shape [T>=2,2]") + if frames < 2: + raise ValueError("frames must be >=2") + sampled = F.interpolate( + history.T.unsqueeze(0), size=frames, mode="linear", align_corners=True + )[0].T + relative = sampled - sampled[-1] + spatial_scale = max(image.shape[-2:]) / 24.0 + output: list[torch.Tensor] = [] + for index, point in enumerate(relative): + dx = int(torch.round(point[1] * spatial_scale).item()) + dy = int(torch.round(-point[0] * spatial_scale * 0.25).item()) + dx = max(-image.shape[-1] // 6, min(image.shape[-1] // 6, dx)) + dy = max(-image.shape[-2] // 6, min(image.shape[-2] // 6, dy)) + # The separate current image is appended by MEM, so even the newest + # proxy frame remains a slightly older exposure rather than an exact + # duplicate of current_pixel_values. + exposure = 0.90 + 0.08 * index / (frames - 1) + output.append(torch.roll(image, shifts=(dy, dx), dims=(-2, -1)) * exposure) + return torch.stack(output).clamp(0.0, 1.0) + + +def build_ambiguous_history_pair( + history_xy: torch.Tensor, *, frames: int = 6 +) -> torch.Tensor: + """Return two distinct histories that share exactly the same current state. + + This supplies a minimal history-disambiguation probe: a current-frame-only + policy sees the same endpoint, whereas a memory encoder can distinguish the + two past paths. Construction depends only on observed history. + """ + + history = torch.as_tensor(history_xy, dtype=torch.float32) + if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 2: + raise ValueError("history_xy must have shape [T>=2,2]") + sampled = F.interpolate( + history.T.unsqueeze(0), size=frames, mode="linear", align_corners=True + )[0].T + alpha = torch.linspace(0.0, 1.0, frames, dtype=sampled.dtype).unsqueeze(-1) + chord = sampled[:1] + alpha * (sampled[-1:] - sampled[:1]) + mirrored = 2.0 * chord - sampled + if torch.allclose(mirrored, sampled, atol=1e-4, rtol=0.0): + direction = sampled[-1] - sampled[0] + perpendicular = torch.stack((-direction[1], direction[0])) + perpendicular = perpendicular / torch.linalg.vector_norm(perpendicular).clamp_min(1e-6) + amplitude = torch.linalg.vector_norm(direction).clamp_min(1.0) * 0.08 + hump = torch.sin(torch.linspace(0.0, torch.pi, frames)).unsqueeze(-1) + mirrored = sampled + hump * perpendicular * amplitude + mirrored[-1] = sampled[-1] + return torch.stack((sampled, mirrored)) diff --git a/src/pi07mini/data.py b/src/pi07mini/data.py new file mode 100644 index 0000000..877ac25 --- /dev/null +++ b/src/pi07mini/data.py @@ -0,0 +1,418 @@ +from __future__ import annotations + +import json +import shutil +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, 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 .context import ( + PolicyMetadata, + build_ambiguous_history_pair, + drop_prompt_components, + endpoint_subgoal_heatmap, + format_policy_prompt, + make_history_proxy_frames, + prompt_components, +) +from .rtc import MAX_RTC_DELAY_STEPS, align_actions_for_rtc, sample_rtc_delay + + +DEFAULT_SOURCE_MANIFEST = ( + Path(__file__).resolve().parents[4] + / "alpamayo" + / "alpamayo-qwen05b-validation" + / "data" + / "processed" + / "overfit_manifest.jsonl" +) + +PolicyKind = Literal["expert", "fast", "slow", "understeer", "stop"] + + +@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: + 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) + 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) + ] + ) + return cls(flat.mean(0), flat.std(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) + deltas = actions * self.scale.to(actions.device) + self.mean.to(actions.device) + start = torch.as_tensor(history_xy, dtype=torch.float32, device=actions.device)[-1:] + return start + torch.cumsum(deltas, dim=-2) + + def as_dict(self) -> dict[str, list[float]]: + return self.statistics.as_dict() + + +@dataclass(frozen=True) +class PolicyTrajectory: + kind: PolicyKind + history_xy: torch.Tensor + future_xy: torch.Tensor + metadata: PolicyMetadata + + +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).expanduser().resolve() + records: list[SampleRecord] = [] + for line in manifest.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, row["image_path"]), + trajectory_path=_resolve_data_path(manifest, row["trajectory_path"]), + reasoning=str(row.get("reasoning", "")), + ) + ) + if not records: + raise ValueError(f"no samples found in {manifest}") + return records + + +def materialize_subset( + source_manifest: str | Path = DEFAULT_SOURCE_MANIFEST, + output_dir: str | Path = "data/overfit", + *, + max_samples: int = 1, +) -> Path: + if max_samples < 1: + raise ValueError("max_samples must be positive") + destination = Path(output_dir).expanduser().resolve() + (destination / "images").mkdir(parents=True, exist_ok=True) + (destination / "trajectories").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 = destination / "images" / f"{stem}{record.image_path.suffix.lower()}" + trajectory_path = destination / "trajectories" / 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") + return F.interpolate(xy.T.unsqueeze(0), size=steps, mode="linear", align_corners=True)[0].T + + +def build_policy_trajectories( + history_xy: torch.Tensor, + expert_future_xy: torch.Tensor, +) -> list[PolicyTrajectory]: + """Create steerable expert/speed/failure modes from one PAI demonstration.""" + + 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) < 2: + raise ValueError("history_xy must have shape [T>=2,2]") + if expert.ndim != 2 or expert.shape[-1] != 2 or len(expert) < 4: + raise ValueError("expert_future_xy must have shape [H>=4,2]") + anchor = history[-1] + displacement = expert - anchor + progress = torch.linspace(0.0, 1.0, len(expert)).unsqueeze(-1) + + # Speed steering changes progress along the same demonstrated geometry. + fast_index = (progress.squeeze(-1) * 1.22).clamp(max=1.0) * (len(expert) - 1) + slow_index = progress.squeeze(-1) * 0.68 * (len(expert) - 1) + coordinates = torch.arange(len(expert), dtype=expert.dtype) + + def interpolate_at(indices: torch.Tensor) -> torch.Tensor: + left = indices.floor().long().clamp(max=len(expert) - 1) + right = (left + 1).clamp(max=len(expert) - 1) + weight = (indices - left).unsqueeze(-1) + return expert[left] * (1.0 - weight) + expert[right] * weight + + del coordinates # documents that interpolation is index based, not metric based + fast = interpolate_at(fast_index) + slow = interpolate_at(slow_index) + understeer = anchor + displacement * expert.new_tensor([1.0, 0.22]) + stop_index = max(2, len(expert) // 3) + stopped = expert.clone() + stopped[stop_index:] = stopped[stop_index - 1] + return [ + PolicyTrajectory( + "expert", + history, + expert, + PolicyMetadata( + "normal", "expert", "none", "human", 1_000, 5, "ee", "demonstration" + ), + ), + PolicyTrajectory( + "fast", + history, + fast, + PolicyMetadata( + "fast", "good", "none", "autonomous", 500, 4, "ee", "autonomous" + ), + ), + PolicyTrajectory( + "slow", + history, + slow, + PolicyMetadata( + "slow", "good", "none", "autonomous", 2_000, 4, "ee", "autonomous" + ), + ), + PolicyTrajectory( + "understeer", + history, + understeer, + PolicyMetadata( + "normal", "poor", "understeer", "autonomous", 1_000, 1, "ee", "autonomous" + ), + ), + PolicyTrajectory( + "stop", + history, + stopped, + PolicyMetadata( + "slow", "poor", "stop", "autonomous", 2_000, 1, "ee", "autonomous" + ), + ), + ] + + +class Pi07SteerableDataset(Dataset[dict[str, Any]]): + """Five-policy PAI overfit set with history, subgoal, prompt and RTC inputs.""" + + def __init__( + self, + manifest_path: str | Path, + *, + image_size: int = 64, + augmentation_factor: int = 1, + history_frames: int = 6, + action_horizon: int = 16, + prompt_dropout: float = 0.25, + max_rtc_delay_steps: int = MAX_RTC_DELAY_STEPS, + seed: int = 7, + action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None, + ) -> None: + if image_size < 2 or augmentation_factor < 1 or history_frames < 2: + raise ValueError("image_size, augmentation_factor and history_frames are invalid") + self.records = load_manifest(manifest_path) + self.image_size = int(image_size) + self.augmentation_factor = int(augmentation_factor) + self.history_frames = int(history_frames) + self.action_horizon = int(action_horizon) + self.prompt_dropout = float(prompt_dropout) + self.max_rtc_delay_steps = int(max_rtc_delay_steps) + self.seed = int(seed) + + policies: list[PolicyTrajectory] = [] + 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_frames, + ) + future = _resample_xy( + torch.from_numpy(np.asarray(arrays["future_xyz"], dtype=np.float32))[:, :2], + self.action_horizon, + ) + generated = build_policy_trajectories(history, future) + policies.extend(generated) + record_indices.extend([record_index] * len(generated)) + self.policies = policies + self._record_indices = record_indices + if action_statistics is None: + self.action_normalizer = ActionNormalizer.fit( + [policy.history_xy for policy in policies], + [policy.future_xy for policy in policies], + ) + 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 __len__(self) -> int: + return len(self.policies) * self.augmentation_factor + + def _indices(self, index: int) -> tuple[int, int]: + if index < 0: + index += len(self) + if not 0 <= 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), 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]: + policy_index, variant = self._indices(index) + policy = self.policies[policy_index] + record_index = self._record_indices[policy_index] + image = self._load_image(record_index, variant) + actions = self.action_normalizer.encode(policy.history_xy, policy.future_xy) + generator = torch.Generator().manual_seed(self.seed * 100_003 + index) + components = prompt_components(policy.metadata) + visible_components, component_mask = drop_prompt_components( + components, + drop_probability=self.prompt_dropout, + generator=generator, + ) + delay = sample_rtc_delay( + generator=generator, + max_delay_steps=self.max_rtc_delay_steps, + action_horizon=self.action_horizon, + ) + rtc = align_actions_for_rtc( + actions, delay, max_delay_steps=self.max_rtc_delay_steps + ) + subgoal = endpoint_subgoal_heatmap( + policy.future_xy[-1], size=self.image_size, sigma=max(1.0, self.image_size / 24.0) + ) + history_pair = build_ambiguous_history_pair( + policy.history_xy, frames=self.history_frames + ) + history_images = make_history_proxy_frames( + image, policy.history_xy, frames=self.history_frames + ) + return { + "image": image, + "pixel_values": image, + "history_images": history_images, + "history_pixel_values": history_images, + "history_xy": policy.history_xy.clone(), + "history_states": policy.history_xy.clone(), + "current_state": policy.history_xy[-1].clone(), + "ambiguous_history_pair": history_pair, + "ambiguous_pair_id": f"{self.records[record_index].sample_id}:current-state", + "future_xy": policy.future_xy.clone(), + "actions": actions, + "policy_kind": policy.kind, + "metadata": policy.metadata.as_dict(), + "token_metadata": policy.metadata.as_token_metadata(), + "speed": policy.metadata.speed, + "quality": policy.metadata.quality, + "mistake": policy.metadata.mistake, + "control": policy.metadata.control, + "speed_steps": policy.metadata.speed_steps, + "quality_score": policy.metadata.quality_score, + "has_mistake": policy.metadata.mistake != "none", + "control_mode": policy.metadata.control_mode, + "source_type": policy.metadata.source_type, + "prompt_components": visible_components, + "prompt_component_mask": component_mask, + "prompt_text": format_policy_prompt(visible_components), + "subgoal_xy": policy.future_xy[-1].clone(), + "subgoal_heatmap": subgoal, + "subgoal_pixel_values": subgoal.expand(3, -1, -1).clone(), + "rtc_delay_steps": delay, + "rtc_target_actions": rtc.delayed_target_actions, + "rtc_valid_mask": rtc.valid_mask, + "rtc_executed_prefix_actions": rtc.executed_prefix_actions, + "reasoning": self.records[record_index].reasoning, + "sample_id": f"{self.records[record_index].sample_id}:{policy.kind}:aug{variant}", + } diff --git a/src/pi07mini/flow.py b/src/pi07mini/flow.py new file mode 100644 index 0000000..eb93db6 --- /dev/null +++ b/src/pi07mini/flow.py @@ -0,0 +1,18 @@ +from __future__ import annotations +from collections.abc import Callable +import torch +from torch import Tensor + + +class FlowMatching: + 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]: + u=torch.rand(target.shape[0],device=target.device,dtype=target.dtype) + t=self.cutoff*(1-u.pow(1/1.5));noise=torch.randn_like(target) + return t[:,None,None]*target+(1-t[:,None,None])*noise,t,target-noise + @torch.no_grad() + def sample(self,fn:Callable[[Tensor,Tensor],Tensor],shape:tuple[int,...],device,steps:int|None=None)->Tensor: + count=steps or self.inference_steps;x=torch.randn(shape,device=device);grid=torch.linspace(0,1,count+1,device=device) + for i in range(count): x=x+(grid[i+1]-grid[i])*fn(x,grid[i].expand(shape[0])) + return x diff --git a/src/pi07mini/infer.py b/src/pi07mini/infer.py new file mode 100644 index 0000000..eaf9079 --- /dev/null +++ b/src/pi07mini/infer.py @@ -0,0 +1,40 @@ +from __future__ import annotations +import argparse +from pathlib import Path +import matplotlib.pyplot as plt +import numpy as np,torch +from .data import Pi07SteerableDataset +from .flow import FlowMatching +from .model import Pi07Config,Pi07Mini,HighLevelPolicy,TinyGoalGenerator +from .runtime import load_config,choose_device,write_json +from .tokenizer import MiniTokenizer,EpisodeMetadata + +@torch.no_grad() +def evaluate(checkpoint_path): + ck=torch.load(checkpoint_path,map_location='cpu',weights_only=False);cfg=ck['config'];root=Path(__file__).resolve().parents[2];dev=choose_device(cfg['device']);tok=MiniTokenizer();ds=Pi07SteerableDataset(root/'data/overfit/overfit_manifest.jsonl',image_size=cfg['data']['image_size'],history_frames=cfg['data']['history_frames'],action_horizon=cfg['data']['action_horizon'],action_statistics=ck['action_statistics']) + pcfg=Pi07Config(**ck['policy_config']);policy=Pi07Mini(pcfg);policy.load_state_dict(ck['policy']);policy.to(dev).eval();high=HighLevelPolicy(pcfg);high.load_state_dict(ck['high_level']);high.to(dev).eval();goal=TinyGoalGenerator(tok.vocab_size);goal.load_state_dict(ck['goal']);goal.to(dev).eval();flow=FlowMatching(cfg['flow']['integration_steps']);s=next(ds[i] for i in range(len(ds)) if ds[i]['policy_kind']=='expert');meta=EpisodeMetadata(**s['token_metadata']);image=s['image'][None].to(dev) + prompt,_=tok.encode_prompt('drive safely',None,meta,include_subtask=False);hl=high(image,prompt[None].to(dev),s['history_images'][None].to(dev),s['history_states'][None].to(dev),s['current_state'][None].to(dev));pred_subtask=tok.subtasks.decode(int(hl.logits.argmax(-1).cpu())) + full=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta);ids=full.input_ids[:full.action_start_index][None].to(dev) + torch.manual_seed(cfg['seed']);g=torch.randn_like(s['subgoal_pixel_values'][None]).to(dev);grid=torch.linspace(0,1,6,device=dev) + for i in range(5): g=g+(grid[i+1]-grid[i])*goal(image,ids,g,grid[i].expand(1)).vector_field + generated_goal=g.cpu()[0] + def sample(name,include_meta=True,use_goal=True,use_history=True,cfg_beta=None,generated=False): + seq=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=include_meta);xids=seq.input_ids[:seq.action_start_index][None].to(dev) + goal_pixels=(generated_goal if generated else s['subgoal_pixel_values'])[None].to(dev) if use_goal else None + common=dict(pixel_values=image,history_pixel_values=s['history_images'][None].to(dev) if use_history else None,history_states=s['history_states'][None].to(dev) if use_history else None,current_state=s['current_state'][None].to(dev),subgoal_pixel_values=goal_pixels) + def field(x,t): + if cfg_beta is None:return policy(input_ids=xids,noisy_actions=x,timesteps=t,**common).vector_field + un=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=False);uids=un.input_ids[:un.action_start_index][None].to(dev) + return policy.metadata_cfg(xids,uids,guidance_scale=cfg_beta,noisy_actions=x,timesteps=t,**common) + torch.manual_seed(cfg['evaluation']['fixed_noise_seed']);a=flow.sample(field,(1,ds.action_horizon,2),dev,cfg['flow']['integration_steps'])[0].cpu();return ds.action_normalizer.decode(a,s['history_xy']).numpy() + curves={'full_oracle_goal':sample('full'),'generated_goal':sample('generated',generated=True),'no_metadata':sample('no_metadata',False),'no_goal':sample('no_goal',use_goal=False),'no_history':sample('no_history',use_history=False),'cfg_beta2':sample('cfg',cfg_beta=2.0)} + target=ds.action_normalizer.decode(s['rtc_target_actions'],s['history_xy']).numpy();valid=s['rtc_valid_mask'].numpy().astype(bool);metrics={k:float(np.mean((v[valid]-target[valid])**2)) for k,v in curves.items()} + out=Path(checkpoint_path).parent;fig,ax=plt.subplots(figsize=(9,6));ax.plot(target[:,0],target[:,1],'o-',label='expert target',linewidth=2) + for k,v in curves.items():ax.plot(v[:,0],v[:,1],'x--',label=f'{k} ({metrics[k]:.3f})') + 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.7 context steering and ablations');fig.tight_layout();traj=out/'context_ablation.png';fig.savefig(traj,dpi=170);plt.close(fig) + fig,axes=plt.subplots(1,3,figsize=(10,3));axes[0].imshow(image[0].cpu().permute(1,2,0));axes[0].set_title('current observation');axes[1].imshow(s['subgoal_pixel_values'].permute(1,2,0));axes[1].set_title('endpoint goal target');axes[2].imshow(generated_goal.permute(1,2,0).clamp(0,1));axes[2].set_title('generated goal'); + for a in axes:a.axis('off') + fig.tight_layout();goalpath=out/'goal_generation.png';fig.savefig(goalpath,dpi=170);plt.close(fig);write_json(out/'evaluation.json',{'predicted_subtask':pred_subtask,'trajectory_metrics':metrics});print({'predicted_subtask':pred_subtask,'trajectory_metrics':metrics});return traj,goalpath +def main(): + p=argparse.ArgumentParser();p.add_argument('--checkpoint',default='outputs/overfit/checkpoint.pt');evaluate(p.parse_args().checkpoint) +if __name__=='__main__':main() diff --git a/src/pi07mini/memory.py b/src/pi07mini/memory.py new file mode 100644 index 0000000..8bd73b3 --- /dev/null +++ b/src/pi07mini/memory.py @@ -0,0 +1,215 @@ +"""A tiny MEM-style video encoder used by the pi0.7 architecture test. + +The real MEM encoder modifies a pretrained ViT by alternating spatial +attention within each frame and causal temporal attention across the same +patch in different frames. It then discards all past-frame representations +and forwards only the temporally enriched current-frame patches to the VLM. + +This module preserves those information-flow contracts in pure PyTorch. It +is intentionally small and does not attempt to reproduce the pretrained ViT. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +from torch import Tensor, nn + + +@dataclass(frozen=True) +class MEMConfig: + image_channels: int = 3 + image_size: int = 64 + patch_size: int = 16 + d_model: int = 128 + n_heads: int = 4 + n_layers: int = 2 + mlp_ratio: int = 4 + max_history_frames: int = 6 + dropout: float = 0.0 + + def __post_init__(self) -> None: + if self.image_size % self.patch_size: + raise ValueError("image_size must be divisible by patch_size") + if self.d_model % self.n_heads: + raise ValueError("d_model must be divisible by n_heads") + if self.max_history_frames < 0: + raise ValueError("max_history_frames must be non-negative") + + @property + def output_tokens(self) -> int: + return (self.image_size // self.patch_size) ** 2 + + +class _MEMLayer(nn.Module): + """Factorized spatial-bidirectional and temporal-causal attention.""" + + def __init__(self, config: MEMConfig) -> None: + super().__init__() + d = config.d_model + self.n_heads = config.n_heads + self.spatial_norm = nn.LayerNorm(d) + self.spatial_attention = nn.MultiheadAttention( + d, config.n_heads, dropout=config.dropout, batch_first=True + ) + self.temporal_norm = nn.LayerNorm(d) + self.temporal_attention = nn.MultiheadAttention( + d, config.n_heads, dropout=config.dropout, batch_first=True + ) + hidden = d * config.mlp_ratio + self.mlp_norm = nn.LayerNorm(d) + self.mlp = nn.Sequential( + nn.Linear(d, hidden), + nn.GELU(approximate="tanh"), + nn.Dropout(config.dropout), + nn.Linear(hidden, d), + nn.Dropout(config.dropout), + ) + + def forward(self, tokens: Tensor, frame_valid: Tensor) -> Tensor: + batch, frames, patches, width = tokens.shape + + # Standard ViT spatial attention is bidirectional inside each frame. + spatial = self.spatial_norm(tokens).reshape(batch * frames, patches, width) + spatial_update, _ = self.spatial_attention(spatial, spatial, spatial, need_weights=False) + tokens = tokens + spatial_update.reshape(batch, frames, patches, width) + + # MEM factorizes time: each spatial patch attends causally over frames. + temporal = self.temporal_norm(tokens).permute(0, 2, 1, 3) + temporal = temporal.reshape(batch * patches, frames, width) + causal_allowed = torch.tril( + torch.ones(frames, frames, dtype=torch.bool, device=tokens.device) + ) + allowed = causal_allowed[None] & frame_valid[:, None] + # An invalid query receives a harmless self edge, preventing all-masked + # rows from producing NaNs. Its output is zeroed after the layer. + identity = torch.eye(frames, dtype=torch.bool, device=tokens.device)[None] + allowed = torch.where(frame_valid[:, :, None], allowed, identity) + blocked = (~allowed)[:, None].expand(batch, patches, frames, frames) + blocked = blocked.reshape(batch * patches, frames, frames) + blocked = blocked[:, None].expand(-1, self.n_heads, -1, -1) + blocked = blocked.reshape(batch * patches * self.n_heads, frames, frames) + temporal_update, _ = self.temporal_attention( + temporal, + temporal, + temporal, + attn_mask=blocked, + need_weights=False, + ) + temporal_update = temporal_update.reshape(batch, patches, frames, width) + tokens = tokens + temporal_update.permute(0, 2, 1, 3) + tokens = tokens + self.mlp(self.mlp_norm(tokens)) + + # Invalid history frames must not leak through residual paths. + return tokens * frame_valid[:, :, None, None].to(tokens.dtype) + + +class MEMVideoEncoder(nn.Module): + """Compress arbitrary valid history lengths to one frame's patch count. + + ``history_pixel_values`` is ordered oldest-to-newest. The current image is + appended internally and is always valid. A fixed sinusoidal temporal + encoding and causal temporal attention make the representation sensitive + to frame order, while returning only current-frame tokens keeps the output + size independent of history length. + """ + + def __init__(self, config: MEMConfig) -> None: + super().__init__() + self.config = config + self.patch_encoder = nn.Conv2d( + config.image_channels, + config.d_model, + config.patch_size, + stride=config.patch_size, + ) + self.spatial_position = nn.Parameter( + torch.randn(1, 1, config.output_tokens, config.d_model) * 0.01 + ) + self.layers = nn.ModuleList(_MEMLayer(config) for _ in range(config.n_layers)) + self.output_norm = nn.LayerNorm(config.d_model) + + @staticmethod + def temporal_position(ages: Tensor, width: int) -> Tensor: + """Fixed sinusoidal encoding whose current-frame value (age 0) is zero.""" + + if width <= 0: + raise ValueError("temporal embedding width must be positive") + half = (width + 1) // 2 + exponent = torch.arange(half, device=ages.device, dtype=torch.float32) + denominator = torch.exp(exponent * (-math.log(10_000.0) / max(half - 1, 1))) + angles = ages.float()[:, None] * denominator[None] + # Subtracting one from cosine ensures a zero embedding at age zero, + # matching MEM's invariant single-frame initialization contract. + embedding = torch.cat((angles.sin(), angles.cos() - 1.0), dim=-1) + return embedding[:, :width] + + def _validate_images(self, images: Tensor, name: str) -> None: + if images.ndim != 4: + raise ValueError(f"{name} must have shape [batch, channels, height, width]") + expected = (self.config.image_channels, self.config.image_size, self.config.image_size) + if tuple(images.shape[1:]) != expected: + raise ValueError(f"{name} must end in shape {expected}") + + def forward( + self, + current_pixel_values: Tensor, + history_pixel_values: Tensor | None = None, + history_mask: Tensor | None = None, + ) -> Tensor: + self._validate_images(current_pixel_values, "current_pixel_values") + batch = current_pixel_values.shape[0] + + if history_pixel_values is None: + history_pixel_values = current_pixel_values.new_empty( + batch, + 0, + self.config.image_channels, + self.config.image_size, + self.config.image_size, + ) + if history_pixel_values.ndim != 5: + raise ValueError( + "history_pixel_values must have shape [batch, time, channels, height, width]" + ) + if history_pixel_values.shape[0] != batch: + raise ValueError("history and current images must have the same batch size") + if tuple(history_pixel_values.shape[2:]) != tuple(current_pixel_values.shape[1:]): + raise ValueError("history and current images must have the same image shape") + history = history_pixel_values.shape[1] + if history > self.config.max_history_frames: + raise ValueError("history length exceeds configured maximum") + + if history_mask is None: + history_mask = torch.ones( + batch, history, dtype=torch.bool, device=current_pixel_values.device + ) + if history_mask.shape != (batch, history): + raise ValueError("history_mask must have shape [batch, history]") + history_mask = history_mask.bool() + current_valid = torch.ones(batch, 1, dtype=torch.bool, device=current_pixel_values.device) + frame_valid = torch.cat((history_mask, current_valid), dim=1) + + frames = torch.cat((history_pixel_values, current_pixel_values[:, None]), dim=1) + flat = frames.reshape( + batch * (history + 1), + self.config.image_channels, + self.config.image_size, + self.config.image_size, + ) + tokens = self.patch_encoder(flat).flatten(2).transpose(1, 2) + tokens = tokens.reshape(batch, history + 1, self.config.output_tokens, -1) + tokens = tokens + self.spatial_position + + # Oldest frame has the largest age; current frame has age zero. + ages = torch.arange(history, -1, -1, device=tokens.device) + temporal_position = self.temporal_position(ages, self.config.d_model) + tokens = tokens + temporal_position[None, :, None].to(tokens.dtype) + tokens = tokens * frame_valid[:, :, None, None].to(tokens.dtype) + for layer in self.layers: + tokens = layer(tokens, frame_valid) + + # Fixed number of tokens regardless of how many history frames entered. + return self.output_norm(tokens[:, -1]) diff --git a/src/pi07mini/model.py b/src/pi07mini/model.py new file mode 100644 index 0000000..4c66a2a --- /dev/null +++ b/src/pi07mini/model.py @@ -0,0 +1,816 @@ +r"""Minimal :math:`\pi_{0.7}` policy with MEM and diverse context. + +The implementation is a small architecture validation model, not a scaled +reproduction. It preserves the paper's essential paths: + +* a MEM-style encoder compresses ordered visual history to a fixed token count; +* current and historical continuous proprioceptive states are individual tokens; +* observations, visual subgoals, and causal text form distinct attention blocks; +* the VLM predicts FAST-lite tokens with cross entropy; +* an AdaRMSNorm flow expert predicts continuous actions; +* FAST tokens and flow actions cannot read one another; and +* flow-to-VLM gradients are stopped by Knowledge Insulation. +""" + +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 + +from .memory import MEMConfig, MEMVideoEncoder + + +@dataclass +class Pi07Config: + vocab_size: int = 368 + action_dim: int = 2 + state_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 + memory_layers: int = 2 + mlp_ratio: int = 4 + max_history_frames: int = 6 + max_text_tokens: int = 96 + max_action_tokens: int = 50 + action_start_token_id: int = 302 + 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 image_tokens(self) -> int: + return (self.image_size // self.patch_size) ** 2 + + @property + def max_prefix_tokens(self) -> int: + # MEM observation + all states + one subgoal view + textual prompt. + return ( + self.image_tokens + + self.max_history_frames + + 1 + + self.image_tokens + + self.max_text_tokens + ) + + +@dataclass(frozen=True) +class PrefixLayout: + observation_tokens: int + goal_tokens: int + text_tokens: int + + @property + def goal_start(self) -> int: + return self.observation_tokens + + @property + def text_start(self) -> int: + return self.observation_tokens + self.goal_tokens + + @property + def total_tokens(self) -> int: + return self.text_start + self.text_tokens + + +@dataclass +class Pi07Output: + text_logits: Tensor + vector_field: Tensor + prefix_hidden: Tensor + action_hidden: Tensor + layout: PrefixLayout + loss: Optional[Tensor] = None + text_loss: Optional[Tensor] = None + flow_loss: Optional[Tensor] = None + + +@dataclass +class HighLevelOutput: + logits: Tensor + loss: Optional[Tensor] = None + + +@dataclass +class GoalGeneratorOutput: + vector_field: Tensor + 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): + def __init__(self, dim: int, condition_dim: int, eps: float = 1e-6) -> None: + super().__init__() + self.norm = RMSNorm(dim, eps) + self.modulation = nn.Linear(condition_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): + 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: Tensor, + dropout: float, + training: bool, +) -> Tensor: + batch, query_length, 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 = q @ k.transpose(-2, -1) / math.sqrt(head_dim) + 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, query_length, width) + + +class KnowledgeInsulatedBlock(nn.Module): + """Independent VLM/flow weights with stop-gradient VLM keys and values.""" + + def __init__(self, config: Pi07Config) -> 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: continuous flow loss can condition on VLM + # features but cannot update the VLM or its K/V projections. + key = torch.cat((p.k(prefix_input).detach(), a.k(action_input)), dim=1) + value = torch.cat((p.v(prefix_input).detach(), 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 Pi07Mini(nn.Module): + """Small pi0.7-style VLA for structural and overfitting experiments.""" + + def __init__(self, config: Optional[Pi07Config] = None) -> None: + super().__init__() + self.config = config or Pi07Config() + c = self.config + self.memory_encoder = MEMVideoEncoder( + MEMConfig( + image_channels=c.image_channels, + image_size=c.image_size, + patch_size=c.patch_size, + d_model=c.d_model, + n_heads=c.n_heads, + n_layers=c.memory_layers, + mlp_ratio=c.mlp_ratio, + max_history_frames=c.max_history_frames, + dropout=c.dropout, + ) + ) + self.state_projection = nn.Linear(c.state_dim, c.d_model) + self.token_embedding = nn.Embedding(c.vocab_size, c.d_model) + self.segment_embedding = nn.Embedding(3, c.d_model) + self.prefix_position = nn.Parameter( + torch.randn(1, c.max_prefix_tokens, c.d_model) * 0.01 + ) + self.action_projection = nn.Linear(c.action_dim, c.d_model) + 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 _adapt_goal_channels(self, pixels: Tensor) -> Tensor: + if pixels.shape[1] == self.config.image_channels: + return pixels + if pixels.shape[1] == 1 and self.config.image_channels == 3: + return pixels.expand(-1, 3, -1, -1) + raise ValueError("subgoal image has unsupported channel count") + + def encode_prefix( + self, + pixel_values: Tensor, + history_pixel_values: Tensor | None, + history_states: Tensor | None, + current_state: Tensor, + input_ids: Tensor, + subgoal_pixel_values: Tensor | None = None, + history_mask: Tensor | None = None, + subgoal_mask: Tensor | None = None, + attention_mask: Tensor | None = None, + ) -> tuple[Tensor, Tensor, PrefixLayout]: + observation_images = self.memory_encoder( + pixel_values, history_pixel_values, history_mask + ) + batch = pixel_values.shape[0] + if current_state.shape != (batch, self.config.state_dim): + raise ValueError("current_state must have shape [batch, state_dim]") + + if history_states is None: + history_states = current_state.new_empty(batch, 0, self.config.state_dim) + if history_states.ndim != 3 or history_states.shape[0] != batch: + raise ValueError("history_states must have shape [batch, time, state_dim]") + if history_states.shape[-1] != self.config.state_dim: + raise ValueError("history_states has the wrong state dimension") + history_length = history_states.shape[1] + if history_length > self.config.max_history_frames: + raise ValueError("state history exceeds configured maximum") + if history_pixel_values is not None and history_pixel_values.shape[1] != history_length: + raise ValueError("visual and proprioceptive history lengths must match") + if history_mask is None: + history_mask = torch.ones( + batch, history_length, dtype=torch.bool, device=pixel_values.device + ) + if history_mask.shape != (batch, history_length): + raise ValueError("history_mask must have shape [batch, history]") + + all_states = torch.cat((history_states, current_state[:, None]), dim=1) + state_tokens = self.state_projection(all_states) + ages = torch.arange(history_length, -1, -1, device=state_tokens.device) + state_tokens = state_tokens + MEMVideoEncoder.temporal_position( + ages, self.config.d_model + )[None].to(state_tokens.dtype) + state_valid = torch.cat( + ( + history_mask.bool(), + torch.ones(batch, 1, dtype=torch.bool, device=pixel_values.device), + ), + dim=1, + ) + state_tokens = state_tokens * state_valid[:, :, None].to(state_tokens.dtype) + observation = torch.cat((observation_images, state_tokens), dim=1) + observation_valid = torch.cat( + ( + torch.ones( + batch, + observation_images.shape[1], + dtype=torch.bool, + device=pixel_values.device, + ), + state_valid, + ), + dim=1, + ) + + if subgoal_pixel_values is None: + goal = observation.new_empty(batch, 0, self.config.d_model) + goal_valid = torch.empty(batch, 0, dtype=torch.bool, device=pixel_values.device) + else: + goal_pixels = self._adapt_goal_channels(subgoal_pixel_values) + goal = self.memory_encoder(goal_pixels) + if subgoal_mask is None: + subgoal_mask = torch.ones(batch, dtype=torch.bool, device=pixel_values.device) + if subgoal_mask.shape != (batch,): + raise ValueError("subgoal_mask must have shape [batch]") + goal_valid = subgoal_mask.bool()[:, None].expand(batch, goal.shape[1]) + goal = goal * goal_valid[:, :, None].to(goal.dtype) + + if input_ids.ndim != 2 or input_ids.shape[0] != batch: + raise ValueError("input_ids must have shape [batch, text]") + text = self.token_embedding(input_ids) + if attention_mask is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + if attention_mask.shape != input_ids.shape: + raise ValueError("attention_mask must match input_ids") + text_valid = attention_mask.bool() + + layout = PrefixLayout(observation.shape[1], goal.shape[1], text.shape[1]) + if layout.total_tokens > self.config.max_prefix_tokens: + raise ValueError("prefix exceeds configured maximum length") + prefix = torch.cat((observation, goal, text), dim=1) + segments = torch.cat( + ( + torch.zeros(layout.observation_tokens, dtype=torch.long, device=prefix.device), + torch.ones(layout.goal_tokens, dtype=torch.long, device=prefix.device), + torch.full( + (layout.text_tokens,), 2, dtype=torch.long, device=prefix.device + ), + ) + ) + prefix = prefix + self.segment_embedding(segments)[None] + prefix = prefix + self.prefix_position[:, : layout.total_tokens] + valid = torch.cat((observation_valid, goal_valid, text_valid), dim=1) + return prefix, valid, layout + + @staticmethod + def build_prefix_attention_mask(layout: PrefixLayout, valid: Tensor) -> Tensor: + """Build obs -> goal -> causal-text block attention. + + Observation queries only read observations. Goal-image queries read + observations and all goal tokens bidirectionally. Text queries read + both visual blocks and earlier text tokens causally. + """ + + length = layout.total_tokens + if valid.shape[1] != length: + raise ValueError("valid mask does not match PrefixLayout") + device = valid.device + q = torch.arange(length, device=device)[:, None] + k = torch.arange(length, device=device)[None] + obs = q < layout.goal_start + goal = (q >= layout.goal_start) & (q < layout.text_start) + text = q >= layout.text_start + structural = ( + (obs & (k < layout.goal_start)) + | (goal & (k < layout.text_start)) + | (text & ((k < layout.text_start) | ((k >= layout.text_start) & (k <= q)))) + ) + edges = structural[None] & valid[:, None] + identity = torch.eye(length, dtype=torch.bool, device=device)[None] + return torch.where(valid[:, :, None], edges, identity) + + def _derive_action_prefix_mask(self, input_ids: Tensor) -> Tensor: + is_fast = input_ids.eq(self.config.action_start_token_id) + return is_fast.long().cumsum(dim=1).eq(0) + + def _action_attention_mask( + self, + valid: Tensor, + layout: PrefixLayout, + action_length: int, + action_prefix_mask: Tensor, + ) -> Tensor: + batch = valid.shape[0] + supplied = action_prefix_mask.bool() + if supplied.shape == (batch, layout.text_tokens): + supplied = torch.cat( + ( + torch.ones( + batch, + layout.text_start, + dtype=torch.bool, + device=valid.device, + ), + supplied, + ), + dim=1, + ) + if supplied.shape != valid.shape: + raise ValueError("action_prefix_mask must cover text or the full prefix") + readable = valid & supplied + prefix_keys = readable[:, None].expand(batch, action_length, layout.total_tokens) + action_keys = torch.ones( + batch, action_length, action_length, dtype=torch.bool, device=valid.device + ) + return torch.cat((prefix_keys, action_keys), dim=-1) + + def vlm_parameters(self) -> Iterator[nn.Parameter]: + modules: list[nn.Module] = [ + self.memory_encoder, + self.state_projection, + self.token_embedding, + self.segment_embedding, + self.prefix_norm, + self.text_head, + ] + for layer in self.layers: + modules.extend( + (layer.vlm_norm1, layer.vlm_attention, layer.vlm_norm2, layer.vlm_mlp) + ) + seen: set[int] = set() + if id(self.prefix_position) not in seen: + seen.add(id(self.prefix_position)) + yield self.prefix_position + 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]: + 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, + history_pixel_values: Tensor | None = None, + history_states: Tensor | None = None, + current_state: Tensor | None = None, + subgoal_pixel_values: Tensor | None = None, + history_mask: Tensor | None = None, + subgoal_mask: Tensor | None = None, + attention_mask: Tensor | None = None, + action_prefix_mask: Tensor | None = None, + generated_token_start: int | None = None, + text_labels: Tensor | None = None, + target_velocity: Tensor | None = None, + text_loss_weight: float = 1.0, + flow_loss_weight: float = 1.0, + ) -> Pi07Output: + if generated_token_start not in (None, 0): + raise ValueError("pi0.7 prompt text is causal from its first token") + 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") + if current_state is None: + current_state = noisy_actions.new_zeros( + noisy_actions.shape[0], self.config.state_dim + ) + + prefix, valid, layout = self.encode_prefix( + pixel_values, + history_pixel_values, + history_states, + current_state, + input_ids, + subgoal_pixel_values, + history_mask, + subgoal_mask, + attention_mask, + ) + 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 = self.build_prefix_attention_mask(layout, valid) + action_edges = self._action_attention_mask( + valid, layout, horizon, action_prefix_mask + ) + 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[:, layout.text_start :]) + 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 Pi07Output( + text_logits=text_logits, + vector_field=vector_field, + prefix_hidden=prefix_hidden, + action_hidden=action_hidden, + layout=layout, + loss=loss, + text_loss=text_loss, + flow_loss=flow_loss, + ) + + @staticmethod + def combine_metadata_cfg( + conditioned: Tensor, unconditioned: Tensor, guidance_scale: float = 1.5 + ) -> Tensor: + """pi0.7 metadata CFG: conditional + beta*(conditional-unconditional).""" + + if guidance_scale < 0: + raise ValueError("guidance_scale must be non-negative") + if conditioned.shape != unconditioned.shape: + raise ValueError("conditioned and unconditioned predictions must match") + return conditioned + guidance_scale * (conditioned - unconditioned) + + def metadata_cfg( + self, + conditioned_input_ids: Tensor, + unconditioned_input_ids: Tensor, + *, + guidance_scale: float = 1.5, + conditioned_attention_mask: Tensor | None = None, + unconditioned_attention_mask: Tensor | None = None, + conditioned_action_prefix_mask: Tensor | None = None, + unconditioned_action_prefix_mask: Tensor | None = None, + **common_inputs: Tensor | None, + ) -> Tensor: + """Run metadata-positive and metadata-dropped branches and combine them.""" + + conditioned = self( + input_ids=conditioned_input_ids, + attention_mask=conditioned_attention_mask, + action_prefix_mask=conditioned_action_prefix_mask, + **common_inputs, + ).vector_field + unconditioned = self( + input_ids=unconditioned_input_ids, + attention_mask=unconditioned_attention_mask, + action_prefix_mask=unconditioned_action_prefix_mask, + **common_inputs, + ).vector_field + return self.combine_metadata_cfg(conditioned, unconditioned, guidance_scale) + + +def flow_matching_sample(actions: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Construct the linear conditional-flow target used by the action expert.""" + + batch = actions.shape[0] + timestep = torch.rand(batch, device=actions.device, dtype=actions.dtype) + noise = torch.randn_like(actions) + broadcast = timestep.reshape(batch, *([1] * (actions.ndim - 1))) + noisy_actions = (1 - broadcast) * noise + broadcast * actions + return noisy_actions, timestep, actions - noise + + +class HighLevelPolicy(nn.Module): + """Tiny semantic policy that predicts the next subtask prompt token. + + The paper uses a separate policy based on the same VLM family. This proxy + shares the same MEM input contract while replacing the 4B VLM with a small + Transformer encoder suitable for a few-example overfit test. + """ + + def __init__(self, config: Optional[Pi07Config] = None) -> None: + super().__init__() + self.config = config or Pi07Config() + c = self.config + self.memory_encoder = MEMVideoEncoder( + MEMConfig( + image_channels=c.image_channels, + image_size=c.image_size, + patch_size=c.patch_size, + d_model=c.d_model, + n_heads=c.n_heads, + n_layers=c.memory_layers, + mlp_ratio=c.mlp_ratio, + max_history_frames=c.max_history_frames, + dropout=c.dropout, + ) + ) + self.state_projection = nn.Linear(c.state_dim, c.d_model) + self.token_embedding = nn.Embedding(c.vocab_size, c.d_model) + layer = nn.TransformerEncoderLayer( + c.d_model, + c.n_heads, + c.d_model * c.mlp_ratio, + c.dropout, + activation="gelu", + batch_first=True, + norm_first=True, + ) + self.transformer = nn.TransformerEncoder(layer, c.n_layers) + self.norm = nn.LayerNorm(c.d_model) + self.head = nn.Linear(c.d_model, c.vocab_size, bias=False) + self.head.weight = self.token_embedding.weight + + def forward( + self, + pixel_values: Tensor, + input_ids: Tensor, + history_pixel_values: Tensor | None = None, + history_states: Tensor | None = None, + current_state: Tensor | None = None, + history_mask: Tensor | None = None, + attention_mask: Tensor | None = None, + target_token: Tensor | None = None, + ) -> HighLevelOutput: + batch = pixel_values.shape[0] + if current_state is None: + current_state = pixel_values.new_zeros(batch, self.config.state_dim) + visual = self.memory_encoder(pixel_values, history_pixel_values, history_mask) + if history_states is None: + history_states = current_state.new_empty(batch, 0, self.config.state_dim) + states = self.state_projection(torch.cat((history_states, current_state[:, None]), 1)) + text = self.token_embedding(input_ids) + tokens = torch.cat((visual, states, text), dim=1) + observation_length = visual.shape[1] + states.shape[1] + length = tokens.shape[1] + q = torch.arange(length, device=tokens.device)[:, None] + k = torch.arange(length, device=tokens.device)[None] + allowed = torch.where(q < observation_length, k < observation_length, k <= q) + blocked = ~allowed + if attention_mask is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + state_valid = torch.ones(batch, states.shape[1], dtype=torch.bool, device=tokens.device) + if history_mask is not None and states.shape[1] > 1: + state_valid[:, :-1] = history_mask.bool() + valid = torch.cat( + ( + torch.ones(batch, visual.shape[1], dtype=torch.bool, device=tokens.device), + state_valid, + attention_mask.bool(), + ), + dim=1, + ) + hidden = self.transformer(tokens, mask=blocked, src_key_padding_mask=~valid) + text_hidden = self.norm(hidden[:, observation_length:]) + lengths = attention_mask.long().sum(dim=1).clamp_min(1) - 1 + pooled = text_hidden[torch.arange(batch, device=tokens.device), lengths] + logits = self.head(pooled) + loss = None if target_token is None else F.cross_entropy(logits, target_token.long()) + return HighLevelOutput(logits, loss) + + +class TinyGoalGenerator(nn.Module): + """Small conditional-flow image model standing in for the BAGEL world model.""" + + def __init__( + self, + vocab_size: int, + image_channels: int = 3, + hidden_dim: int = 48, + text_dim: int = 64, + ) -> None: + super().__init__() + self.image_channels = image_channels + self.token_embedding = nn.Embedding(vocab_size, text_dim) + self.text_projection = nn.Linear(text_dim, hidden_dim) + self.time_embedding = TimeEmbedding(hidden_dim) + self.input_projection = nn.Conv2d(2 * image_channels, hidden_dim, 3, padding=1) + self.network = nn.Sequential( + nn.SiLU(), + nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), + nn.SiLU(), + nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), + nn.SiLU(), + nn.Conv2d(hidden_dim, image_channels, 3, padding=1), + ) + + def forward( + self, + current_pixel_values: Tensor, + input_ids: Tensor, + noisy_subgoal: Tensor, + timesteps: Tensor, + attention_mask: Tensor | None = None, + target_velocity: Tensor | None = None, + ) -> GoalGeneratorOutput: + if current_pixel_values.shape != noisy_subgoal.shape: + raise ValueError("current image and noisy subgoal must have identical shapes") + if current_pixel_values.ndim != 4 or current_pixel_values.shape[1] != self.image_channels: + raise ValueError("goal-generator images have the wrong shape") + if attention_mask is None: + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + embedded = self.token_embedding(input_ids) + weights = attention_mask.to(embedded.dtype)[..., None] + pooled = (embedded * weights).sum(1) / weights.sum(1).clamp_min(1) + condition = self.text_projection(pooled) + self.time_embedding(timesteps) + hidden = self.input_projection( + torch.cat((current_pixel_values, noisy_subgoal), dim=1) + ) + hidden = hidden + condition[:, :, None, None].to(hidden.dtype) + vector_field = self.network(hidden) + loss = None + if target_velocity is not None: + loss = F.mse_loss(vector_field, target_velocity) + return GoalGeneratorOutput(vector_field, loss) diff --git a/src/pi07mini/rtc.py b/src/pi07mini/rtc.py new file mode 100644 index 0000000..f0297f0 --- /dev/null +++ b/src/pi07mini/rtc.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +MAX_RTC_DELAY_STEPS = 12 + + +@dataclass(frozen=True) +class RTCAlignment: + """A fixed-width target aligned to an inference/control delay.""" + + delayed_target_actions: torch.Tensor + valid_mask: torch.Tensor + executed_prefix_actions: torch.Tensor + delay_steps: int + + +def align_actions_for_rtc( + actions: torch.Tensor, + delay_steps: int, + *, + max_delay_steps: int = MAX_RTC_DELAY_STEPS, +) -> RTCAlignment: + """Shift targets left so prediction step zero starts after control delay. + + For delay ``d``, target position ``i`` supervises original action ``i+d``. + Vacated tail positions are zero padded and excluded by ``valid_mask``. + The already-executed prefix is returned for optional action-prefix input. + """ + + actions = torch.as_tensor(actions, dtype=torch.float32) + if actions.ndim < 2 or actions.shape[-2] < 1: + raise ValueError("actions must have shape [..., horizon, action_dim]") + if not 0 <= delay_steps <= max_delay_steps: + raise ValueError(f"delay_steps must lie in [0, {max_delay_steps}]") + horizon = actions.shape[-2] + if delay_steps >= horizon: + raise ValueError("delay_steps must be smaller than the action horizon") + aligned = torch.zeros_like(actions) + valid = torch.zeros(actions.shape[:-1], dtype=torch.bool, device=actions.device) + remaining = horizon - delay_steps + aligned[..., :remaining, :] = actions[..., delay_steps:, :] + valid[..., :remaining] = True + prefix = actions[..., :delay_steps, :].clone() + return RTCAlignment(aligned, valid, prefix, int(delay_steps)) + + +def sample_rtc_delay( + *, + generator: torch.Generator | None = None, + max_delay_steps: int = MAX_RTC_DELAY_STEPS, + action_horizon: int, +) -> int: + if action_horizon < 1: + raise ValueError("action_horizon must be positive") + maximum = min(int(max_delay_steps), action_horizon - 1) + if maximum < 0: + raise ValueError("no valid delay for an empty action horizon") + return int(torch.randint(maximum + 1, (), generator=generator).item()) diff --git a/src/pi07mini/runtime.py b/src/pi07mini/runtime.py new file mode 100644 index 0000000..d88b554 --- /dev/null +++ b/src/pi07mini/runtime.py @@ -0,0 +1,14 @@ +from __future__ import annotations +import json,random +from pathlib import Path +import numpy as np,torch,yaml +def load_config(path): + with Path(path).open() as f:return yaml.safe_load(f) +def seed_everything(seed:int):random.seed(seed);np.random.seed(seed);torch.manual_seed(seed) +def choose_device(name="auto"): + 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,value): + p=Path(path);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(value,indent=2,ensure_ascii=False)) diff --git a/src/pi07mini/tokenizer.py b/src/pi07mini/tokenizer.py new file mode 100644 index 0000000..47392f0 --- /dev/null +++ b/src/pi07mini/tokenizer.py @@ -0,0 +1,393 @@ +"""Prompt and FAST-lite tokenization for the minimal pi0.7 policy. + +pi0.7 conditions the low-level policy on a richer prompt than earlier pi +models: task and subtask language, episode speed/quality/mistake metadata, +control mode, and data provenance. This tokenizer represents those concepts +with explicit marker/value pairs so the architecture can be validated without +shipping a full language tokenizer. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Iterable, Literal + +import torch +from torch import Tensor + + +ControlMode = Literal["joint", "ee"] + + +@dataclass(frozen=True) +class EpisodeMetadata: + speed_steps: int | None = None + quality: int | None = None + mistake: bool | None = None + control_mode: ControlMode = "joint" + source_type: str | None = None + + def __post_init__(self) -> None: + if self.speed_steps is not None and self.speed_steps <= 0: + raise ValueError("speed_steps must be positive") + if self.quality is not None and not 1 <= self.quality <= 5: + raise ValueError("quality must be in the inclusive range 1..5") + if self.control_mode not in ("joint", "ee"): + raise ValueError("control_mode must be 'joint' or 'ee'") + + +@dataclass(frozen=True) +class PolicySequence: + input_ids: Tensor + labels: Tensor + action_prefix_mask: Tensor + generated_token_start: int + action_start_index: int + prompt_length: int + metadata_indices: tuple[int, ...] + + +class _LabelVocabulary: + def __init__(self, labels: tuple[str, ...], first_token_id: int) -> None: + self.labels = labels + self.first_token_id = first_token_id + self._to_token = { + label: first_token_id + index for index, label in enumerate(labels) + } + self._to_label = {token: label for label, token in self._to_token.items()} + + @property + def token_ids(self) -> tuple[int, ...]: + return tuple(self._to_token[label] for label in self.labels) + + def encode(self, label: str) -> int: + try: + return self._to_token[label] + except KeyError as error: + raise ValueError(f"unknown categorical label: {label}") from error + + def decode(self, token: int | Tensor) -> str: + try: + return self._to_label[int(token)] + except KeyError as error: + raise ValueError(f"unknown categorical token: {int(token)}") from error + + +class MiniTokenizer: + """Explicit rich-prompt tokens plus a DCT-based FAST-lite action stream.""" + + BOS = 256 + TASK_MARKER = 257 + SUBTASK_MARKER = 258 + SPEED_MARKER = 259 + QUALITY_MARKER = 260 + MISTAKE_MARKER = 261 + CONTROL_MARKER = 262 + SOURCE_MARKER = 263 + + TASK_LABELS = ( + "drive", + "pick_place", + "fold", + "open_close", + "clean", + "cook", + "unknown", + ) + TASK_TOKEN_START = 264 + TASK_TOKEN_END = TASK_TOKEN_START + len(TASK_LABELS) - 1 + + SUBTASK_LABELS = ( + "keep_lane", + "steer_left", + "steer_right", + "slow_down", + "stop", + "accelerate", + "pick_place", + "fold", + "open_close", + "unknown", + ) + SUBTASK_TOKEN_START = TASK_TOKEN_END + 1 + SUBTASK_TOKEN_END = SUBTASK_TOKEN_START + len(SUBTASK_LABELS) - 1 + + SPEED_BINS = (500, 1000, 2000, 4000, 8000, 16000) + SPEED_TOKEN_START = SUBTASK_TOKEN_END + 1 + SPEED_TOKEN_END = SPEED_TOKEN_START + len(SPEED_BINS) - 1 + QUALITY_TOKEN_START = SPEED_TOKEN_END + 1 + QUALITY_TOKEN_END = QUALITY_TOKEN_START + 4 + MISTAKE_FALSE = QUALITY_TOKEN_END + 1 + MISTAKE_TRUE = MISTAKE_FALSE + 1 + CONTROL_JOINT = MISTAKE_TRUE + 1 + CONTROL_EE = CONTROL_JOINT + 1 + + SOURCE_LABELS = ( + "demonstration", + "autonomous", + "intervention", + "human_video", + "web", + "unknown", + ) + SOURCE_TOKEN_START = CONTROL_EE + 1 + SOURCE_TOKEN_END = SOURCE_TOKEN_START + len(SOURCE_LABELS) - 1 + + ACTION_START = SOURCE_TOKEN_END + 1 + ACTION_END = ACTION_START + 1 + ACTION_TOKEN_START = ACTION_END + 1 + ACTION_BINS = 64 + ACTION_TOKEN_END = ACTION_TOKEN_START + ACTION_BINS - 1 + ACTION_MINIMUM = -4.0 + ACTION_MAXIMUM = 4.0 + + def __init__(self) -> None: + self.tasks = _LabelVocabulary(self.TASK_LABELS, self.TASK_TOKEN_START) + self.subtasks = _LabelVocabulary(self.SUBTASK_LABELS, self.SUBTASK_TOKEN_START) + self.sources = _LabelVocabulary(self.SOURCE_LABELS, self.SOURCE_TOKEN_START) + + @property + def vocab_size(self) -> int: + return self.ACTION_TOKEN_END + 1 + + @property + def action_token_ids(self) -> range: + return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1) + + @staticmethod + def _classify_task(description: str) -> str: + text = description.casefold() + if any(word in text for word in ("drive", "lane", "road", "steer", "vehicle")): + return "drive" + if any(word in text for word in ("fold", "shirt", "towel", "laundry")): + return "fold" + if any(word in text for word in ("open", "close", "door", "drawer", "fridge")): + return "open_close" + if any(word in text for word in ("clean", "wipe", "wash", "tidy")): + return "clean" + if any(word in text for word in ("cook", "food", "recipe", "microwave")): + return "cook" + if any(word in text for word in ("pick", "place", "move", "put")): + return "pick_place" + return description if description in MiniTokenizer.TASK_LABELS else "unknown" + + @staticmethod + def _classify_subtask(description: str) -> str: + if description in MiniTokenizer.SUBTASK_LABELS: + return description + text = description.casefold() + if any(word in text for word in ("stop", "halt", "stationary")): + return "stop" + if any(word in text for word in ("slow", "brake", "deceler")): + return "slow_down" + if any(word in text for word in ("accelerat", "speed up")): + return "accelerate" + if "left" in text: + return "steer_left" + if "right" in text: + return "steer_right" + if any(word in text for word in ("lane", "straight", "continue", "follow")): + return "keep_lane" + if "fold" in text: + return "fold" + if any(word in text for word in ("open", "close")): + return "open_close" + if any(word in text for word in ("pick", "place", "put")): + return "pick_place" + return "unknown" + + @staticmethod + def _classify_source(source: str) -> str: + if source in MiniTokenizer.SOURCE_LABELS: + return source + text = source.casefold() + if any(word in text for word in ("demo", "teleop")): + return "demonstration" + if any(word in text for word in ("autonomous", "rollout", "policy")): + return "autonomous" + if any(word in text for word in ("intervention", "correction")): + return "intervention" + if any(word in text for word in ("human", "ego", "video")): + return "human_video" + if any(word in text for word in ("web", "internet")): + return "web" + return "unknown" + + def encode_prompt( + self, + task: str, + subtask: str | None, + metadata: EpisodeMetadata, + *, + include_episode_metadata: bool = True, + include_subtask: bool = True, + ) -> tuple[Tensor, tuple[int, ...]]: + """Encode prompt and return positions controlled by metadata CFG. + + Control mode is deliberately retained even in the unconditional branch, + matching pi0.7 training where control-mode dropout is not applied. + """ + + ids = [self.BOS, self.TASK_MARKER, self.tasks.encode(self._classify_task(task))] + if include_subtask and subtask is not None: + ids.extend( + (self.SUBTASK_MARKER, self.subtasks.encode(self._classify_subtask(subtask))) + ) + + metadata_indices: list[int] = [] + + def add_metadata(marker: int, value: int) -> None: + metadata_indices.extend((len(ids), len(ids) + 1)) + ids.extend((marker, value)) + + if include_episode_metadata: + if metadata.speed_steps is not None: + bin_index = min( + range(len(self.SPEED_BINS)), + key=lambda index: abs(self.SPEED_BINS[index] - metadata.speed_steps), + ) + add_metadata(self.SPEED_MARKER, self.SPEED_TOKEN_START + bin_index) + if metadata.quality is not None: + add_metadata( + self.QUALITY_MARKER, self.QUALITY_TOKEN_START + metadata.quality - 1 + ) + if metadata.mistake is not None: + add_metadata( + self.MISTAKE_MARKER, + self.MISTAKE_TRUE if metadata.mistake else self.MISTAKE_FALSE, + ) + if metadata.source_type is not None: + source = self._classify_source(metadata.source_type) + add_metadata(self.SOURCE_MARKER, self.sources.encode(source)) + + control = self.CONTROL_JOINT if metadata.control_mode == "joint" else self.CONTROL_EE + ids.extend((self.CONTROL_MARKER, control)) + return torch.tensor(ids, dtype=torch.long), tuple(metadata_indices) + + @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 FAST-lite 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, + task: str, + subtask: str | None, + actions: Tensor, + metadata: EpisodeMetadata, + *, + include_episode_metadata: bool = True, + include_subtask: bool = True, + ) -> PolicySequence: + prompt, metadata_indices = self.encode_prompt( + task, + subtask, + metadata, + include_episode_metadata=include_episode_metadata, + include_subtask=include_subtask, + ) + action_tokens = self.encode_actions(actions).cpu() + input_ids = torch.cat((prompt, action_tokens)) + action_start_index = prompt.numel() + labels = torch.full_like(input_ids, -100) + 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, + action_prefix_mask=action_prefix_mask, + generated_token_start=0, + action_start_index=action_start_index, + prompt_length=prompt.numel(), + metadata_indices=metadata_indices, + ) + + def build_metadata_cfg_pair( + self, + task: str, + subtask: str | None, + actions: Tensor, + metadata: EpisodeMetadata, + ) -> tuple[PolicySequence, PolicySequence]: + """Return conditioned and episode-metadata-dropped policy sequences.""" + + conditioned = self.build_policy_sequence(task, subtask, actions, metadata) + unconditional = self.build_policy_sequence( + task, + subtask, + actions, + metadata, + include_episode_metadata=False, + ) + return conditioned, unconditional + diff --git a/src/pi07mini/train.py b/src/pi07mini/train.py new file mode 100644 index 0000000..df6c8c3 --- /dev/null +++ b/src/pi07mini/train.py @@ -0,0 +1,40 @@ +from __future__ import annotations +import argparse,copy,json +from pathlib import Path +import torch +from .data import Pi07SteerableDataset,materialize_subset +from .flow import FlowMatching +from .model import Pi07Config,Pi07Mini,HighLevelPolicy,TinyGoalGenerator +from .runtime import load_config,seed_everything,choose_device,write_json +from .tokenizer import MiniTokenizer,EpisodeMetadata + +def build_config(cfg,tok): + m=cfg['model'];return Pi07Config(vocab_size=tok.vocab_size,action_dim=2,state_dim=2,image_size=cfg['data']['image_size'],patch_size=m['patch_size'],d_model=m['d_model'],n_heads=m['n_heads'],n_layers=m['n_layers'],mlp_ratio=m['mlp_ratio'],max_history_frames=cfg['data']['history_frames'],max_text_tokens=m['max_text_tokens'],max_action_tokens=cfg['data']['action_horizon'],action_start_token_id=tok.ACTION_START) + +def train(config_path): + cfg=load_config(config_path);root=Path(__file__).resolve().parents[2];seed_everything(cfg['seed']);dev=choose_device(cfg['device']) + manifest=materialize_subset(cfg['data']['source_manifest'],root/'data/overfit',max_samples=cfg['data']['max_samples']) + ds=Pi07SteerableDataset(manifest,image_size=cfg['data']['image_size'],history_frames=cfg['data']['history_frames'],action_horizon=cfg['data']['action_horizon'],augmentation_factor=cfg['data']['augmentation_factor']) + tok=MiniTokenizer();pcfg=build_config(cfg,tok);policy=Pi07Mini(pcfg).to(dev);high=HighLevelPolicy(pcfg).to(dev);goal=TinyGoalGenerator(tok.vocab_size).to(dev);flow=FlowMatching(cfg['flow']['integration_steps']);out=root/cfg['output_dir'];out.mkdir(parents=True,exist_ok=True);history=[] + # Goal generator + opt=torch.optim.AdamW(goal.parameters(),lr=cfg['train']['learning_rate']) + for step in range(1,cfg['train']['goal_steps']+1): + s=ds[(step-1)%len(ds)];meta=EpisodeMetadata(**s['token_metadata']);seq=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta);ids=seq.input_ids[:seq.action_start_index][None].to(dev);target=s['subgoal_pixel_values'][None].to(dev);t=torch.rand(1,device=dev);noise=torch.randn_like(target);noisy=t[:,None,None,None]*target+(1-t[:,None,None,None])*noise;o=goal(s['image'][None].to(dev),ids,noisy,t,target_velocity=target-noise);opt.zero_grad();o.loss.backward();opt.step() + if step==1 or step%cfg['train']['log_every']==0:row={'stage':'goal','step':step,'loss':float(o.loss.detach().cpu())};history.append(row);print(json.dumps(row)) + # High-level semantic policy + opt=torch.optim.AdamW(high.parameters(),lr=cfg['train']['learning_rate']) + for step in range(1,cfg['train']['high_level_steps']+1): + s=ds[(step-1)%len(ds)];meta=EpisodeMetadata(**s['token_metadata']);prompt,_=tok.encode_prompt('drive safely',None,meta,include_subtask=False);target=torch.tensor([tok.subtasks.encode(tok._classify_subtask(s['reasoning']))],device=dev);o=high(s['image'][None].to(dev),prompt[None].to(dev),s['history_images'][None].to(dev),s['history_states'][None].to(dev),s['current_state'][None].to(dev),target_token=target);opt.zero_grad();o.loss.backward();opt.step() + if step==1 or step%cfg['train']['log_every']==0:row={'stage':'high_level','step':step,'loss':float(o.loss.detach().cpu())};history.append(row);print(json.dumps(row)) + # Full steerable policy, alternating prompt-dropout ablations and RTC targets. + opt=torch.optim.AdamW(policy.parameters(),lr=cfg['train']['learning_rate']) + for step in range(1,cfg['train']['policy_steps']+1): + s=ds[(step-1)%len(ds)];meta=EpisodeMetadata(**s['token_metadata']);include_meta=(step%4)!=0;seq=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=include_meta);ids=seq.input_ids[None].to(dev);labels=seq.labels[None].to(dev);target=s['rtc_target_actions'][None].to(dev);noisy,t,vel=flow.training_batch(target);use_goal=(step%4)!=1;use_history=(step%5)!=1 + o=policy(s['image'][None].to(dev),ids,noisy,t,history_pixel_values=s['history_images'][None].to(dev) if use_history else None,history_states=s['history_states'][None].to(dev) if use_history else None,current_state=s['current_state'][None].to(dev),subgoal_pixel_values=s['subgoal_pixel_values'][None].to(dev) if use_goal else None,action_prefix_mask=seq.action_prefix_mask[None].to(dev),text_labels=labels) + valid=s['rtc_valid_mask'][None].to(dev)[:,:,None];flow_loss=((o.vector_field-vel).square()*valid).sum()/valid.sum().clamp_min(1);loss=o.text_loss+cfg['train']['flow_loss_weight']*flow_loss;opt.zero_grad();loss.backward();torch.nn.utils.clip_grad_norm_(policy.parameters(),cfg['train']['grad_clip_norm']);opt.step() + if step==1 or step%cfg['train']['log_every']==0:row={'stage':'policy','step':step,'loss':float(loss.detach().cpu()),'text_loss':float(o.text_loss.detach().cpu()),'flow_loss':float(flow_loss.detach().cpu())};history.append(row);print(json.dumps(row)) + ck={'policy':policy.state_dict(),'high_level':high.state_dict(),'goal':goal.state_dict(),'policy_config':pcfg.__dict__,'config':cfg,'action_statistics':ds.action_normalizer.as_dict()};torch.save(ck,out/'checkpoint.pt');write_json(out/'training_history.json',history);write_json(out/'summary.json',{'device':str(dev),'samples':len(ds),'policy_parameters':sum(p.numel() for p in policy.parameters()),'goal_parameters':sum(p.numel() for p in goal.parameters()),'high_level_parameters':sum(p.numel() for p in high.parameters())});return out/'checkpoint.pt' + +def main(): + p=argparse.ArgumentParser();p.add_argument('--config',default='configs/overfit.yaml');train(p.parse_args().config) +if __name__=='__main__':main() diff --git a/tests/test_context.py b/tests/test_context.py new file mode 100644 index 0000000..0951622 --- /dev/null +++ b/tests/test_context.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import torch + +from pi07mini.context import ( + PolicyMetadata, + build_ambiguous_history_pair, + drop_prompt_components, + endpoint_subgoal_heatmap, + format_policy_prompt, + make_history_proxy_frames, + prompt_components, +) + + +def _history() -> torch.Tensor: + x = torch.linspace(-8.0, 0.0, 8) + return torch.stack((x, 0.04 * x.square()), -1) + + +def test_subgoal_heatmap_depends_on_endpoint_only() -> None: + endpoint = torch.tensor([20.0, 4.0]) + first = endpoint_subgoal_heatmap(endpoint, size=24) + # Two future paths with radically different interiors produce exactly the + # same condition because only their shared endpoint is passed to the API. + path_a = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1) + path_b = path_a.clone() + path_b[1:-1, 1] += 100.0 + second = endpoint_subgoal_heatmap(path_b[-1], size=24) + torch.testing.assert_close(first, second, atol=0.0, rtol=0.0) + assert first.shape == (1, 24, 24) + assert float(first.max()) == 1.0 + assert torch.isfinite(first).all() + + +def test_history_proxy_is_oldest_to_newest_and_future_free() -> None: + image = torch.linspace(0.0, 1.0, 3 * 20 * 20).reshape(3, 20, 20) + frames = make_history_proxy_frames(image, _history(), frames=6) + assert frames.shape == (6, 3, 20, 20) + torch.testing.assert_close(frames[-1], image * 0.98) + assert not torch.equal(frames[0], frames[-1]) + + +def test_ambiguous_histories_are_distinct_but_share_current_state() -> None: + pair = build_ambiguous_history_pair(_history(), frames=6) + assert pair.shape == (2, 6, 2) + torch.testing.assert_close(pair[0, -1], pair[1, -1]) + assert not torch.allclose(pair[0, :-1], pair[1, :-1]) + + +def test_prompt_component_dropout_is_independent_and_deterministic() -> None: + metadata = PolicyMetadata("fast", "good", "none", "autonomous") + components = prompt_components(metadata, task="Drive safely.") + generator_a = torch.Generator().manual_seed(42) + generator_b = torch.Generator().manual_seed(42) + visible_a, mask_a = drop_prompt_components( + components, drop_probability=0.5, generator=generator_a + ) + visible_b, mask_b = drop_prompt_components( + components, drop_probability=0.5, generator=generator_b + ) + assert visible_a == visible_b + torch.testing.assert_close(mask_a, mask_b) + assert visible_a["task"] == "Drive safely." + assert mask_a.shape == (5,) + text = format_policy_prompt(visible_a) + assert "Drive safely." in text + for dropped in set(components) - set(visible_a): + assert f"{dropped}=" not in text + + +def test_full_metadata_dropout_keeps_only_task() -> None: + components = prompt_components(PolicyMetadata("slow", "poor", "stop", "autonomous")) + visible, mask = drop_prompt_components(components, drop_probability=1.0) + assert set(visible) == {"task"} + assert mask.tolist() == [True, False, False, False, False] diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..f5a7c1c --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + +from pi07mini.data import ( + Pi07SteerableDataset, + build_policy_trajectories, + 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_policy_variants_and_metadata() -> None: + x = torch.linspace(-7.0, 0.0, 8) + history = torch.stack((x, torch.zeros_like(x)), -1) + fx = torch.linspace(1.0, 24.0, 16) + future = torch.stack((fx, 0.02 * fx.square()), -1) + policies = build_policy_trajectories(history, future) + assert [policy.kind for policy in policies] == [ + "expert", + "fast", + "slow", + "understeer", + "stop", + ] + metadata = {policy.kind: policy.metadata.as_dict() for policy in policies} + assert metadata["expert"] == { + "speed": "normal", + "quality": "expert", + "mistake": "none", + "control": "human", + } + assert metadata["understeer"]["mistake"] == "understeer" + assert metadata["stop"]["quality"] == "poor" + assert policies[0].metadata.as_token_metadata() == { + "speed_steps": 1_000, + "quality": 5, + "mistake": False, + "control_mode": "ee", + "source_type": "demonstration", + } + assert torch.allclose(policies[-1].future_xy[-1], policies[-1].future_xy[4]) + + +def test_dataset_contract_and_roundtrip(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 = Pi07SteerableDataset( + manifest, + image_size=32, + prompt_dropout=0.0, + max_rtc_delay_steps=12, + ) + assert len(dataset) == 5 + assert {policy.kind for policy in dataset.policies} == { + "expert", + "fast", + "slow", + "understeer", + "stop", + } + item = dataset[0] + assert item["pixel_values"].shape == (3, 32, 32) + assert item["history_pixel_values"].shape == (6, 3, 32, 32) + assert item["history_states"].shape == (6, 2) + assert item["current_state"].shape == (2,) + assert item["actions"].shape == (16, 2) + assert item["subgoal_heatmap"].shape == (1, 32, 32) + assert item["subgoal_pixel_values"].shape == (3, 32, 32) + assert item["ambiguous_history_pair"].shape == (2, 6, 2) + assert 0 <= item["rtc_delay_steps"] <= 12 + assert item["rtc_target_actions"].shape == (16, 2) + assert item["rtc_valid_mask"].shape == (16,) + assert item["reasoning"] == "Steer left." + assert item["prompt_component_mask"].all() + assert item["token_metadata"]["quality"] == 5 + assert item["control_mode"] == "ee" + decoded = dataset.action_normalizer.decode(item["actions"], item["history_xy"]) + torch.testing.assert_close(decoded, item["future_xy"], atol=1e-5, rtol=1e-5) + + +def test_dataset_is_deterministic_per_index(tmp_path: Path) -> None: + source = _make_source_project(tmp_path / "source") + dataset = Pi07SteerableDataset( + source, + image_size=24, + prompt_dropout=0.5, + seed=11, + ) + first = dataset[3] + second = dataset[3] + assert first["prompt_components"] == second["prompt_components"] + assert first["rtc_delay_steps"] == second["rtc_delay_steps"] + torch.testing.assert_close(first["rtc_target_actions"], second["rtc_target_actions"]) diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 0000000..f41bf62 --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,63 @@ +import torch + +from pi07mini.memory import MEMConfig, MEMVideoEncoder + + +def tiny_encoder() -> MEMVideoEncoder: + return MEMVideoEncoder( + MEMConfig( + image_size=16, + patch_size=8, + d_model=16, + n_heads=4, + n_layers=2, + max_history_frames=6, + ) + ).eval() + + +def test_mem_output_token_count_is_independent_of_history_length() -> None: + encoder = tiny_encoder() + current = torch.randn(2, 3, 16, 16) + with torch.no_grad(): + no_history = encoder(current) + long_history = encoder(current, torch.randn(2, 6, 3, 16, 16)) + assert no_history.shape == (2, 4, 16) + assert long_history.shape == no_history.shape + + +def test_mem_is_sensitive_to_temporal_order() -> None: + torch.manual_seed(7) + encoder = tiny_encoder() + current = torch.randn(1, 3, 16, 16) + history = torch.stack( + ( + torch.zeros(3, 16, 16), + torch.ones(3, 16, 16), + torch.full((3, 16, 16), 2.0), + ) + )[None] + with torch.no_grad(): + chronological = encoder(current, history) + reversed_history = encoder(current, history.flip(1)) + assert not torch.allclose(chronological, reversed_history) + + +def test_masked_history_frame_cannot_change_memory() -> None: + torch.manual_seed(9) + encoder = tiny_encoder() + current = torch.randn(1, 3, 16, 16) + history = torch.randn(1, 2, 3, 16, 16) + changed = history.clone() + changed[:, 0] = 1000.0 + mask = torch.tensor([[False, True]]) + with torch.no_grad(): + first = encoder(current, history, mask) + second = encoder(current, changed, mask) + torch.testing.assert_close(first, second) + + +def test_current_frame_temporal_position_is_exactly_zero() -> None: + encoding = MEMVideoEncoder.temporal_position(torch.tensor([3, 2, 1, 0]), 15) + assert torch.count_nonzero(encoding[-1]) == 0 + diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..01acfed --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,180 @@ +import torch + +from pi07mini.model import ( + HighLevelPolicy, + Pi07Config, + Pi07Mini, + PrefixLayout, + TinyGoalGenerator, +) +from pi07mini.tokenizer import EpisodeMetadata, MiniTokenizer + + +def tiny_model() -> Pi07Mini: + tokenizer = MiniTokenizer() + return Pi07Mini( + Pi07Config( + vocab_size=tokenizer.vocab_size, + image_size=16, + patch_size=16, + d_model=32, + n_heads=4, + n_layers=1, + memory_layers=1, + max_history_frames=3, + max_text_tokens=64, + max_action_tokens=4, + action_start_token_id=tokenizer.ACTION_START, + ) + ) + + +def sample_sequence(): + tokenizer = MiniTokenizer() + metadata = EpisodeMetadata(2000, 5, False, "ee", "demonstration") + return tokenizer, tokenizer.build_policy_sequence( + "drive", "steer_left", torch.zeros(2, 2), metadata + ) + + +def common_inputs(sequence): + return dict( + pixel_values=torch.randn(1, 3, 16, 16), + history_pixel_values=torch.randn(1, 2, 3, 16, 16), + history_states=torch.randn(1, 2, 2), + current_state=torch.randn(1, 2), + subgoal_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], + ) + + +def test_block_causal_attention_matches_pi07_modal_order() -> None: + layout = PrefixLayout(observation_tokens=2, goal_tokens=2, text_tokens=3) + valid = torch.ones(1, layout.total_tokens, dtype=torch.bool) + edges = Pi07Mini.build_prefix_attention_mask(layout, valid)[0] + assert edges[0, :2].all() + assert not edges[0, 2:].any() # observation cannot read goal or text + assert edges[2, :4].all() + assert not edges[2, 4:].any() # goal reads observation + goal, not text + assert edges[4, :5].all() + assert not edges[4, 5:].any() # first text token is causal + assert edges[6].all() + + +def test_prefix_has_fixed_mem_tokens_and_one_token_per_state() -> None: + _, sequence = sample_sequence() + model = tiny_model() + inputs = common_inputs(sequence) + prefix, _, layout = model.encode_prefix( + inputs["pixel_values"], + inputs["history_pixel_values"], + inputs["history_states"], + inputs["current_state"], + inputs["input_ids"], + inputs["subgoal_pixel_values"], + ) + assert layout.observation_tokens == 1 + 2 + 1 # MEM + history states + current + assert layout.goal_tokens == 1 + assert prefix.shape[1] == layout.total_tokens + + +def test_flow_cannot_read_teacher_forced_fast_tokens() -> None: + torch.manual_seed(12) + tokenizer, sequence = sample_sequence() + model = tiny_model().eval() + changed_ids = sequence.input_ids.clone() + changed_ids[sequence.action_start_index + 1 : -1] = tokenizer.ACTION_TOKEN_START + inputs = common_inputs(sequence) + with torch.no_grad(): + first = model(**inputs).vector_field + inputs["input_ids"] = changed_ids[None] + second = model(**inputs).vector_field + torch.testing.assert_close(first, second) + + +def test_flow_loss_is_knowledge_insulated_from_all_vlm_parameters() -> None: + torch.manual_seed(13) + _, sequence = sample_sequence() + model = tiny_model() + output = model( + **common_inputs(sequence), target_velocity=torch.randn(1, 2, 2) + ) + assert output.flow_loss is not None + output.flow_loss.backward() + assert all( + parameter.grad is None or torch.count_nonzero(parameter.grad) == 0 + for parameter in model.vlm_parameters() + ) + assert any( + parameter.grad is not None and torch.count_nonzero(parameter.grad) > 0 + for parameter in model.action_parameters() + ) + + +def test_fast_cross_entropy_still_trains_vlm() -> None: + _, sequence = sample_sequence() + model = tiny_model() + output = model( + **common_inputs(sequence), text_labels=sequence.labels[None] + ) + assert output.text_loss is not None + output.text_loss.backward() + assert model.token_embedding.weight.grad is not None + assert torch.count_nonzero(model.token_embedding.weight.grad) > 0 + + +def test_visual_subgoal_conditions_flow_expert() -> None: + torch.manual_seed(14) + _, sequence = sample_sequence() + model = tiny_model().eval() + inputs = common_inputs(sequence) + with torch.no_grad(): + first = model(**inputs).vector_field + inputs["subgoal_pixel_values"] = inputs["subgoal_pixel_values"] + 3.0 + second = model(**inputs).vector_field + assert not torch.allclose(first, second) + + +def test_metadata_cfg_uses_paper_extrapolation_formula() -> None: + conditioned = torch.tensor([1.0, 2.0]) + unconditional = torch.tensor([0.5, 1.5]) + torch.testing.assert_close( + Pi07Mini.combine_metadata_cfg(conditioned, unconditional, 0.0), conditioned + ) + torch.testing.assert_close( + Pi07Mini.combine_metadata_cfg(conditioned, unconditional, 1.0), + torch.tensor([1.5, 2.5]), + ) + + +def test_high_level_policy_predicts_next_subtask_token() -> None: + model = HighLevelPolicy(tiny_model().config) + output = model( + pixel_values=torch.randn(2, 3, 16, 16), + history_pixel_values=torch.randn(2, 2, 3, 16, 16), + history_states=torch.randn(2, 2, 2), + current_state=torch.randn(2, 2), + input_ids=torch.tensor([[256, 257, 264], [256, 257, 264]]), + target_token=torch.tensor([271, 272]), + ) + assert output.logits.shape == (2, tiny_model().config.vocab_size) + assert output.loss is not None + + +def test_tiny_goal_generator_predicts_image_flow() -> None: + tokenizer = MiniTokenizer() + generator = TinyGoalGenerator(tokenizer.vocab_size) + current = torch.randn(2, 3, 16, 16) + noisy = torch.randn_like(current) + output = generator( + current, + torch.tensor([[256, 257, 264], [256, 257, 264]]), + noisy, + torch.rand(2), + target_velocity=torch.randn_like(current), + ) + assert output.vector_field.shape == current.shape + assert output.loss is not None diff --git a/tests/test_rtc.py b/tests/test_rtc.py new file mode 100644 index 0000000..c233e8a --- /dev/null +++ b/tests/test_rtc.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import pytest +import torch + +from pi07mini.rtc import align_actions_for_rtc, sample_rtc_delay + + +def _actions() -> torch.Tensor: + return torch.arange(32, dtype=torch.float32).reshape(16, 2) + + +def test_zero_delay_preserves_entire_chunk() -> None: + actions = _actions() + aligned = align_actions_for_rtc(actions, 0) + torch.testing.assert_close(aligned.delayed_target_actions, actions) + assert aligned.valid_mask.all() + assert aligned.executed_prefix_actions.shape == (0, 2) + + +def test_delay_shifts_target_and_masks_padded_tail() -> None: + actions = _actions() + aligned = align_actions_for_rtc(actions, 12) + torch.testing.assert_close(aligned.delayed_target_actions[:4], actions[12:]) + assert torch.count_nonzero(aligned.delayed_target_actions[4:]) == 0 + assert aligned.valid_mask.tolist() == [True] * 4 + [False] * 12 + torch.testing.assert_close(aligned.executed_prefix_actions, actions[:12]) + + +def test_rtc_alignment_preserves_batch_dimensions() -> None: + actions = torch.stack((_actions(), _actions() + 100.0)) + aligned = align_actions_for_rtc(actions, 3) + assert aligned.delayed_target_actions.shape == (2, 16, 2) + assert aligned.valid_mask.shape == (2, 16) + torch.testing.assert_close(aligned.delayed_target_actions[:, :13], actions[:, 3:]) + + +def test_delay_sampling_covers_closed_zero_to_twelve_range() -> None: + generator = torch.Generator().manual_seed(9) + values = { + sample_rtc_delay(generator=generator, action_horizon=16) for _ in range(300) + } + assert min(values) == 0 + assert max(values) == 12 + + +@pytest.mark.parametrize("delay", [-1, 13, 16]) +def test_invalid_delay_is_rejected(delay: int) -> None: + with pytest.raises(ValueError): + align_actions_for_rtc(_actions(), delay) diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py new file mode 100644 index 0000000..55f299c --- /dev/null +++ b/tests/test_tokenizer.py @@ -0,0 +1,65 @@ +import torch + +from pi07mini.tokenizer import EpisodeMetadata, MiniTokenizer + + +def metadata() -> EpisodeMetadata: + return EpisodeMetadata( + speed_steps=2100, + quality=5, + mistake=False, + control_mode="ee", + source_type="autonomous_rollout", + ) + + +def test_prompt_contains_rich_metadata_and_control_tokens() -> None: + tokenizer = MiniTokenizer() + prompt, metadata_indices = tokenizer.encode_prompt( + "drive on the road", "steer_left", metadata() + ) + assert tokenizer.SPEED_MARKER in prompt + assert tokenizer.QUALITY_MARKER in prompt + assert tokenizer.MISTAKE_MARKER in prompt + assert tokenizer.SOURCE_MARKER in prompt + assert tokenizer.CONTROL_MARKER in prompt + assert tokenizer.CONTROL_EE in prompt + assert len(metadata_indices) == 8 + + +def test_metadata_cfg_pair_drops_metadata_but_retains_control_mode() -> None: + tokenizer = MiniTokenizer() + actions = torch.zeros(3, 2) + conditioned, unconditional = tokenizer.build_metadata_cfg_pair( + "drive", "keep_lane", actions, metadata() + ) + assert tokenizer.QUALITY_MARKER in conditioned.input_ids + assert tokenizer.QUALITY_MARKER not in unconditional.input_ids + assert tokenizer.CONTROL_EE in conditioned.input_ids + assert tokenizer.CONTROL_EE in unconditional.input_ids + assert conditioned.metadata_indices + assert not unconditional.metadata_indices + + +def test_flow_prefix_ends_before_fast_stream() -> None: + tokenizer = MiniTokenizer() + sequence = tokenizer.build_policy_sequence( + "drive", "keep_lane", torch.zeros(4, 2), metadata() + ) + assert sequence.input_ids[sequence.action_start_index] == tokenizer.ACTION_START + assert sequence.action_prefix_mask[: sequence.action_start_index].all() + assert not sequence.action_prefix_mask[sequence.action_start_index :].any() + assert sequence.labels[sequence.action_start_index - 1] == 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 +