diff --git a/README.md b/README.md index 046811e..7e4f4e5 100644 --- a/README.md +++ b/README.md @@ -1,246 +1,125 @@ # π0.7 Minimal Validation -这是一个面向**结构验证和小样本过拟合**的 π0.7-style 教学项目。它在前面的 -π0.6/RECAP 最小实验之上，重点验证一条可解释的 steerable policy 链路：利用 -短时历史形成 memory，在高层预测视觉 subgoal，通过可独立缺省的 policy -metadata 控制动作风格，并在存在控制延迟时用 RTC 对齐动作 chunk。 +这是一个面向**结构验证和极小样本过拟合**的 π0.7 最小实现。目标不是复现参数量、预训练权重或论文指标，而是保留官方论文中的关键模块、模态顺序、注意力信息流和训练/推理闭环。 -本项目不是 Physical Intelligence 官方实现，也不能用于真实车辆或机器人控制。 +实现依据： -## 公开资料边界 +- **π0.7: a Steerable Generalist Robotic Foundation Model with Emergent Capabilities**； +- **MEM: Multi-Scale Embodied Memory for Vision Language Action Models**； +- **Training-Time Action Conditioning for Efficient Real-Time Chunking**； +- π0.5/π0.6 的 FAST、flow action expert 与 Knowledge Insulation 结构。 -截至本项目实现时，作者没有在 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 ───────────────┐ +history images --MEM video encoder--┐ +history/current states -------------> observation block │ -current image + task ────────────────┼→ TinyGoalGenerator → predicted subgoal +task + predicted subtask + metadata -> causal text prompt │ -optional speed/quality/mistake/control metadata - ↓ - HighLevelPolicy - ↓ -noisy action + flow time + RTC prefix → Pi07Mini action expert - ↓ - continuous action chunk +generated future observation -------> image-goal block + │ +FAST action tokens -----------------> causal discrete action branch + │ +clean RTC prefix + noisy postfix ---> flow action expert ``` -## 重点结构 - -### 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 的开头： +官方 Figure 19 对应的 prefix 顺序被明确实现为： ```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 +memory observations -> causal text -> bidirectional image goals -> causal FAST ``` -默认随机采样 `d=0..12`，返回已执行 prefix，并对 padding tail 屏蔽损失。RTC 在 -这里表示训练/评估时的离线对齐机制，不包含实际控制器、异步推理线程或硬实时保证。 +- observation 只能读取 observation； +- text 读取 observation，并在 text 内因果注意； +- image goals 读取 observation、完整 text 和 goal block； +- FAST 读取此前全部上下文，并在 FAST 内因果注意； +- flow actions 读取 observation、text、image goals 和彼此，但不能读取 teacher-forced FAST； +- flow loss 通过 Knowledge Insulation 不更新 VLM/MEM stream。 -### 5. CFG +## High-level 到 low-level 的闭环 -推理时在相同 observation、goal、memory 和初始噪声下计算 metadata conditional -与 metadata-omitted 预测： +高层策略输入当前 observation、MEM history、总体任务、metadata，以及过去的 subtask history，输出下一条 subtask。推理时该**预测 subtask**同时用于： + +1. 条件化 visual subgoal generator； +2. 条件化低层 flow/FAST policy。 + +评估不再使用 ground-truth subtask 替代高层预测。 + +## Visual subgoal + +论文使用 segment 末端的多视角未来观测作为 world-model target。当前极小数据只有一张当前图像，因此本项目生成一个确定性的 `future_observation_proxy`：根据 segment 末端状态对当前相机图像进行小幅图像空间变换，得到与未来相机帧相同的 RGB 接口。 + +这比旧版 endpoint heatmap 更接近论文结构：low-level policy 接收的是生成的 RGB future observation，而不是直接编码的二维终点坐标。它仍只是缺少真实未来视频时的最小替代，不能证明真实图像 world modeling 能力。 + +## Training-time RTC + +训练严格采用 action-prefix conditioning： ```text -v_cfg = v_uncond + beta * (v_cond - v_uncond) +prefix positions: ground-truth clean action, flow time = 1, no loss +postfix positions: noisy action, sampled flow time, flow-matching loss ``` -默认 `beta=2.0`。必须同时报告 `beta=0,1,2`，因为高 guidance 可能放大过拟合、 -产生分布外动作，不能只展示最好看的曲线。 +推理时已经 committed 的 action prefix 在每个 denoising step 被重新写回，且对应 token 的 flow time 固定为 1。RTC 不再通过简单左移 target 来近似。 -## 数据与防泄漏约束 +## 官方 prompt dropout -- history proxy 和 ambiguous pair 只依赖当前图像与过去状态； -- subgoal condition 只使用 endpoint target，不能编码完整 action chunk； -- flow expert 训练时不能读取 teacher-forced 真值动作 token； -- 所有模型对比使用同一 initial noise、归一化统计和 RTC delay； -- 数据增强的随机数由 sample index 和 seed 决定，保证复现实验； -- 仓库不包含或重新分发原始 PAI 数据。 +- visual subgoal 出现在 25% policy examples； +- 有 visual subgoal 时，subtask 以 30% 概率移除； +- episode metadata 整体以 15% 概率移除； +- speed、quality、mistake 分别再以 5% 概率移除； +- control mode 始终保留； +- visual history 整体以 30% 概率移除。 -## 安装 +metadata CFG 使用论文公式： + +```text +v_cfg = v_cond + beta * (v_cond - v_uncond) +``` + +## 运行 ```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 -``` +1. high-quality segments 上的 visual goal flow； +2. next-subtask high-level policy； +3. 带 prompt dropout、FAST、flow、KI 和 training-time RTC 的低层 policy。 -最终展示不能只有 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 subtask； +- generated visual subgoal； +- full/no-metadata/no-goal/no-history/CFG 消融； +- delay `0/4/8/12` 的 RTC postfix action error； +- trajectory 和 normalized action MSE。 -还必须提供： +## 当前最小训练结果 -- 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 上的差异指标。 +默认配置在 Apple MPS 上完成 `800 / 200 / 900` 步训练： -小样本过拟合中，某个消融偶然更好并不证明模块无效；应展示全部固定设置结果， -不调换 noise seed 或只挑选最佳样本。 +- high-level subtask：`steer_left`，与目标一致； +- generated future-observation MSE：`0.0246`； +- full generated-goal normalized action MSE：`0.2848`； +- no-goal normalized action MSE：`0.2938`； +- no-metadata normalized action MSE：`0.3739`； +- RTC delay `8/12` 的 postfix action MSE：`0.1210 / 0.0874`。 -## 安全限制 +在这个极小数据实验中，metadata CFG 的 trajectory MSE 仍然退化，且 generated goal 的视觉细节较粗糙。结果被完整保留，不挑选更有利的 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 -编号或官方结果表。 +- 本项目只允许离线训练和可视化，严禁连接真实车辆或机器人执行器。 +- future observation 是单帧数据上的代理构造，不是真实多视角未来视频。 +- 小样本过拟合、loss 下降或轨迹重合不代表泛化、闭环稳定性或安全性。 +- 成功只说明论文中的模块接口和信息流能在缩小模型上形成可训练闭环。 diff --git a/configs/overfit.yaml b/configs/overfit.yaml index 5203fac..a6a871c 100644 --- a/configs/overfit.yaml +++ b/configs/overfit.yaml @@ -19,11 +19,9 @@ data: 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 + memory_layers: 2 mlp_ratio: 3 dropout: 0.0 max_text_tokens: 96 @@ -34,14 +32,12 @@ model: memory: enabled: true history_frames: 6 + history_dropout: 0.3 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 + target: future_observation_proxy high_level: enabled: true @@ -54,10 +50,11 @@ high_level: - control steering: - metadata_dropout: 0.25 - preserve_task_prompt: true - cfg_training_dropout: 0.25 - cfg_beta: 2.0 + subgoal_probability: 0.25 + subtask_dropout_with_goal: 0.3 + metadata_full_dropout: 0.15 + metadata_component_dropout: 0.05 + cfg_beta: 1.7 rtc: enabled: true @@ -68,10 +65,10 @@ rtc: flow: timestep_beta_alpha: 1.5 timestep_beta_beta: 1.0 - integration_steps: 10 + integration_steps: 5 train: - goal_steps: 300 + goal_steps: 800 high_level_steps: 200 policy_steps: 900 batch_size: 5 @@ -86,7 +83,8 @@ train: evaluation: fixed_noise_seed: 1 - cfg_betas: [0.0, 1.0, 2.0] + goal_flow_steps: 10 + cfg_betas: [0.0, 1.3, 1.7, 2.2] ablations: - full - no_metadata diff --git a/outputs/overfit/context_ablation.png b/outputs/overfit/context_ablation.png index 8d01cf1..d35268d 100644 Binary files a/outputs/overfit/context_ablation.png and b/outputs/overfit/context_ablation.png differ diff --git a/outputs/overfit/evaluation.json b/outputs/overfit/evaluation.json index 4b36128..02a14f1 100644 --- a/outputs/overfit/evaluation.json +++ b/outputs/overfit/evaluation.json @@ -1,11 +1,39 @@ { + "target_subtask": "steer_left", "predicted_subtask": "steer_left", + "goal_image_mse": 0.0246433112770319, "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 + "full_generated_goal": 9.339226722717285, + "full_oracle_goal": 9.344671249389648, + "no_metadata": 4.243795394897461, + "no_goal": 9.066949844360352, + "no_history": 10.404154777526855, + "metadata_cfg": 19.987443923950195 + }, + "normalized_action_metrics": { + "full_generated_goal": 0.2848391532897949, + "full_oracle_goal": 0.2848338186740875, + "no_metadata": 0.37385478615760803, + "no_goal": 0.29379916191101074, + "no_history": 0.28792211413383484, + "metadata_cfg": 0.2725481688976288 + }, + "rtc_metrics": { + "0": { + "postfix_action_mse": 0.2848391532897949, + "trajectory_mse": 9.339226722717285 + }, + "4": { + "postfix_action_mse": 1.2005467414855957, + "trajectory_mse": 2.2003140449523926 + }, + "8": { + "postfix_action_mse": 0.12100008875131607, + "trajectory_mse": 0.10110154747962952 + }, + "12": { + "postfix_action_mse": 0.08735787123441696, + "trajectory_mse": 0.11811599135398865 + } } } \ No newline at end of file diff --git a/outputs/overfit/goal_generation.png b/outputs/overfit/goal_generation.png index a38255f..5dc9537 100644 Binary files a/outputs/overfit/goal_generation.png and b/outputs/overfit/goal_generation.png differ diff --git a/outputs/overfit/summary.json b/outputs/overfit/summary.json index 1f6e681..503e491 100644 --- a/outputs/overfit/summary.json +++ b/outputs/overfit/summary.json @@ -1,7 +1,7 @@ { "device": "mps", "samples": 5, - "policy_parameters": 1088930, - "goal_parameters": 75731, - "high_level_parameters": 652800 + "policy_parameters": 1088354, + "goal_parameters": 75347, + "high_level_parameters": 652224 } \ No newline at end of file diff --git a/outputs/overfit/training_history.json b/outputs/overfit/training_history.json index 01f5f9d..0b20283 100644 --- a/outputs/overfit/training_history.json +++ b/outputs/overfit/training_history.json @@ -2,87 +2,212 @@ { "stage": "goal", "step": 1, - "loss": 1.0040572881698608 + "loss": 1.1256178617477417 }, { "stage": "goal", "step": 20, - "loss": 0.34230467677116394 + "loss": 0.8757796287536621 }, { "stage": "goal", "step": 40, - "loss": 0.37451064586639404 + "loss": 0.4848751425743103 }, { "stage": "goal", "step": 60, - "loss": 0.26792070269584656 + "loss": 0.2024349421262741 }, { "stage": "goal", "step": 80, - "loss": 0.10219617933034897 + "loss": 0.31732437014579773 }, { "stage": "goal", "step": 100, - "loss": 0.2711445987224579 + "loss": 0.22155670821666718 }, { "stage": "goal", "step": 120, - "loss": 0.2289590686559677 + "loss": 0.32919183373451233 }, { "stage": "goal", "step": 140, - "loss": 0.36456483602523804 + "loss": 0.171668142080307 }, { "stage": "goal", "step": 160, - "loss": 0.0842418223619461 + "loss": 1.0167959928512573 }, { "stage": "goal", "step": 180, - "loss": 0.9889341592788696 + "loss": 0.8417221903800964 }, { "stage": "goal", "step": 200, - "loss": 0.5776903629302979 + "loss": 0.10604722797870636 }, { "stage": "goal", "step": 220, - "loss": 0.44620272517204285 + "loss": 0.07537326216697693 }, { "stage": "goal", "step": 240, - "loss": 0.05634145438671112 + "loss": 0.1885790228843689 }, { "stage": "goal", "step": 260, - "loss": 0.05701519176363945 + "loss": 0.06766514480113983 }, { "stage": "goal", "step": 280, - "loss": 0.06764699518680573 + "loss": 0.14665386080741882 }, { "stage": "goal", "step": 300, - "loss": 0.10973378270864487 + "loss": 0.10101648420095444 + }, + { + "stage": "goal", + "step": 320, + "loss": 0.07647702097892761 + }, + { + "stage": "goal", + "step": 340, + "loss": 0.08067559450864792 + }, + { + "stage": "goal", + "step": 360, + "loss": 0.12979857623577118 + }, + { + "stage": "goal", + "step": 380, + "loss": 0.899062991142273 + }, + { + "stage": "goal", + "step": 400, + "loss": 0.12218281626701355 + }, + { + "stage": "goal", + "step": 420, + "loss": 0.07381283491849899 + }, + { + "stage": "goal", + "step": 440, + "loss": 0.13227435946464539 + }, + { + "stage": "goal", + "step": 460, + "loss": 0.08732561767101288 + }, + { + "stage": "goal", + "step": 480, + "loss": 0.09666546434164047 + }, + { + "stage": "goal", + "step": 500, + "loss": 0.06472445279359818 + }, + { + "stage": "goal", + "step": 520, + "loss": 0.08121205866336823 + }, + { + "stage": "goal", + "step": 540, + "loss": 0.06035961955785751 + }, + { + "stage": "goal", + "step": 560, + "loss": 0.07124190032482147 + }, + { + "stage": "goal", + "step": 580, + "loss": 0.14373362064361572 + }, + { + "stage": "goal", + "step": 600, + "loss": 0.07122757285833359 + }, + { + "stage": "goal", + "step": 620, + "loss": 0.8890293836593628 + }, + { + "stage": "goal", + "step": 640, + "loss": 0.062469229102134705 + }, + { + "stage": "goal", + "step": 660, + "loss": 0.05835632607340813 + }, + { + "stage": "goal", + "step": 680, + "loss": 0.8515450358390808 + }, + { + "stage": "goal", + "step": 700, + "loss": 0.06296383589506149 + }, + { + "stage": "goal", + "step": 720, + "loss": 0.06271275877952576 + }, + { + "stage": "goal", + "step": 740, + "loss": 0.09368501603603363 + }, + { + "stage": "goal", + "step": 760, + "loss": 0.07982345670461655 + }, + { + "stage": "goal", + "step": 780, + "loss": 0.06414016336202621 + }, + { + "stage": "goal", + "step": 800, + "loss": 0.08895712345838547 }, { "stage": "high_level", "step": 1, - "loss": 76.15934753417969 + "loss": 84.32572937011719 }, { "stage": "high_level", @@ -137,323 +262,507 @@ { "stage": "policy", "step": 1, - "loss": 51.04338836669922, - "text_loss": 44.21609115600586, - "flow_loss": 3.4136481285095215 + "loss": 51.165855407714844, + "text_loss": 46.97157287597656, + "flow_loss": 2.0971415042877197, + "rtc_delay": 5, + "use_goal": true, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 20, - "loss": 26.20608901977539, - "text_loss": 5.128565311431885, - "flow_loss": 10.538762092590332 + "loss": 11.228275299072266, + "text_loss": 4.843214988708496, + "flow_loss": 3.192530393600464, + "rtc_delay": 1, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 40, - "loss": 13.018914222717285, - "text_loss": 3.846059799194336, - "flow_loss": 4.586427211761475 + "loss": 6.2679572105407715, + "text_loss": 2.9056875705718994, + "flow_loss": 1.681134819984436, + "rtc_delay": 4, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 60, - "loss": 11.344329833984375, - "text_loss": 3.335771083831787, - "flow_loss": 4.004279136657715 + "loss": 6.005075931549072, + "text_loss": 2.709343671798706, + "flow_loss": 1.647866129875183, + "rtc_delay": 2, + "use_goal": true, + "use_history": true, + "include_subtask": false }, { "stage": "policy", "step": 80, - "loss": 6.586862087249756, - "text_loss": 2.5455708503723145, - "flow_loss": 2.0206456184387207 + "loss": 9.903374671936035, + "text_loss": 2.334010362625122, + "flow_loss": 3.784682273864746, + "rtc_delay": 4, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 100, - "loss": 11.200895309448242, - "text_loss": 2.1739449501037598, - "flow_loss": 4.513474941253662 + "loss": 5.646769046783447, + "text_loss": 2.3033981323242188, + "flow_loss": 1.6716854572296143, + "rtc_delay": 12, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 120, - "loss": 17.273035049438477, - "text_loss": 2.1746408939361572, - "flow_loss": 7.549196720123291 + "loss": 6.867326259613037, + "text_loss": 1.850703239440918, + "flow_loss": 2.5083115100860596, + "rtc_delay": 12, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 140, - "loss": 11.353687286376953, - "text_loss": 2.2218563556671143, - "flow_loss": 4.565915584564209 + "loss": 3.6174087524414062, + "text_loss": 2.0811572074890137, + "flow_loss": 0.7681257128715515, + "rtc_delay": 11, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 160, - "loss": 14.110220909118652, - "text_loss": 2.3123977184295654, - "flow_loss": 5.898911476135254 + "loss": 4.300071716308594, + "text_loss": 1.9097824096679688, + "flow_loss": 1.1951446533203125, + "rtc_delay": 8, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 180, - "loss": 19.999370574951172, - "text_loss": 2.749119997024536, - "flow_loss": 8.62512493133545 + "loss": 8.075389862060547, + "text_loss": 1.6174044609069824, + "flow_loss": 3.2289929389953613, + "rtc_delay": 2, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 200, - "loss": 8.102465629577637, - "text_loss": 1.631801962852478, - "flow_loss": 3.2353317737579346 + "loss": 5.758746147155762, + "text_loss": 1.5621910095214844, + "flow_loss": 2.0982775688171387, + "rtc_delay": 11, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 220, - "loss": 21.903854370117188, - "text_loss": 1.6063504219055176, - "flow_loss": 10.148752212524414 + "loss": 2.9783213138580322, + "text_loss": 1.389180302619934, + "flow_loss": 0.7945705056190491, + "rtc_delay": 11, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 240, - "loss": 10.197904586791992, - "text_loss": 1.5640149116516113, - "flow_loss": 4.3169450759887695 + "loss": 5.832153797149658, + "text_loss": 1.4576387405395508, + "flow_loss": 2.1872575283050537, + "rtc_delay": 0, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 260, - "loss": 6.899936676025391, - "text_loss": 1.3081820011138916, - "flow_loss": 2.795877456665039 + "loss": 3.4007065296173096, + "text_loss": 0.8826411366462708, + "flow_loss": 1.2590327262878418, + "rtc_delay": 5, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 280, - "loss": 10.051774978637695, - "text_loss": 1.2130903005599976, - "flow_loss": 4.419342517852783 + "loss": 4.018008232116699, + "text_loss": 0.9588160514831543, + "flow_loss": 1.529596209526062, + "rtc_delay": 11, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 300, - "loss": 7.963553428649902, - "text_loss": 1.1586928367614746, - "flow_loss": 3.402430295944214 + "loss": 1.492077112197876, + "text_loss": 0.9946327209472656, + "flow_loss": 0.24872222542762756, + "rtc_delay": 9, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 320, - "loss": 5.177718639373779, - "text_loss": 1.2673476934432983, - "flow_loss": 1.9551855325698853 + "loss": 5.338322639465332, + "text_loss": 0.8486652970314026, + "flow_loss": 2.244828701019287, + "rtc_delay": 8, + "use_goal": true, + "use_history": true, + "include_subtask": false }, { "stage": "policy", "step": 340, - "loss": 4.487771034240723, - "text_loss": 1.0361082553863525, - "flow_loss": 1.7258312702178955 + "loss": 2.0994033813476562, + "text_loss": 0.8355546593666077, + "flow_loss": 0.6319243907928467, + "rtc_delay": 7, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 360, - "loss": 4.356203556060791, - "text_loss": 1.3223515748977661, - "flow_loss": 1.5169259309768677 + "loss": 2.2518701553344727, + "text_loss": 0.5455462336540222, + "flow_loss": 0.8531619906425476, + "rtc_delay": 10, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 380, - "loss": 5.289548397064209, - "text_loss": 1.4994667768478394, - "flow_loss": 1.89504075050354 + "loss": 7.911384105682373, + "text_loss": 0.6893162727355957, + "flow_loss": 3.6110339164733887, + "rtc_delay": 7, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 400, - "loss": 8.761286735534668, - "text_loss": 1.4391229152679443, - "flow_loss": 3.6610820293426514 + "loss": 2.3619725704193115, + "text_loss": 0.5672968626022339, + "flow_loss": 0.8973378539085388, + "rtc_delay": 10, + "use_goal": true, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 420, - "loss": 2.400188446044922, - "text_loss": 1.0106937885284424, - "flow_loss": 0.694747269153595 + "loss": 1.233857274055481, + "text_loss": 0.661384642124176, + "flow_loss": 0.28623631596565247, + "rtc_delay": 12, + "use_goal": true, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 440, - "loss": 9.683706283569336, - "text_loss": 1.044368863105774, - "flow_loss": 4.319668769836426 + "loss": 3.295243740081787, + "text_loss": 0.4776339828968048, + "flow_loss": 1.4088048934936523, + "rtc_delay": 9, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 460, - "loss": 10.443679809570312, - "text_loss": 0.7045990228652954, - "flow_loss": 4.869540214538574 + "loss": 1.7052180767059326, + "text_loss": 0.5183412432670593, + "flow_loss": 0.5934383869171143, + "rtc_delay": 8, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 480, - "loss": 6.514384746551514, - "text_loss": 1.185524582862854, - "flow_loss": 2.6644301414489746 + "loss": 8.46796703338623, + "text_loss": 0.6230667233467102, + "flow_loss": 3.922450065612793, + "rtc_delay": 0, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 500, - "loss": 2.255455255508423, - "text_loss": 1.0136079788208008, - "flow_loss": 0.620923638343811 + "loss": 2.5979185104370117, + "text_loss": 0.27633967995643616, + "flow_loss": 1.1607893705368042, + "rtc_delay": 1, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 520, - "loss": 4.1700439453125, - "text_loss": 1.1829472780227661, - "flow_loss": 1.4935483932495117 + "loss": 2.012068510055542, + "text_loss": 0.30755916237831116, + "flow_loss": 0.8522546291351318, + "rtc_delay": 1, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 540, - "loss": 6.956455230712891, - "text_loss": 1.1096038818359375, - "flow_loss": 2.9234256744384766 + "loss": 5.146071910858154, + "text_loss": 0.4677533805370331, + "flow_loss": 2.3391592502593994, + "rtc_delay": 0, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 560, - "loss": 4.5965800285339355, - "text_loss": 1.5257400274276733, - "flow_loss": 1.5354200601577759 + "loss": 1.9480068683624268, + "text_loss": 0.9942646026611328, + "flow_loss": 0.476871132850647, + "rtc_delay": 7, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 580, - "loss": 1.7948781251907349, - "text_loss": 0.835124135017395, - "flow_loss": 0.4798769950866699 + "loss": 7.889751434326172, + "text_loss": 0.6651982665061951, + "flow_loss": 3.612276554107666, + "rtc_delay": 3, + "use_goal": true, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 600, - "loss": 11.411609649658203, - "text_loss": 0.7508105635643005, - "flow_loss": 5.330399513244629 + "loss": 1.454559087753296, + "text_loss": 0.5516519546508789, + "flow_loss": 0.4514535963535309, + "rtc_delay": 11, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 620, - "loss": 14.173189163208008, - "text_loss": 0.553065836429596, - "flow_loss": 6.810061454772949 + "loss": 4.760444641113281, + "text_loss": 0.3332846760749817, + "flow_loss": 2.2135798931121826, + "rtc_delay": 7, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 640, - "loss": 2.423476457595825, - "text_loss": 0.6566628813743591, - "flow_loss": 0.8834067583084106 + "loss": 6.838594913482666, + "text_loss": 0.33531704545021057, + "flow_loss": 3.251638889312744, + "rtc_delay": 3, + "use_goal": true, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 660, - "loss": 7.33753776550293, - "text_loss": 0.5569484829902649, - "flow_loss": 3.3902945518493652 + "loss": 4.01173210144043, + "text_loss": 0.842496395111084, + "flow_loss": 1.5846179723739624, + "rtc_delay": 1, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 680, - "loss": 7.2152018547058105, - "text_loss": 0.6606324911117554, - "flow_loss": 3.277284622192383 + "loss": 5.610707759857178, + "text_loss": 0.45004716515541077, + "flow_loss": 2.5803303718566895, + "rtc_delay": 6, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 700, - "loss": 2.5858616828918457, - "text_loss": 0.471678227186203, - "flow_loss": 1.0570917129516602 + "loss": 2.415992021560669, + "text_loss": 0.2757202684879303, + "flow_loss": 1.0701358318328857, + "rtc_delay": 3, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 720, - "loss": 5.87769889831543, - "text_loss": 0.5711548924446106, - "flow_loss": 2.6532719135284424 + "loss": 3.6413190364837646, + "text_loss": 0.15377745032310486, + "flow_loss": 1.7437708377838135, + "rtc_delay": 2, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 740, - "loss": 6.787323951721191, - "text_loss": 0.5064371228218079, - "flow_loss": 3.1404433250427246 + "loss": 2.4192750453948975, + "text_loss": 0.7513293027877808, + "flow_loss": 0.8339728713035583, + "rtc_delay": 4, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 760, - "loss": 11.166234970092773, - "text_loss": 0.5802748203277588, - "flow_loss": 5.292980194091797 + "loss": 1.3346192836761475, + "text_loss": 0.2767994999885559, + "flow_loss": 0.5289098620414734, + "rtc_delay": 10, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 780, - "loss": 3.563460350036621, - "text_loss": 0.6946841478347778, - "flow_loss": 1.4343880414962769 + "loss": 2.4135358333587646, + "text_loss": 0.16017527878284454, + "flow_loss": 1.1266802549362183, + "rtc_delay": 3, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 800, - "loss": 3.2471487522125244, - "text_loss": 0.44556593894958496, - "flow_loss": 1.4007914066314697 + "loss": 2.966399908065796, + "text_loss": 0.2566848397254944, + "flow_loss": 1.3548575639724731, + "rtc_delay": 3, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 820, - "loss": 3.1103086471557617, - "text_loss": 0.7132699489593506, - "flow_loss": 1.1985193490982056 + "loss": 1.8670570850372314, + "text_loss": 0.23726917803287506, + "flow_loss": 0.8148939609527588, + "rtc_delay": 4, + "use_goal": false, + "use_history": true, + "include_subtask": true }, { "stage": "policy", "step": 840, - "loss": 4.753307342529297, - "text_loss": 0.48397523164749146, - "flow_loss": 2.1346659660339355 + "loss": 1.6185007095336914, + "text_loss": 0.27975988388061523, + "flow_loss": 0.6693704128265381, + "rtc_delay": 5, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 860, - "loss": 3.5492496490478516, - "text_loss": 0.366057813167572, - "flow_loss": 1.5915958881378174 + "loss": 0.9501249194145203, + "text_loss": 0.4190976619720459, + "flow_loss": 0.2655136287212372, + "rtc_delay": 7, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 880, - "loss": 4.0379414558410645, - "text_loss": 0.6511648893356323, - "flow_loss": 1.6933883428573608 + "loss": 1.5665624141693115, + "text_loss": 0.317945659160614, + "flow_loss": 0.6243083477020264, + "rtc_delay": 5, + "use_goal": false, + "use_history": false, + "include_subtask": true }, { "stage": "policy", "step": 900, - "loss": 1.7170453071594238, - "text_loss": 0.2673894464969635, - "flow_loss": 0.7248279452323914 + "loss": 1.8227444887161255, + "text_loss": 0.32235953211784363, + "flow_loss": 0.7501924633979797, + "rtc_delay": 8, + "use_goal": false, + "use_history": true, + "include_subtask": true } ] \ No newline at end of file diff --git a/src/pi07mini/context.py b/src/pi07mini/context.py index 2e05286..3e19a84 100644 --- a/src/pi07mini/context.py +++ b/src/pi07mini/context.py @@ -21,7 +21,6 @@ class PolicyMetadata: 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: @@ -47,7 +46,6 @@ class PolicyMetadata: "quality": self.quality_score, "mistake": self.mistake != "none", "control_mode": self.control_mode, - "source_type": self.source_type, } @@ -106,36 +104,34 @@ def format_policy_prompt(components: Mapping[str, str]) -> str: 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), +def future_observation_proxy( + image: torch.Tensor, + history_xy: torch.Tensor, + future_xy: torch.Tensor, ) -> torch.Tensor: - """Render a normalized Gaussian subgoal heatmap from one endpoint only. + """Create a tiny future-camera proxy instead of encoding coordinates as a heatmap. - 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. + The source dataset contains only the current frame. For the architecture + validation we therefore warp that frame according to the segment endpoint, + producing an image-space target with the same contract as ``o_tend`` in the + paper. The policy only receives the generated image, never the coordinates. """ - 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) + image = torch.as_tensor(image, dtype=torch.float32) + history = torch.as_tensor(history_xy, dtype=torch.float32) + future = torch.as_tensor(future_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 future.ndim != 2 or history.shape[-1] != 2 or future.shape[-1] != 2: + raise ValueError("history_xy and future_xy must have shape [T,2]") + displacement = future[-1] - history[-1] + scale = max(image.shape[-2:]) / 32.0 + dx = int(torch.round(displacement[1] * scale).clamp(-image.shape[-1] // 5, image.shape[-1] // 5)) + dy = int(torch.round(-displacement[0] * scale * 0.08).clamp(-image.shape[-2] // 6, image.shape[-2] // 6)) + distance = torch.linalg.vector_norm(displacement) + exposure = float((1.0 + 0.02 * torch.tanh(distance / 8.0)).item()) + shifted = torch.roll(image, shifts=(dy, dx), dims=(-2, -1)) * exposure + return F.avg_pool2d(shifted[None], kernel_size=3, stride=1, padding=1)[0].clamp(0.0, 1.0) def make_history_proxy_frames( diff --git a/src/pi07mini/data.py b/src/pi07mini/data.py index 877ac25..0bf18e9 100644 --- a/src/pi07mini/data.py +++ b/src/pi07mini/data.py @@ -16,12 +16,12 @@ from .context import ( PolicyMetadata, build_ambiguous_history_pair, drop_prompt_components, - endpoint_subgoal_heatmap, format_policy_prompt, + future_observation_proxy, make_history_proxy_frames, prompt_components, ) -from .rtc import MAX_RTC_DELAY_STEPS, align_actions_for_rtc, sample_rtc_delay +from .rtc import MAX_RTC_DELAY_STEPS, build_rtc_conditioning, sample_rtc_delay DEFAULT_SOURCE_MANIFEST = ( @@ -232,7 +232,7 @@ def build_policy_trajectories( history, expert, PolicyMetadata( - "normal", "expert", "none", "human", 1_000, 5, "ee", "demonstration" + "normal", "expert", "none", "human", 1_000, 5, "ee" ), ), PolicyTrajectory( @@ -240,7 +240,7 @@ def build_policy_trajectories( history, fast, PolicyMetadata( - "fast", "good", "none", "autonomous", 500, 4, "ee", "autonomous" + "fast", "good", "none", "autonomous", 500, 4, "ee" ), ), PolicyTrajectory( @@ -248,7 +248,7 @@ def build_policy_trajectories( history, slow, PolicyMetadata( - "slow", "good", "none", "autonomous", 2_000, 4, "ee", "autonomous" + "slow", "good", "none", "autonomous", 2_000, 4, "ee" ), ), PolicyTrajectory( @@ -256,7 +256,7 @@ def build_policy_trajectories( history, understeer, PolicyMetadata( - "normal", "poor", "understeer", "autonomous", 1_000, 1, "ee", "autonomous" + "normal", "poor", "understeer", "autonomous", 1_000, 1, "ee" ), ), PolicyTrajectory( @@ -264,7 +264,7 @@ def build_policy_trajectories( history, stopped, PolicyMetadata( - "slow", "poor", "stop", "autonomous", 2_000, 1, "ee", "autonomous" + "slow", "poor", "stop", "autonomous", 2_000, 1, "ee" ), ), ] @@ -367,11 +367,11 @@ class Pi07SteerableDataset(Dataset[dict[str, Any]]): max_delay_steps=self.max_rtc_delay_steps, action_horizon=self.action_horizon, ) - rtc = align_actions_for_rtc( + rtc = build_rtc_conditioning( 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) + future_observation = future_observation_proxy( + image, policy.history_xy, policy.future_xy ) history_pair = build_ambiguous_history_pair( policy.history_xy, frames=self.history_frames @@ -402,17 +402,17 @@ class Pi07SteerableDataset(Dataset[dict[str, Any]]): "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(), + "future_observation": future_observation, + "subgoal_pixel_values": future_observation.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, + "rtc_action_prefix": rtc.action_prefix, + "rtc_prefix_mask": rtc.prefix_mask, + "rtc_postfix_mask": rtc.postfix_mask, + "past_subtasks": ("keep_lane",), "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 index eb93db6..4325aa1 100644 --- a/src/pi07mini/flow.py +++ b/src/pi07mini/flow.py @@ -1,18 +1,101 @@ from __future__ import annotations + from collections.abc import Callable + import torch from torch import Tensor class FlowMatching: + """Rectified flow with optional paper-faithful training-time RTC.""" + 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 + self.inference_steps = int(inference_steps) + self.cutoff = float(cutoff) + + def training_batch( + self, + target: Tensor, + *, + delays: Tensor | None = None, + generator: torch.Generator | None = None, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + batch, horizon = target.shape[:2] + uniform = torch.rand( + batch, device=target.device, dtype=target.dtype, generator=generator + ) + base_time = self.cutoff * (1.0 - uniform.pow(1.0 / 1.5)) + noise = torch.randn( + target.shape, + device=target.device, + dtype=target.dtype, + generator=generator, + ) + + if delays is None: + times = base_time + expanded = times[:, None, None] + loss_mask = torch.ones( + batch, horizon, dtype=torch.bool, device=target.device + ) + else: + delays = torch.as_tensor(delays, dtype=torch.long, device=target.device) + if delays.shape != (batch,): + raise ValueError("delays must have shape [batch]") + if torch.any(delays < 0) or torch.any(delays >= horizon): + raise ValueError("each RTC delay must lie in [0, horizon)") + prefix_mask = torch.arange(horizon, device=target.device)[None] < delays[:, None] + times = torch.where(prefix_mask, torch.ones_like(base_time[:, None]), base_time[:, None]) + expanded = times[:, :, None] + loss_mask = ~prefix_mask + + noisy = expanded * target + (1.0 - expanded) * noise + target_velocity = target - noise + return noisy, times, target_velocity, loss_mask + + @staticmethod + def masked_loss(prediction: Tensor, target: Tensor, mask: Tensor) -> Tensor: + if mask.shape != prediction.shape[:2]: + raise ValueError("flow loss mask must have shape [batch, horizon]") + per_token = (prediction.float() - target.float()).square().mean(dim=-1) + weights = mask.to(per_token.dtype) + return (per_token * weights).sum() / weights.sum().clamp_min(1.0) + @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])) + def sample( + self, + model_fn: Callable[[Tensor, Tensor], Tensor], + shape: tuple[int, ...], + device: torch.device | str, + steps: int | None = None, + *, + action_prefix: Tensor | None = None, + delays: Tensor | None = None, + generator: torch.Generator | None = None, + ) -> Tensor: + count = steps or self.inference_steps + x = torch.randn(shape, device=device, generator=generator) + prefix_mask = None + if action_prefix is not None or delays is not None: + if action_prefix is None or delays is None: + raise ValueError("action_prefix and delays must be supplied together") + action_prefix = torch.as_tensor(action_prefix, dtype=x.dtype, device=x.device) + delays = torch.as_tensor(delays, dtype=torch.long, device=x.device) + if action_prefix.shape != x.shape or delays.shape != (shape[0],): + raise ValueError("RTC action prefix or delay has the wrong shape") + prefix_mask = torch.arange(shape[1], device=x.device)[None] < delays[:, None] + + grid = torch.linspace(0.0, 1.0, count + 1, device=device) + for index in range(count): + timestep = grid[index].expand(shape[0]) + if prefix_mask is not None: + x = torch.where(prefix_mask[:, :, None], action_prefix, x) + timestep = torch.where( + prefix_mask, + torch.ones_like(timestep[:, None]), + timestep[:, None], + ) + x = x + (grid[index + 1] - grid[index]) * model_fn(x, timestep) + if prefix_mask is not None: + x = torch.where(prefix_mask[:, :, None], action_prefix, x) return x diff --git a/src/pi07mini/infer.py b/src/pi07mini/infer.py index eaf9079..0ac9565 100644 --- a/src/pi07mini/infer.py +++ b/src/pi07mini/infer.py @@ -1,40 +1,258 @@ from __future__ import annotations + import argparse from pathlib import Path + import matplotlib.pyplot as plt -import numpy as np,torch +import numpy as np +import 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 +from .model import HighLevelPolicy, Pi07Config, Pi07Mini, TinyGoalGenerator +from .rtc import build_rtc_conditioning +from .runtime import choose_device, write_json +from .tokenizer import EpisodeMetadata, MiniTokenizer + @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() +def evaluate(checkpoint_path: str) -> tuple[Path, Path]: + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + cfg = checkpoint["config"] + root = Path(__file__).resolve().parents[2] + device = choose_device(cfg["device"]) + tokenizer = MiniTokenizer() + dataset = 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=checkpoint["action_statistics"], + max_rtc_delay_steps=cfg["rtc"]["max_delay_steps"], + ) + policy_config = Pi07Config(**checkpoint["policy_config"]) + policy = Pi07Mini(policy_config) + policy.load_state_dict(checkpoint["policy"]) + policy.to(device).eval() + high_level = HighLevelPolicy(policy_config) + high_level.load_state_dict(checkpoint["high_level"]) + high_level.to(device).eval() + goal_generator = TinyGoalGenerator(tokenizer.vocab_size) + goal_generator.load_state_dict(checkpoint["goal"]) + goal_generator.to(device).eval() + flow = FlowMatching(cfg["flow"]["integration_steps"]) + + sample = next( + dataset[index] + for index in range(len(dataset)) + if dataset[index]["policy_kind"] == "expert" + ) + metadata = EpisodeMetadata(**sample["token_metadata"]) + image = sample["image"][None].to(device) + high_level_prompt = tokenizer.encode_high_level_prompt( + "drive safely", sample["past_subtasks"], metadata + ) + high_level_output = high_level( + image, + high_level_prompt[None].to(device), + sample["history_images"][None].to(device), + sample["history_states"][None].to(device), + sample["current_state"][None].to(device), + ) + predicted_subtask = tokenizer.subtasks.decode( + int(high_level_output.logits.argmax(-1).cpu()) + ) + + goal_prompt, _ = tokenizer.encode_prompt( + "drive safely", predicted_subtask, metadata + ) + goal_noise = torch.randn( + sample["future_observation"][None].shape, + generator=torch.Generator().manual_seed(cfg["evaluation"]["fixed_noise_seed"]), + ).to(device) + goal_grid = torch.linspace(0.0, 1.0, cfg["evaluation"]["goal_flow_steps"] + 1, device=device) + generated_goal = goal_noise + for index in range(cfg["evaluation"]["goal_flow_steps"]): + time = goal_grid[index].expand(1) + generated_goal = generated_goal + (goal_grid[index + 1] - goal_grid[index]) * goal_generator( + image, goal_prompt[None].to(device), generated_goal, time + ).vector_field + generated_goal_cpu = generated_goal[0].cpu().clamp(0.0, 1.0) + + def sample_policy( + *, + include_metadata: bool = True, + use_goal: bool = True, + use_history: bool = True, + guidance_beta: float | None = None, + rtc_delay: int = 0, + oracle_goal: bool = False, + ) -> tuple[np.ndarray, torch.Tensor]: + conditioned = tokenizer.build_policy_sequence( + "drive safely", + predicted_subtask, + sample["actions"], + metadata, + include_episode_metadata=include_metadata, + ) + conditioned_ids = conditioned.input_ids[: conditioned.action_start_index][None].to(device) + common = { + "pixel_values": image, + "history_pixel_values": sample["history_images"][None].to(device) + if use_history + else None, + "history_states": sample["history_states"][None].to(device) + if use_history + else None, + "current_state": sample["current_state"][None].to(device), + "subgoal_pixel_values": ( + sample["future_observation"][None].to(device) + if oracle_goal + else generated_goal + ) + if use_goal + else None, + } + + if guidance_beta is None: + def field(actions: torch.Tensor, times: torch.Tensor) -> torch.Tensor: + return policy( + input_ids=conditioned_ids, + noisy_actions=actions, + timesteps=times, + **common, + ).vector_field + else: + unconditional = tokenizer.build_policy_sequence( + "drive safely", + predicted_subtask, + sample["actions"], + metadata, + include_episode_metadata=False, + ) + unconditional_ids = unconditional.input_ids[ + : unconditional.action_start_index + ][None].to(device) + + def field(actions: torch.Tensor, times: torch.Tensor) -> torch.Tensor: + return policy.metadata_cfg( + conditioned_ids, + unconditional_ids, + guidance_scale=guidance_beta, + noisy_actions=actions, + timesteps=times, + **common, + ) + + rtc = build_rtc_conditioning( + sample["actions"], + rtc_delay, + max_delay_steps=cfg["rtc"]["max_delay_steps"], + ) + torch.manual_seed(cfg["evaluation"]["fixed_noise_seed"]) + normalized = flow.sample( + field, + (1, dataset.action_horizon, 2), + device, + cfg["flow"]["integration_steps"], + action_prefix=rtc.action_prefix[None].to(device), + delays=torch.tensor([rtc_delay], device=device), + )[0].cpu() + trajectory = dataset.action_normalizer.decode( + normalized, sample["history_xy"] + ).numpy() + return trajectory, normalized + + curves: dict[str, np.ndarray] = {} + normalized_predictions: dict[str, torch.Tensor] = {} + configurations = { + "full_generated_goal": {}, + "full_oracle_goal": {"oracle_goal": True}, + "no_metadata": {"include_metadata": False}, + "no_goal": {"use_goal": False}, + "no_history": {"use_history": False}, + "metadata_cfg": {"guidance_beta": cfg["steering"]["cfg_beta"]}, + } + for name, arguments in configurations.items(): + curves[name], normalized_predictions[name] = sample_policy(**arguments) + + target = sample["future_xy"].numpy() + target_actions = sample["actions"] + trajectory_metrics = { + name: float(np.mean((trajectory - target) ** 2)) + for name, trajectory in curves.items() + } + action_metrics = { + name: float(torch.mean((actions - target_actions) ** 2)) + for name, actions in normalized_predictions.items() + } + + rtc_metrics: dict[str, dict[str, float]] = {} + for delay in cfg["evaluation"]["compare_rtc_delays"]: + trajectory, actions = sample_policy(rtc_delay=delay) + postfix = torch.arange(dataset.action_horizon) >= delay + rtc_metrics[str(delay)] = { + "postfix_action_mse": float( + torch.mean((actions[postfix] - target_actions[postfix]) ** 2) + ), + "trajectory_mse": float(np.mean((trajectory - target) ** 2)), + } + + output_dir = Path(checkpoint_path).parent + figure, axis = plt.subplots(figsize=(10, 7)) + axis.plot(target[:, 0], target[:, 1], "o-", label="expert target", linewidth=2) + for name, trajectory in curves.items(): + axis.plot( + trajectory[:, 0], + trajectory[:, 1], + "x--", + label=f"{name} ({trajectory_metrics[name]:.3f})", + ) + axis.axis("equal") + axis.grid(True, alpha=0.3) + axis.legend(fontsize=8) + axis.set_xlabel("forward x (m)") + axis.set_ylabel("left y (m)") + axis.set_title("pi0.7 predicted-subtask context and RTC evaluation") + figure.tight_layout() + trajectory_path = output_dir / "context_ablation.png" + figure.savefig(trajectory_path, dpi=170) + plt.close(figure) + + figure, 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(sample["future_observation"].permute(1, 2, 0)) + axes[1].set_title("future observation target") + axes[2].imshow(generated_goal_cpu.permute(1, 2, 0)) + axes[2].set_title("generated visual subgoal") + for axis in axes: + axis.axis("off") + figure.tight_layout() + goal_path = output_dir / "goal_generation.png" + figure.savefig(goal_path, dpi=170) + plt.close(figure) + + results = { + "target_subtask": tokenizer._classify_subtask(sample["reasoning"]), + "predicted_subtask": predicted_subtask, + "goal_image_mse": float( + torch.mean((generated_goal_cpu - sample["future_observation"]) ** 2) + ), + "trajectory_metrics": trajectory_metrics, + "normalized_action_metrics": action_metrics, + "rtc_metrics": rtc_metrics, + } + write_json(output_dir / "evaluation.json", results) + print(results) + return trajectory_path, goal_path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", default="outputs/overfit/checkpoint.pt") + evaluate(parser.parse_args().checkpoint) + + +if __name__ == "__main__": + main() diff --git a/src/pi07mini/model.py b/src/pi07mini/model.py index 4c66a2a..347c7d7 100644 --- a/src/pi07mini/model.py +++ b/src/pi07mini/model.py @@ -72,20 +72,29 @@ class Pi07Config: @dataclass(frozen=True) class PrefixLayout: observation_tokens: int + prompt_tokens: int goal_tokens: int - text_tokens: int + fast_tokens: int @property - def goal_start(self) -> int: + def prompt_start(self) -> int: return self.observation_tokens @property - def text_start(self) -> int: - return self.observation_tokens + self.goal_tokens + def goal_start(self) -> int: + return self.prompt_start + self.prompt_tokens + + @property + def fast_start(self) -> int: + return self.goal_start + self.goal_tokens + + @property + def text_tokens(self) -> int: + return self.prompt_tokens + self.fast_tokens @property def total_tokens(self) -> int: - return self.text_start + self.text_tokens + return self.fast_start + self.fast_tokens @dataclass @@ -133,7 +142,15 @@ class AdaRMSNorm(nn.Module): 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] + if condition.ndim == 2: + shift = shift[:, None] + scale = scale[:, None] + elif condition.ndim == 3: + if shift.shape[:2] != x.shape[:2]: + raise ValueError("per-token timestep conditioning must match the action sequence") + else: + raise ValueError("condition must have shape [batch, dim] or [batch, tokens, dim]") + return self.norm(x) * (1 + scale) + shift class TimeEmbedding(nn.Module): @@ -146,6 +163,7 @@ class TimeEmbedding(nn.Module): ) def forward(self, timestep: Tensor) -> Tensor: + original_shape = timestep.shape timestep = timestep.reshape(-1).float() half = self.fourier_dim // 2 frequencies = torch.exp( @@ -158,7 +176,8 @@ class TimeEmbedding(nn.Module): ) ) angles = 2 * math.pi * timestep[:, None] * frequencies[None] - return self.mlp(torch.cat((angles.sin(), angles.cos()), dim=-1)) + embedded = self.mlp(torch.cat((angles.sin(), angles.cos()), dim=-1)) + return embedded.reshape(*original_shape, embedded.shape[-1]) class AttentionWeights(nn.Module): @@ -403,32 +422,47 @@ class Pi07Mini(nn.Module): raise ValueError("attention_mask must match input_ids") text_valid = attention_mask.bool() - layout = PrefixLayout(observation.shape[1], goal.shape[1], text.shape[1]) + action_start = input_ids.eq(self.config.action_start_token_id) + if action_start.any(): + first_action = action_start.long().argmax(dim=1) + if not torch.equal(first_action, first_action[:1].expand_as(first_action)): + raise ValueError("all batch items must place ACTION_START at the same position") + prompt_tokens = int(first_action[0]) + else: + prompt_tokens = input_ids.shape[1] + fast_tokens = input_ids.shape[1] - prompt_tokens + prompt = text[:, :prompt_tokens] + prompt_valid = text_valid[:, :prompt_tokens] + fast = text[:, prompt_tokens:] + fast_valid = text_valid[:, prompt_tokens:] + + layout = PrefixLayout( + observation_tokens=observation.shape[1], + prompt_tokens=prompt_tokens, + goal_tokens=goal.shape[1], + fast_tokens=fast_tokens, + ) if layout.total_tokens > self.config.max_prefix_tokens: raise ValueError("prefix exceeds configured maximum length") - prefix = torch.cat((observation, goal, text), dim=1) + # Official pi0.7 Figure 19 order: + # memory observations -> causal text prompt -> image-goal block -> FAST. + prefix = torch.cat((observation, prompt, goal, fast), 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 - ), + torch.ones(layout.prompt_tokens, dtype=torch.long, device=prefix.device), + torch.full((layout.goal_tokens,), 2, dtype=torch.long, device=prefix.device), + torch.ones(layout.fast_tokens, 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) + valid = torch.cat((observation_valid, prompt_valid, goal_valid, fast_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. - """ + """Build the paper's observation -> text -> image-goal -> FAST mask.""" length = layout.total_tokens if valid.shape[1] != length: @@ -436,13 +470,21 @@ class Pi07Mini(nn.Module): 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 + obs = q < layout.prompt_start + prompt = (q >= layout.prompt_start) & (q < layout.goal_start) + goal = (q >= layout.goal_start) & (q < layout.fast_start) + fast = q >= layout.fast_start structural = ( - (obs & (k < layout.goal_start)) - | (goal & (k < layout.text_start)) - | (text & ((k < layout.text_start) | ((k >= layout.text_start) & (k <= q)))) + (obs & (k < layout.prompt_start)) + | ( + prompt + & ((k < layout.prompt_start) | ((k >= layout.prompt_start) & (k <= q))) + ) + | (goal & (k < layout.fast_start)) + | ( + fast + & ((k < layout.fast_start) | ((k >= layout.fast_start) & (k <= q))) + ) ) edges = structural[None] & valid[:, None] identity = torch.eye(length, dtype=torch.bool, device=device)[None] @@ -462,18 +504,12 @@ class Pi07Mini(nn.Module): 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, - ) + supplied = torch.cat(( + torch.ones(batch, layout.observation_tokens, dtype=torch.bool, device=valid.device), + supplied[:, : layout.prompt_tokens], + torch.ones(batch, layout.goal_tokens, dtype=torch.bool, device=valid.device), + supplied[:, layout.prompt_tokens :], + ), dim=1) if supplied.shape != valid.shape: raise ValueError("action_prefix_mask must cover text or the full prefix") readable = valid & supplied @@ -573,6 +609,11 @@ class Pi07Mini(nn.Module): horizon = noisy_actions.shape[1] action = self.action_projection(noisy_actions) action = action + self.action_position[:, :horizon] + if timesteps.ndim == 1: + if timesteps.shape != (noisy_actions.shape[0],): + raise ValueError("timesteps must have shape [batch] or [batch, horizon]") + elif timesteps.shape != noisy_actions.shape[:2]: + raise ValueError("per-action timesteps must have shape [batch, 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) @@ -587,7 +628,14 @@ class Pi07Mini(nn.Module): prefix_hidden = self.prefix_norm(prefix) action_hidden = self.action_norm(action, condition) - text_logits = self.text_head(prefix_hidden[:, layout.text_start :]) + text_hidden = torch.cat( + ( + prefix_hidden[:, layout.prompt_start : layout.goal_start], + prefix_hidden[:, layout.fast_start :], + ), + dim=1, + ) + text_logits = self.text_head(text_hidden) vector_field = self.vector_field_head(action_hidden) text_loss = None diff --git a/src/pi07mini/rtc.py b/src/pi07mini/rtc.py index f0297f0..f6e144e 100644 --- a/src/pi07mini/rtc.py +++ b/src/pi07mini/rtc.py @@ -9,27 +9,22 @@ MAX_RTC_DELAY_STEPS = 12 @dataclass(frozen=True) -class RTCAlignment: - """A fixed-width target aligned to an inference/control delay.""" +class RTCConditioning: + """Hard action-prefix conditioning used by training-time RTC.""" - delayed_target_actions: torch.Tensor - valid_mask: torch.Tensor - executed_prefix_actions: torch.Tensor + action_prefix: torch.Tensor + prefix_mask: torch.Tensor + postfix_mask: torch.Tensor delay_steps: int -def align_actions_for_rtc( +def build_rtc_conditioning( 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. - """ +) -> RTCConditioning: + """Pad the committed action prefix to the full prediction horizon.""" actions = torch.as_tensor(actions, dtype=torch.float32) if actions.ndim < 2 or actions.shape[-2] < 1: @@ -39,13 +34,17 @@ def align_actions_for_rtc( 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)) + + prefix_mask = torch.arange(horizon, device=actions.device) < delay_steps + prefix_mask = prefix_mask.expand(actions.shape[:-2] + (horizon,)) + action_prefix = torch.zeros_like(actions) + action_prefix[..., :delay_steps, :] = actions[..., :delay_steps, :] + return RTCConditioning( + action_prefix=action_prefix, + prefix_mask=prefix_mask, + postfix_mask=~prefix_mask, + delay_steps=int(delay_steps), + ) def sample_rtc_delay( @@ -57,6 +56,4 @@ def sample_rtc_delay( 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 index d88b554..56cecef 100644 --- a/src/pi07mini/runtime.py +++ b/src/pi07mini/runtime.py @@ -1,14 +1,39 @@ from __future__ import annotations -import json,random + +import json +import 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") + +import numpy as np +import torch +import yaml + + +def load_config(path: str | Path) -> dict: + with Path(path).open(encoding="utf-8") as stream: + return yaml.safe_load(stream) + + +def seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def choose_device(name: str = "auto") -> torch.device: + if name != "auto": + return torch.device(name) + if torch.cuda.is_available(): + return torch.device("cuda") + if torch.backends.mps.is_available(): + return torch.device("mps") return torch.device("cpu") -def write_json(path,value): - p=Path(path);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(value,indent=2,ensure_ascii=False)) + + +def write_json(path: str | Path, value: object) -> None: + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps(value, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) diff --git a/src/pi07mini/tokenizer.py b/src/pi07mini/tokenizer.py index 47392f0..86badde 100644 --- a/src/pi07mini/tokenizer.py +++ b/src/pi07mini/tokenizer.py @@ -2,7 +2,7 @@ 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 +and control mode. This tokenizer represents those concepts with explicit marker/value pairs so the architecture can be validated without shipping a full language tokenizer. """ @@ -26,7 +26,6 @@ class EpisodeMetadata: 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: @@ -84,7 +83,7 @@ class MiniTokenizer: QUALITY_MARKER = 260 MISTAKE_MARKER = 261 CONTROL_MARKER = 262 - SOURCE_MARKER = 263 + NEXT_SUBTASK_MARKER = 263 TASK_LABELS = ( "drive", @@ -123,18 +122,7 @@ class MiniTokenizer: 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_START = CONTROL_EE + 1 ACTION_END = ACTION_START + 1 ACTION_TOKEN_START = ACTION_END + 1 ACTION_BINS = 64 @@ -145,7 +133,6 @@ class MiniTokenizer: 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: @@ -197,23 +184,6 @@ class MiniTokenizer: 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, @@ -257,14 +227,40 @@ class MiniTokenizer: 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) + def encode_high_level_prompt( + self, + task: str, + past_subtasks: Iterable[str], + metadata: EpisodeMetadata, + ) -> Tensor: + """Encode task, past semantic decisions, metadata, and a next-subtask query.""" + + prompt, _ = self.encode_prompt( + task, + subtask=None, + metadata=metadata, + include_subtask=False, + ) + control_pair = prompt[-2:] + body = prompt[:-2] + history: list[int] = [] + for subtask in past_subtasks: + history.extend( + (self.SUBTASK_MARKER, self.subtasks.encode(self._classify_subtask(subtask))) + ) + return torch.cat( + ( + body, + torch.tensor(history, dtype=torch.long), + control_pair, + torch.tensor([self.NEXT_SUBTASK_MARKER], dtype=torch.long), + ) + ) + @staticmethod def _orthonormal_dct_matrix( horizon: int, *, device: torch.device, dtype: torch.dtype @@ -390,4 +386,3 @@ class MiniTokenizer: include_episode_metadata=False, ) return conditioned, unconditional - diff --git a/src/pi07mini/train.py b/src/pi07mini/train.py index df6c8c3..6095f46 100644 --- a/src/pi07mini/train.py +++ b/src/pi07mini/train.py @@ -1,40 +1,275 @@ from __future__ import annotations -import argparse,copy,json + +import argparse +import json from pathlib import Path + import torch -from .data import Pi07SteerableDataset,materialize_subset + +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 +from .model import HighLevelPolicy, Pi07Config, Pi07Mini, TinyGoalGenerator +from .rtc import sample_rtc_delay +from .runtime import choose_device, load_config, seed_everything, write_json +from .tokenizer import EpisodeMetadata, MiniTokenizer -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 build_config(cfg: dict, tokenizer: MiniTokenizer) -> Pi07Config: + model = cfg["model"] + return Pi07Config( + vocab_size=tokenizer.vocab_size, + action_dim=2, + state_dim=2, + image_size=cfg["data"]["image_size"], + patch_size=model["patch_size"], + d_model=model["d_model"], + n_heads=model["n_heads"], + n_layers=model["n_layers"], + memory_layers=model.get("memory_layers", 2), + mlp_ratio=model["mlp_ratio"], + max_history_frames=cfg["data"]["history_frames"], + max_text_tokens=model["max_text_tokens"], + max_action_tokens=cfg["data"]["action_horizon"], + action_start_token_id=tokenizer.ACTION_START, + dropout=model["dropout"], + ) -def main(): - p=argparse.ArgumentParser();p.add_argument('--config',default='configs/overfit.yaml');train(p.parse_args().config) -if __name__=='__main__':main() + +def _metadata_with_paper_dropout( + metadata: EpisodeMetadata, + cfg: dict, + *, + generator: torch.Generator, +) -> EpisodeMetadata: + steering = cfg["steering"] + if torch.rand((), generator=generator) < steering["metadata_full_dropout"]: + return EpisodeMetadata(control_mode=metadata.control_mode) + + probability = steering["metadata_component_dropout"] + return EpisodeMetadata( + speed_steps=None + if torch.rand((), generator=generator) < probability + else metadata.speed_steps, + quality=None + if torch.rand((), generator=generator) < probability + else metadata.quality, + mistake=None + if torch.rand((), generator=generator) < probability + else metadata.mistake, + control_mode=metadata.control_mode, + ) + + +def train(config_path: str) -> Path: + cfg = load_config(config_path) + root = Path(__file__).resolve().parents[2] + seed_everything(cfg["seed"]) + device = choose_device(cfg["device"]) + generator = torch.Generator().manual_seed(cfg["seed"]) + + manifest = materialize_subset( + cfg["data"]["source_manifest"], + root / "data/overfit", + max_samples=cfg["data"]["max_samples"], + ) + dataset = 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"], + max_rtc_delay_steps=cfg["rtc"]["max_delay_steps"], + ) + tokenizer = MiniTokenizer() + policy_config = build_config(cfg, tokenizer) + policy = Pi07Mini(policy_config).to(device) + high_level = HighLevelPolicy(policy_config).to(device) + goal_generator = TinyGoalGenerator(tokenizer.vocab_size).to(device) + flow = FlowMatching(cfg["flow"]["integration_steps"]) + output_dir = root / cfg["output_dir"] + output_dir.mkdir(parents=True, exist_ok=True) + history: list[dict] = [] + + # The paper trains the world model on segments with high-quality subtask labels. + goal_samples = [ + dataset[index] + for index in range(len(dataset)) + if dataset[index]["quality_score"] >= 4 + ] + optimizer = torch.optim.AdamW( + goal_generator.parameters(), lr=cfg["train"]["learning_rate"] + ) + goal_generator.train() + for step in range(1, cfg["train"]["goal_steps"] + 1): + sample = goal_samples[(step - 1) % len(goal_samples)] + metadata = EpisodeMetadata(**sample["token_metadata"]) + prompt, _ = tokenizer.encode_prompt( + "drive safely", sample["reasoning"], metadata + ) + target = sample["future_observation"][None].to(device) + time = torch.rand(1, device=device) + noise = torch.randn_like(target) + noisy = time[:, None, None, None] * target + (1 - time[:, None, None, None]) * noise + output = goal_generator( + sample["image"][None].to(device), + prompt[None].to(device), + noisy, + time, + target_velocity=target - noise, + ) + optimizer.zero_grad(set_to_none=True) + output.loss.backward() + optimizer.step() + if step == 1 or step % cfg["train"]["log_every"] == 0: + row = {"stage": "goal", "step": step, "loss": float(output.loss.detach().cpu())} + history.append(row) + print(json.dumps(row)) + + optimizer = torch.optim.AdamW( + high_level.parameters(), lr=cfg["train"]["learning_rate"] + ) + high_level.train() + for step in range(1, cfg["train"]["high_level_steps"] + 1): + sample = dataset[(step - 1) % len(dataset)] + metadata = EpisodeMetadata(**sample["token_metadata"]) + prompt = tokenizer.encode_high_level_prompt( + "drive safely", sample["past_subtasks"], metadata + ) + target = torch.tensor( + [tokenizer.subtasks.encode(tokenizer._classify_subtask(sample["reasoning"]))], + device=device, + ) + output = high_level( + sample["image"][None].to(device), + prompt[None].to(device), + sample["history_images"][None].to(device), + sample["history_states"][None].to(device), + sample["current_state"][None].to(device), + target_token=target, + ) + optimizer.zero_grad(set_to_none=True) + output.loss.backward() + optimizer.step() + if step == 1 or step % cfg["train"]["log_every"] == 0: + row = { + "stage": "high_level", + "step": step, + "loss": float(output.loss.detach().cpu()), + } + history.append(row) + print(json.dumps(row)) + + optimizer = torch.optim.AdamW( + policy.parameters(), + lr=cfg["train"]["learning_rate"], + weight_decay=cfg["train"]["weight_decay"], + ) + policy.train() + for step in range(1, cfg["train"]["policy_steps"] + 1): + sample = dataset[(step - 1) % len(dataset)] + metadata = _metadata_with_paper_dropout( + EpisodeMetadata(**sample["token_metadata"]), cfg, generator=generator + ) + use_goal = bool( + torch.rand((), generator=generator) < cfg["steering"]["subgoal_probability"] + ) + include_subtask = not ( + use_goal + and torch.rand((), generator=generator) + < cfg["steering"]["subtask_dropout_with_goal"] + ) + use_history = bool( + torch.rand((), generator=generator) >= cfg["memory"]["history_dropout"] + ) + sequence = tokenizer.build_policy_sequence( + "drive safely", + sample["reasoning"] if include_subtask else None, + sample["actions"], + metadata, + include_episode_metadata=True, + include_subtask=include_subtask, + ) + actions = sample["actions"][None].to(device) + delay = sample_rtc_delay( + generator=generator, + max_delay_steps=cfg["rtc"]["max_delay_steps"], + action_horizon=dataset.action_horizon, + ) + noisy, times, velocity, postfix_mask = flow.training_batch( + actions, delays=torch.tensor([delay], device=device) + ) + output = policy( + sample["image"][None].to(device), + sequence.input_ids[None].to(device), + noisy, + times, + history_pixel_values=sample["history_images"][None].to(device) + if use_history + else None, + history_states=sample["history_states"][None].to(device) + if use_history + else None, + current_state=sample["current_state"][None].to(device), + subgoal_pixel_values=sample["future_observation"][None].to(device) + if use_goal + else None, + action_prefix_mask=sequence.action_prefix_mask[None].to(device), + text_labels=sequence.labels[None].to(device), + ) + flow_loss = flow.masked_loss(output.vector_field, velocity, postfix_mask) + loss = output.text_loss + cfg["train"]["flow_loss_weight"] * flow_loss + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(policy.parameters(), cfg["train"]["grad_clip_norm"]) + optimizer.step() + if step == 1 or step % cfg["train"]["log_every"] == 0: + row = { + "stage": "policy", + "step": step, + "loss": float(loss.detach().cpu()), + "text_loss": float(output.text_loss.detach().cpu()), + "flow_loss": float(flow_loss.detach().cpu()), + "rtc_delay": delay, + "use_goal": use_goal, + "use_history": use_history, + "include_subtask": include_subtask, + } + history.append(row) + print(json.dumps(row)) + + checkpoint = { + "policy": policy.state_dict(), + "high_level": high_level.state_dict(), + "goal": goal_generator.state_dict(), + "policy_config": policy_config.__dict__, + "config": cfg, + "action_statistics": dataset.action_normalizer.as_dict(), + } + checkpoint_path = output_dir / "checkpoint.pt" + torch.save(checkpoint, checkpoint_path) + write_json(output_dir / "training_history.json", history) + write_json( + output_dir / "summary.json", + { + "device": str(device), + "samples": len(dataset), + "policy_parameters": sum(parameter.numel() for parameter in policy.parameters()), + "goal_parameters": sum( + parameter.numel() for parameter in goal_generator.parameters() + ), + "high_level_parameters": sum( + parameter.numel() for parameter in high_level.parameters() + ), + }, + ) + return checkpoint_path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", default="configs/overfit.yaml") + train(parser.parse_args().config) + + +if __name__ == "__main__": + main() diff --git a/tests/test_context.py b/tests/test_context.py index 0951622..46a3be4 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -6,8 +6,8 @@ from pi07mini.context import ( PolicyMetadata, build_ambiguous_history_pair, drop_prompt_components, - endpoint_subgoal_heatmap, format_policy_prompt, + future_observation_proxy, make_history_proxy_frames, prompt_components, ) @@ -18,19 +18,13 @@ def _history() -> torch.Tensor: 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_future_subgoal_has_image_contract_not_coordinate_heatmap() -> None: + image = torch.linspace(0.0, 1.0, 3 * 24 * 24).reshape(3, 24, 24) + future = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1) + subgoal = future_observation_proxy(image, _history(), future) + assert subgoal.shape == image.shape + assert torch.isfinite(subgoal).all() + assert not torch.equal(subgoal, image) def test_history_proxy_is_oldest_to_newest_and_future_free() -> None: diff --git a/tests/test_data.py b/tests/test_data.py index f5a7c1c..bdc407c 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -76,7 +76,6 @@ def test_policy_variants_and_metadata() -> None: "quality": 5, "mistake": False, "control_mode": "ee", - "source_type": "demonstration", } assert torch.allclose(policies[-1].future_xy[-1], policies[-1].future_xy[4]) @@ -105,12 +104,13 @@ def test_dataset_contract_and_roundtrip(tmp_path: Path) -> None: 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["future_observation"].shape == (3, 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["rtc_action_prefix"].shape == (16, 2) + assert item["rtc_prefix_mask"].shape == (16,) + assert item["rtc_postfix_mask"].shape == (16,) assert item["reasoning"] == "Steer left." assert item["prompt_component_mask"].all() assert item["token_metadata"]["quality"] == 5 @@ -131,4 +131,4 @@ def test_dataset_is_deterministic_per_index(tmp_path: Path) -> None: 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"]) + torch.testing.assert_close(first["rtc_action_prefix"], second["rtc_action_prefix"]) diff --git a/tests/test_model.py b/tests/test_model.py index 01acfed..cffc106 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -31,7 +31,7 @@ def tiny_model() -> Pi07Mini: def sample_sequence(): tokenizer = MiniTokenizer() - metadata = EpisodeMetadata(2000, 5, False, "ee", "demonstration") + metadata = EpisodeMetadata(2000, 5, False, "ee") return tokenizer, tokenizer.build_policy_sequence( "drive", "steer_left", torch.zeros(2, 2), metadata ) @@ -52,16 +52,17 @@ def common_inputs(sequence): def test_block_causal_attention_matches_pi07_modal_order() -> None: - layout = PrefixLayout(observation_tokens=2, goal_tokens=2, text_tokens=3) + layout = PrefixLayout( + observation_tokens=2, prompt_tokens=2, goal_tokens=2, fast_tokens=2 + ) 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() + assert not edges[0, 2:].any() + assert edges[2, :3].all() and not edges[2, 3:].any() # causal prompt + assert edges[4, :6].all() and not edges[4, 6:].any() # goal reads obs + text + goal + assert edges[6, :7].all() and not edges[6, 7] # FAST is causal after goal + assert edges[7].all() def test_prefix_has_fixed_mem_tokens_and_one_token_per_state() -> None: @@ -78,9 +79,19 @@ def test_prefix_has_fixed_mem_tokens_and_one_token_per_state() -> None: ) assert layout.observation_tokens == 1 + 2 + 1 # MEM + history states + current assert layout.goal_tokens == 1 + assert layout.prompt_tokens == sequence.action_start_index + assert layout.fast_tokens == sequence.input_ids.numel() - sequence.action_start_index assert prefix.shape[1] == layout.total_tokens +def test_action_expert_accepts_per_token_rtc_timesteps() -> None: + _, sequence = sample_sequence() + inputs = common_inputs(sequence) + inputs["timesteps"] = torch.tensor([[1.0, 0.25]]) + output = tiny_model()(**inputs) + assert output.vector_field.shape == (1, 2, 2) + + def test_flow_cannot_read_teacher_forced_fast_tokens() -> None: torch.manual_seed(12) tokenizer, sequence = sample_sequence() diff --git a/tests/test_rtc.py b/tests/test_rtc.py index c233e8a..c2ab81e 100644 --- a/tests/test_rtc.py +++ b/tests/test_rtc.py @@ -3,36 +3,39 @@ from __future__ import annotations import pytest import torch -from pi07mini.rtc import align_actions_for_rtc, sample_rtc_delay +from pi07mini.flow import FlowMatching +from pi07mini.rtc import build_rtc_conditioning, 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: +def test_zero_delay_has_empty_prefix_and_full_postfix() -> None: + rtc = build_rtc_conditioning(_actions(), 0) + assert not rtc.prefix_mask.any() + assert rtc.postfix_mask.all() + assert torch.count_nonzero(rtc.action_prefix) == 0 + + +def test_delay_keeps_ground_truth_prefix_without_shifting_target() -> 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) + rtc = build_rtc_conditioning(actions, 12) + torch.testing.assert_close(rtc.action_prefix[:12], actions[:12]) + assert torch.count_nonzero(rtc.action_prefix[12:]) == 0 + assert rtc.prefix_mask.tolist() == [True] * 12 + [False] * 4 + assert rtc.postfix_mask.tolist() == [False] * 12 + [True] * 4 -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_flow_training_sets_prefix_time_to_one_and_masks_its_loss() -> None: + actions = _actions()[None] + noisy, times, _, loss_mask = FlowMatching().training_batch( + actions, delays=torch.tensor([3]), generator=torch.Generator().manual_seed(4) + ) + torch.testing.assert_close(noisy[:, :3], actions[:, :3]) + torch.testing.assert_close(times[:, :3], torch.ones(1, 3)) + assert not loss_mask[:, :3].any() + assert loss_mask[:, 3:].all() def test_delay_sampling_covers_closed_zero_to_twelve_range() -> None: @@ -47,4 +50,4 @@ def test_delay_sampling_covers_closed_zero_to_twelve_range() -> None: @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) + build_rtc_conditioning(_actions(), delay) diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py index 55f299c..1a6b234 100644 --- a/tests/test_tokenizer.py +++ b/tests/test_tokenizer.py @@ -9,7 +9,6 @@ def metadata() -> EpisodeMetadata: quality=5, mistake=False, control_mode="ee", - source_type="autonomous_rollout", ) @@ -21,10 +20,18 @@ def test_prompt_contains_rich_metadata_and_control_tokens() -> None: 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 + assert len(metadata_indices) == 6 + + +def test_high_level_prompt_contains_past_subtasks_and_next_query() -> None: + tokenizer = MiniTokenizer() + prompt = tokenizer.encode_high_level_prompt( + "drive", ("keep_lane", "slow_down"), metadata() + ) + assert prompt[-1] == tokenizer.NEXT_SUBTASK_MARKER + assert torch.count_nonzero(prompt == tokenizer.SUBTASK_MARKER) == 2 def test_metadata_cfg_pair_drops_metadata_but_retains_control_mode() -> None: @@ -62,4 +69,3 @@ def test_fast_lite_dct_round_trip_has_bounded_error() -> None: 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 -