Align pi0.7 minimal structure with paper

2026-07-12 20:37:34 +08:00
parent b4c8db846b
commit cd4e8e6a3f
21 changed files with 1485 additions and 660 deletions
@@ -1,246 +1,125 @@
 # π0.7 Minimal Validation
-这是一个面向**结构验证和小样本过拟合**的 π0.7-style 教学项目。它在前面的
+这是一个面向**结构验证和极小样本过拟合**的 π0.7 最小实现。目标不是复现参数量、预训练权重或论文指标，而是保留官方论文中的关键模块、模态顺序、注意力信息流和训练/推理闭环。
 π0.6/RECAP 最小实验之上，重点验证一条可解释的 steerable policy 链路：利用
 短时历史形成 memory，在高层预测视觉 subgoal，通过可独立缺省的 policy
 metadata 控制动作风格，并在存在控制延迟时用 RTC 对齐动作 chunk。
-本项目不是 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
+history images --MEM video encoder--┐
-              ↓
+history/current states -------------> observation block
       TinyMEMEncoder ───────────────┐
                                     │
-current image + task ────────────────┼→ TinyGoalGenerator → predicted subgoal
+task + predicted subtask + metadata -> causal text prompt
                                     │
-optional speed/quality/mistake/control metadata
+generated future observation -------> image-goal block
-                                     ↓
+                                     │
-                              HighLevelPolicy
+FAST action tokens -----------------> causal discrete action branch
-                                     ↓
+                                     │
-noisy action + flow time + RTC prefix → Pi07Mini action expert
+clean RTC prefix + noisy postfix ---> flow action expert
                                     ↓
                          continuous action chunk
 ```
-## 重点结构
+官方 Figure 19 对应的 prefix 顺序被明确实现为：
 ### 1. MEM：短时历史而非单帧
 PAI 的极小 manifest 只有一张相机图像，本项目依据已观测的过去 XY 对该图像做
 确定性的微小平移和曝光变化，构造短时视觉 history proxy。它不是实际多帧视频，
 只能验证 temporal memory 的接口。
 `ambiguous_history_pair` 刻意构造两个“当前状态相同、历史路径不同”的序列。
 完整模型应比 `no_history` 消融更能区分它们。这里的 MEM 仅表示本仓库的 temporal
 memory encoder；在官方资料缺失时，不推断该缩写的官方全称。
 ### 2. Goal generator：只用 endpoint 监督
 训练目标是 future trajectory 的最后一个 XY endpoint，经高斯渲染得到 heatmap。
 API 只接受 endpoint，不接受整条 future trajectory，从而避免把低层动作答案作为
 视觉 subgoal 条件泄漏给 action expert。推理和最终评估必须使用**模型预测的 goal**；
 真值 heatmap 只可用于训练监督和 oracle 诊断。
 ### 3. High-level 与 steerable metadata
 最小数据由一条 expert demonstration 确定性构造五种 policy mode：
 | Mode | speed | quality | mistake | control |
 |---|---|---|---|---|
 | expert | normal | expert | none | human |
 | fast | fast | good | none | autonomous |
 | slow | slow | good | none | autonomous |
 | understeer | normal | poor | understeer | autonomous |
 | stop | slow | poor | stop | autonomous |
 训练时 `speed`、`quality`、`mistake`、`control` 独立做 metadata dropout，任务文本
 始终保留。这既检验属性是否能单独控制行为，也提供 CFG 所需的 omitted condition。
 这些合成标签不是人类偏好数据，也不能证明真实世界可控性。
 ### 4. RTC：控制延迟下的 action chunk 对齐
 若推理耗时等价于已经执行 `d` 个 action，模型不应再次预测旧 chunk 的开头：
 ```text
-original target: a0 a1 ... a15
+memory observations -> causal text -> bidirectional image goals -> causal FAST
 delay d=4:       a4 a5 ... a15 PAD PAD PAD PAD
 valid mask:       1  1 ...   1   0   0   0   0
 ```
-默认随机采样 `d=0..12`，返回已执行 prefix，并对 padding tail 屏蔽损失。RTC 在
+- 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
+高层策略输入当前 observation、MEM history、总体任务、metadata，以及过去的 subtask history，输出下一条 subtask。推理时该**预测 subtask**同时用于：
-与 metadata-omitted 预测：
+
 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 只依赖当前图像与过去状态；
+- visual subgoal 出现在 25% policy examples；
- subgoal condition 只使用 endpoint target，不能编码完整 action chunk；
+- 有 visual subgoal 时，subtask 以 30% 概率移除；
- flow expert 训练时不能读取 teacher-forced 真值动作 token；
+- episode metadata 整体以 15% 概率移除；
- 所有模型对比使用同一 initial noise、归一化统计和 RTC delay；
+- speed、quality、mistake 分别再以 5% 概率移除；
- 数据增强的随机数由 sample index 和 seed 决定，保证复现实验；
+- control mode 始终保留；
- 仓库不包含或重新分发原始 PAI 数据。
+- 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
+1. high-quality segments 上的 visual goal flow；
-python -m pi07mini.infer --checkpoint outputs/overfit/checkpoint.pt
+2. next-subtask high-level policy；
-```
+3. 带 prompt dropout、FAST、flow、KI 和 training-time RTC 的低层 policy。
-最终展示不能只有 full model。相同样本、noise seed 和 RTC delay 下至少比较：
+评估输出：
-| 实验 | 移除内容 | 要回答的问题 |
+- predicted subtask；
-|---|---|---|
+- generated visual subgoal；
-| `full` | 无 | 完整链路能否过拟合并重建对应 mode 的轨迹？ |
+- full/no-metadata/no-goal/no-history/CFG 消融；
-| `no_metadata` | speed/quality/mistake/control | policy mode 差异是否真的来自 steerable condition？ |
+- delay `0/4/8/12` 的 RTC postfix action error；
-| `no_goal` | predicted goal | endpoint condition 是否改善几何方向和终点？ |
+- trajectory 和 normalized action MSE。
 | `no_history` | MEM history，只保留当前 observation | temporal history 是否解决相同当前状态的歧义？ |
-还必须提供：
+## 当前最小训练结果
- predicted goal heatmap 与 ground-truth endpoint heatmap；
+默认配置在 Apple MPS 上完成 `800 / 200 / 900` 步训练：
 - 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 上的差异指标。
-小样本过拟合中，某个消融偶然更好并不证明模块无效；应展示全部固定设置结果，
+- high-level subtask：`steer_left`，与目标一致；
-不调换 noise seed 或只挑选最佳样本。
+- 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 机制，只加载可信来源的文件。
-## 能得出的结论
+- 本项目只允许离线训练和可视化，严禁连接真实车辆或机器人执行器。
-
+- future observation 是单帧数据上的代理构造，不是真实多视角未来视频。
-成功实验最多能说明：本仓库定义的 memory、endpoint goal、high-level metadata、
+- 小样本过拟合、loss 下降或轨迹重合不代表泛化、闭环稳定性或安全性。
-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
 编号或官方结果表。
@@ -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
+  target: future_observation_proxy
  heatmap_sigma: 1.5
  xy_extent: [-8.0, 64.0, -16.0, 16.0]
  endpoint_only_target: true
 high_level:
  enabled: true
@@ -54,10 +50,11 @@ high_level:
    - control
 steering:
-  metadata_dropout: 0.25
+  subgoal_probability: 0.25
-  preserve_task_prompt: true
+  subtask_dropout_with_goal: 0.3
-  cfg_training_dropout: 0.25
+  metadata_full_dropout: 0.15
-  cfg_beta: 2.0
+  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
@@ -1,11 +1,39 @@
 {
  "target_subtask": "steer_left",
  "predicted_subtask": "steer_left",
  "goal_image_mse": 0.0246433112770319,
  "trajectory_metrics": {
-    "full_oracle_goal": 0.47031369805336,
+    "full_generated_goal": 9.339226722717285,
-    "generated_goal": 0.4905388057231903,
+    "full_oracle_goal": 9.344671249389648,
-    "no_metadata": 2.457864284515381,
+    "no_metadata": 4.243795394897461,
-    "no_goal": 4.042560577392578,
+    "no_goal": 9.066949844360352,
-    "no_history": 2.1593737602233887,
+    "no_history": 10.404154777526855,
-    "cfg_beta2": 10.576497077941895
+    "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
    }
  }
 }
@@ -1,7 +1,7 @@
 {
  "device": "mps",
  "samples": 5,
-  "policy_parameters": 1088930,
+  "policy_parameters": 1088354,
-  "goal_parameters": 75731,
+  "goal_parameters": 75347,
-  "high_level_parameters": 652800
+  "high_level_parameters": 652224
 }
@@ -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,
+    "loss": 51.165855407714844,
-    "text_loss": 44.21609115600586,
+    "text_loss": 46.97157287597656,
-    "flow_loss": 3.4136481285095215
+    "flow_loss": 2.0971415042877197,
    "rtc_delay": 5,
    "use_goal": true,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 20,
-    "loss": 26.20608901977539,
+    "loss": 11.228275299072266,
-    "text_loss": 5.128565311431885,
+    "text_loss": 4.843214988708496,
-    "flow_loss": 10.538762092590332
+    "flow_loss": 3.192530393600464,
    "rtc_delay": 1,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 40,
-    "loss": 13.018914222717285,
+    "loss": 6.2679572105407715,
-    "text_loss": 3.846059799194336,
+    "text_loss": 2.9056875705718994,
-    "flow_loss": 4.586427211761475
+    "flow_loss": 1.681134819984436,
    "rtc_delay": 4,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 60,
-    "loss": 11.344329833984375,
+    "loss": 6.005075931549072,
-    "text_loss": 3.335771083831787,
+    "text_loss": 2.709343671798706,
-    "flow_loss": 4.004279136657715
+    "flow_loss": 1.647866129875183,
    "rtc_delay": 2,
    "use_goal": true,
    "use_history": true,
    "include_subtask": false
  },
  {
    "stage": "policy",
    "step": 80,
-    "loss": 6.586862087249756,
+    "loss": 9.903374671936035,
-    "text_loss": 2.5455708503723145,
+    "text_loss": 2.334010362625122,
-    "flow_loss": 2.0206456184387207
+    "flow_loss": 3.784682273864746,
    "rtc_delay": 4,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 100,
-    "loss": 11.200895309448242,
+    "loss": 5.646769046783447,
-    "text_loss": 2.1739449501037598,
+    "text_loss": 2.3033981323242188,
-    "flow_loss": 4.513474941253662
+    "flow_loss": 1.6716854572296143,
    "rtc_delay": 12,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 120,
-    "loss": 17.273035049438477,
+    "loss": 6.867326259613037,
-    "text_loss": 2.1746408939361572,
+    "text_loss": 1.850703239440918,
-    "flow_loss": 7.549196720123291
+    "flow_loss": 2.5083115100860596,
    "rtc_delay": 12,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 140,
-    "loss": 11.353687286376953,
+    "loss": 3.6174087524414062,
-    "text_loss": 2.2218563556671143,
+    "text_loss": 2.0811572074890137,
-    "flow_loss": 4.565915584564209
+    "flow_loss": 0.7681257128715515,
    "rtc_delay": 11,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 160,
-    "loss": 14.110220909118652,
+    "loss": 4.300071716308594,
-    "text_loss": 2.3123977184295654,
+    "text_loss": 1.9097824096679688,
-    "flow_loss": 5.898911476135254
+    "flow_loss": 1.1951446533203125,
    "rtc_delay": 8,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 180,
-    "loss": 19.999370574951172,
+    "loss": 8.075389862060547,
-    "text_loss": 2.749119997024536,
+    "text_loss": 1.6174044609069824,
-    "flow_loss": 8.62512493133545
+    "flow_loss": 3.2289929389953613,
    "rtc_delay": 2,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 200,
-    "loss": 8.102465629577637,
+    "loss": 5.758746147155762,
-    "text_loss": 1.631801962852478,
+    "text_loss": 1.5621910095214844,
-    "flow_loss": 3.2353317737579346
+    "flow_loss": 2.0982775688171387,
    "rtc_delay": 11,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 220,
-    "loss": 21.903854370117188,
+    "loss": 2.9783213138580322,
-    "text_loss": 1.6063504219055176,
+    "text_loss": 1.389180302619934,
-    "flow_loss": 10.148752212524414
+    "flow_loss": 0.7945705056190491,
    "rtc_delay": 11,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 240,
-    "loss": 10.197904586791992,
+    "loss": 5.832153797149658,
-    "text_loss": 1.5640149116516113,
+    "text_loss": 1.4576387405395508,
-    "flow_loss": 4.3169450759887695
+    "flow_loss": 2.1872575283050537,
    "rtc_delay": 0,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 260,
-    "loss": 6.899936676025391,
+    "loss": 3.4007065296173096,
-    "text_loss": 1.3081820011138916,
+    "text_loss": 0.8826411366462708,
-    "flow_loss": 2.795877456665039
+    "flow_loss": 1.2590327262878418,
    "rtc_delay": 5,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 280,
-    "loss": 10.051774978637695,
+    "loss": 4.018008232116699,
-    "text_loss": 1.2130903005599976,
+    "text_loss": 0.9588160514831543,
-    "flow_loss": 4.419342517852783
+    "flow_loss": 1.529596209526062,
    "rtc_delay": 11,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 300,
-    "loss": 7.963553428649902,
+    "loss": 1.492077112197876,
-    "text_loss": 1.1586928367614746,
+    "text_loss": 0.9946327209472656,
-    "flow_loss": 3.402430295944214
+    "flow_loss": 0.24872222542762756,
    "rtc_delay": 9,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 320,
-    "loss": 5.177718639373779,
+    "loss": 5.338322639465332,
-    "text_loss": 1.2673476934432983,
+    "text_loss": 0.8486652970314026,
-    "flow_loss": 1.9551855325698853
+    "flow_loss": 2.244828701019287,
    "rtc_delay": 8,
    "use_goal": true,
    "use_history": true,
    "include_subtask": false
  },
  {
    "stage": "policy",
    "step": 340,
-    "loss": 4.487771034240723,
+    "loss": 2.0994033813476562,
-    "text_loss": 1.0361082553863525,
+    "text_loss": 0.8355546593666077,
-    "flow_loss": 1.7258312702178955
+    "flow_loss": 0.6319243907928467,
    "rtc_delay": 7,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 360,
-    "loss": 4.356203556060791,
+    "loss": 2.2518701553344727,
-    "text_loss": 1.3223515748977661,
+    "text_loss": 0.5455462336540222,
-    "flow_loss": 1.5169259309768677
+    "flow_loss": 0.8531619906425476,
    "rtc_delay": 10,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 380,
-    "loss": 5.289548397064209,
+    "loss": 7.911384105682373,
-    "text_loss": 1.4994667768478394,
+    "text_loss": 0.6893162727355957,
-    "flow_loss": 1.89504075050354
+    "flow_loss": 3.6110339164733887,
    "rtc_delay": 7,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 400,
-    "loss": 8.761286735534668,
+    "loss": 2.3619725704193115,
-    "text_loss": 1.4391229152679443,
+    "text_loss": 0.5672968626022339,
-    "flow_loss": 3.6610820293426514
+    "flow_loss": 0.8973378539085388,
    "rtc_delay": 10,
    "use_goal": true,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 420,
-    "loss": 2.400188446044922,
+    "loss": 1.233857274055481,
-    "text_loss": 1.0106937885284424,
+    "text_loss": 0.661384642124176,
-    "flow_loss": 0.694747269153595
+    "flow_loss": 0.28623631596565247,
    "rtc_delay": 12,
    "use_goal": true,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 440,
-    "loss": 9.683706283569336,
+    "loss": 3.295243740081787,
-    "text_loss": 1.044368863105774,
+    "text_loss": 0.4776339828968048,
-    "flow_loss": 4.319668769836426
+    "flow_loss": 1.4088048934936523,
    "rtc_delay": 9,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 460,
-    "loss": 10.443679809570312,
+    "loss": 1.7052180767059326,
-    "text_loss": 0.7045990228652954,
+    "text_loss": 0.5183412432670593,
-    "flow_loss": 4.869540214538574
+    "flow_loss": 0.5934383869171143,
    "rtc_delay": 8,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 480,
-    "loss": 6.514384746551514,
+    "loss": 8.46796703338623,
-    "text_loss": 1.185524582862854,
+    "text_loss": 0.6230667233467102,
-    "flow_loss": 2.6644301414489746
+    "flow_loss": 3.922450065612793,
    "rtc_delay": 0,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 500,
-    "loss": 2.255455255508423,
+    "loss": 2.5979185104370117,
-    "text_loss": 1.0136079788208008,
+    "text_loss": 0.27633967995643616,
-    "flow_loss": 0.620923638343811
+    "flow_loss": 1.1607893705368042,
    "rtc_delay": 1,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 520,
-    "loss": 4.1700439453125,
+    "loss": 2.012068510055542,
-    "text_loss": 1.1829472780227661,
+    "text_loss": 0.30755916237831116,
-    "flow_loss": 1.4935483932495117
+    "flow_loss": 0.8522546291351318,
    "rtc_delay": 1,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 540,
-    "loss": 6.956455230712891,
+    "loss": 5.146071910858154,
-    "text_loss": 1.1096038818359375,
+    "text_loss": 0.4677533805370331,
-    "flow_loss": 2.9234256744384766
+    "flow_loss": 2.3391592502593994,
    "rtc_delay": 0,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 560,
-    "loss": 4.5965800285339355,
+    "loss": 1.9480068683624268,
-    "text_loss": 1.5257400274276733,
+    "text_loss": 0.9942646026611328,
-    "flow_loss": 1.5354200601577759
+    "flow_loss": 0.476871132850647,
    "rtc_delay": 7,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 580,
-    "loss": 1.7948781251907349,
+    "loss": 7.889751434326172,
-    "text_loss": 0.835124135017395,
+    "text_loss": 0.6651982665061951,
-    "flow_loss": 0.4798769950866699
+    "flow_loss": 3.612276554107666,
    "rtc_delay": 3,
    "use_goal": true,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 600,
-    "loss": 11.411609649658203,
+    "loss": 1.454559087753296,
-    "text_loss": 0.7508105635643005,
+    "text_loss": 0.5516519546508789,
-    "flow_loss": 5.330399513244629
+    "flow_loss": 0.4514535963535309,
    "rtc_delay": 11,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 620,
-    "loss": 14.173189163208008,
+    "loss": 4.760444641113281,
-    "text_loss": 0.553065836429596,
+    "text_loss": 0.3332846760749817,
-    "flow_loss": 6.810061454772949
+    "flow_loss": 2.2135798931121826,
    "rtc_delay": 7,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 640,
-    "loss": 2.423476457595825,
+    "loss": 6.838594913482666,
-    "text_loss": 0.6566628813743591,
+    "text_loss": 0.33531704545021057,
-    "flow_loss": 0.8834067583084106
+    "flow_loss": 3.251638889312744,
    "rtc_delay": 3,
    "use_goal": true,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 660,
-    "loss": 7.33753776550293,
+    "loss": 4.01173210144043,
-    "text_loss": 0.5569484829902649,
+    "text_loss": 0.842496395111084,
-    "flow_loss": 3.3902945518493652
+    "flow_loss": 1.5846179723739624,
    "rtc_delay": 1,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 680,
-    "loss": 7.2152018547058105,
+    "loss": 5.610707759857178,
-    "text_loss": 0.6606324911117554,
+    "text_loss": 0.45004716515541077,
-    "flow_loss": 3.277284622192383
+    "flow_loss": 2.5803303718566895,
    "rtc_delay": 6,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 700,
-    "loss": 2.5858616828918457,
+    "loss": 2.415992021560669,
-    "text_loss": 0.471678227186203,
+    "text_loss": 0.2757202684879303,
-    "flow_loss": 1.0570917129516602
+    "flow_loss": 1.0701358318328857,
    "rtc_delay": 3,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 720,
-    "loss": 5.87769889831543,
+    "loss": 3.6413190364837646,
-    "text_loss": 0.5711548924446106,
+    "text_loss": 0.15377745032310486,
-    "flow_loss": 2.6532719135284424
+    "flow_loss": 1.7437708377838135,
    "rtc_delay": 2,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 740,
-    "loss": 6.787323951721191,
+    "loss": 2.4192750453948975,
-    "text_loss": 0.5064371228218079,
+    "text_loss": 0.7513293027877808,
-    "flow_loss": 3.1404433250427246
+    "flow_loss": 0.8339728713035583,
    "rtc_delay": 4,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 760,
-    "loss": 11.166234970092773,
+    "loss": 1.3346192836761475,
-    "text_loss": 0.5802748203277588,
+    "text_loss": 0.2767994999885559,
-    "flow_loss": 5.292980194091797
+    "flow_loss": 0.5289098620414734,
    "rtc_delay": 10,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 780,
-    "loss": 3.563460350036621,
+    "loss": 2.4135358333587646,
-    "text_loss": 0.6946841478347778,
+    "text_loss": 0.16017527878284454,
-    "flow_loss": 1.4343880414962769
+    "flow_loss": 1.1266802549362183,
    "rtc_delay": 3,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 800,
-    "loss": 3.2471487522125244,
+    "loss": 2.966399908065796,
-    "text_loss": 0.44556593894958496,
+    "text_loss": 0.2566848397254944,
-    "flow_loss": 1.4007914066314697
+    "flow_loss": 1.3548575639724731,
    "rtc_delay": 3,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 820,
-    "loss": 3.1103086471557617,
+    "loss": 1.8670570850372314,
-    "text_loss": 0.7132699489593506,
+    "text_loss": 0.23726917803287506,
-    "flow_loss": 1.1985193490982056
+    "flow_loss": 0.8148939609527588,
    "rtc_delay": 4,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 840,
-    "loss": 4.753307342529297,
+    "loss": 1.6185007095336914,
-    "text_loss": 0.48397523164749146,
+    "text_loss": 0.27975988388061523,
-    "flow_loss": 2.1346659660339355
+    "flow_loss": 0.6693704128265381,
    "rtc_delay": 5,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 860,
-    "loss": 3.5492496490478516,
+    "loss": 0.9501249194145203,
-    "text_loss": 0.366057813167572,
+    "text_loss": 0.4190976619720459,
-    "flow_loss": 1.5915958881378174
+    "flow_loss": 0.2655136287212372,
    "rtc_delay": 7,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 880,
-    "loss": 4.0379414558410645,
+    "loss": 1.5665624141693115,
-    "text_loss": 0.6511648893356323,
+    "text_loss": 0.317945659160614,
-    "flow_loss": 1.6933883428573608
+    "flow_loss": 0.6243083477020264,
    "rtc_delay": 5,
    "use_goal": false,
    "use_history": false,
    "include_subtask": true
  },
  {
    "stage": "policy",
    "step": 900,
-    "loss": 1.7170453071594238,
+    "loss": 1.8227444887161255,
-    "text_loss": 0.2673894464969635,
+    "text_loss": 0.32235953211784363,
-    "flow_loss": 0.7248279452323914
+    "flow_loss": 0.7501924633979797,
    "rtc_delay": 8,
    "use_goal": false,
    "use_history": true,
    "include_subtask": true
  }
 ]
@@ -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(
+def future_observation_proxy(
-    endpoint_xy: torch.Tensor,
+    image: torch.Tensor,
-    *,
+    history_xy: torch.Tensor,
-    size: int = 32,
+    future_xy: torch.Tensor,
    sigma: float = 1.5,
    xy_extent: tuple[float, float, float, float] = (-8.0, 64.0, -16.0, 16.0),
 ) -> torch.Tensor:
-    """Render a normalized Gaussian subgoal heatmap from one endpoint only.
+    """Create a tiny future-camera proxy instead of encoding coordinates as a heatmap.
-    It deliberately accepts no trajectory tensor. This API makes it difficult
+    The source dataset contains only the current frame.  For the architecture
-    to accidentally leak the entire ground-truth action chunk into the visual
+    validation we therefore warp that frame according to the segment endpoint,
-    subgoal condition.
+    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)
+    image = torch.as_tensor(image, dtype=torch.float32)
-    if size < 2 or sigma <= 0.0:
+    history = torch.as_tensor(history_xy, dtype=torch.float32)
-        raise ValueError("size must be >=2 and sigma must be positive")
+    future = torch.as_tensor(future_xy, dtype=torch.float32)
-    x_min, x_max, y_min, y_max = xy_extent
+    if image.ndim != 3 or image.shape[0] not in {1, 3}:
-    if not x_min < x_max or not y_min < y_max:
+        raise ValueError("image must have shape [C,H,W]")
-        raise ValueError("xy_extent must contain increasing bounds")
+    if history.ndim != 2 or future.ndim != 2 or history.shape[-1] != 2 or future.shape[-1] != 2:
-    px = (endpoint[0] - x_min) / (x_max - x_min) * (size - 1)
+        raise ValueError("history_xy and future_xy must have shape [T,2]")
-    py = (y_max - endpoint[1]) / (y_max - y_min) * (size - 1)
+    displacement = future[-1] - history[-1]
-    yy, xx = torch.meshgrid(
+    scale = max(image.shape[-2:]) / 32.0
-        torch.arange(size, dtype=torch.float32),
+    dx = int(torch.round(displacement[1] * scale).clamp(-image.shape[-1] // 5, image.shape[-1] // 5))
-        torch.arange(size, dtype=torch.float32),
+    dy = int(torch.round(-displacement[0] * scale * 0.08).clamp(-image.shape[-2] // 6, image.shape[-2] // 6))
-        indexing="ij",
+    distance = torch.linalg.vector_norm(displacement)
-    )
+    exposure = float((1.0 + 0.02 * torch.tanh(distance / 8.0)).item())
-    heatmap = torch.exp(-((xx - px) ** 2 + (yy - py) ** 2) / (2.0 * sigma**2))
+    shifted = torch.roll(image, shifts=(dy, dx), dims=(-2, -1)) * exposure
-    heatmap = heatmap / heatmap.amax().clamp_min(torch.finfo(heatmap.dtype).eps)
+    return F.avg_pool2d(shifted[None], kernel_size=3, stride=1, padding=1)[0].clamp(0.0, 1.0)
    return heatmap.unsqueeze(0)
 def make_history_proxy_frames(
@@ -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(
+        future_observation = future_observation_proxy(
-            policy.future_xy[-1], size=self.image_size, sigma=max(1.0, self.image_size / 24.0)
+            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,
+            "future_observation": future_observation,
-            "subgoal_pixel_values": subgoal.expand(3, -1, -1).clone(),
+            "subgoal_pixel_values": future_observation.clone(),
            "rtc_delay_steps": delay,
-            "rtc_target_actions": rtc.delayed_target_actions,
+            "rtc_action_prefix": rtc.action_prefix,
-            "rtc_valid_mask": rtc.valid_mask,
+            "rtc_prefix_mask": rtc.prefix_mask,
-            "rtc_executed_prefix_actions": rtc.executed_prefix_actions,
+            "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}",
        }
@@ -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
+        self.inference_steps = int(inference_steps)
-    def training_batch(self,target:Tensor)->tuple[Tensor,Tensor,Tensor]:
+        self.cutoff = float(cutoff)
-        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)
+    def training_batch(
-        return t[:,None,None]*target+(1-t[:,None,None])*noise,t,target-noise
+        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:
+    def sample(
-        count=steps or self.inference_steps;x=torch.randn(shape,device=device);grid=torch.linspace(0,1,count+1,device=device)
+        self,
-        for i in range(count): x=x+(grid[i+1]-grid[i])*fn(x,grid[i].expand(shape[0]))
+        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
@@ -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 .model import HighLevelPolicy, Pi07Config, Pi07Mini, TinyGoalGenerator
-from .runtime import load_config,choose_device,write_json
+from .rtc import build_rtc_conditioning
-from .tokenizer import MiniTokenizer,EpisodeMetadata
+from .runtime import choose_device, write_json
 from .tokenizer import EpisodeMetadata, MiniTokenizer
@torch.no_grad()
-def evaluate(checkpoint_path):
+def evaluate(checkpoint_path: str) -> tuple[Path, 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'])
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
- 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)
+    cfg = checkpoint["config"]
- 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()))
+    root = Path(__file__).resolve().parents[2]
- full=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta);ids=full.input_ids[:full.action_start_index][None].to(dev)
+    device = choose_device(cfg["device"])
- torch.manual_seed(cfg['seed']);g=torch.randn_like(s['subgoal_pixel_values'][None]).to(dev);grid=torch.linspace(0,1,6,device=dev)
+    tokenizer = MiniTokenizer()
- for i in range(5): g=g+(grid[i+1]-grid[i])*goal(image,ids,g,grid[i].expand(1)).vector_field
+    dataset = Pi07SteerableDataset(
- generated_goal=g.cpu()[0]
+        root / "data/overfit/overfit_manifest.jsonl",
- def sample(name,include_meta=True,use_goal=True,use_history=True,cfg_beta=None,generated=False):
+        image_size=cfg["data"]["image_size"],
-  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)
+        history_frames=cfg["data"]["history_frames"],
-  goal_pixels=(generated_goal if generated else s['subgoal_pixel_values'])[None].to(dev) if use_goal else None
+        action_horizon=cfg["data"]["action_horizon"],
-  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)
+        action_statistics=checkpoint["action_statistics"],
-  def field(x,t):
+        max_rtc_delay_steps=cfg["rtc"]["max_delay_steps"],
-   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)
+    policy_config = Pi07Config(**checkpoint["policy_config"])
-   return policy.metadata_cfg(xids,uids,guidance_scale=cfg_beta,noisy_actions=x,timesteps=t,**common)
+    policy = Pi07Mini(policy_config)
-  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()
+    policy.load_state_dict(checkpoint["policy"])
- 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)}
+    policy.to(device).eval()
- 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()}
+    high_level = HighLevelPolicy(policy_config)
- out=Path(checkpoint_path).parent;fig,ax=plt.subplots(figsize=(9,6));ax.plot(target[:,0],target[:,1],'o-',label='expert target',linewidth=2)
+    high_level.load_state_dict(checkpoint["high_level"])
- for k,v in curves.items():ax.plot(v[:,0],v[:,1],'x--',label=f'{k} ({metrics[k]:.3f})')
+    high_level.to(device).eval()
- 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)
+    goal_generator = TinyGoalGenerator(tokenizer.vocab_size)
- 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');
+    goal_generator.load_state_dict(checkpoint["goal"])
- for a in axes:a.axis('off')
+    goal_generator.to(device).eval()
- 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
+    flow = FlowMatching(cfg["flow"]["integration_steps"])
-def main():
+
- p=argparse.ArgumentParser();p.add_argument('--checkpoint',default='outputs/overfit/checkpoint.pt');evaluate(p.parse_args().checkpoint)
+    sample = next(
-if __name__=='__main__':main()
+        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()
@@ -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:
+    def goal_start(self) -> int:
-        return self.observation_tokens + self.goal_tokens
+        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.ones(layout.prompt_tokens, dtype=torch.long, device=prefix.device),
-                torch.full(
+                torch.full((layout.goal_tokens,), 2, dtype=torch.long, device=prefix.device),
-                    (layout.text_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.
+        """Build the paper's observation -> text -> image-goal -> FAST mask."""
        Observation queries only read observations.  Goal-image queries read
        observations and all goal tokens bidirectionally.  Text queries read
        both visual blocks and earlier text tokens causally.
        """
        length = layout.total_tokens
        if valid.shape[1] != length:
@@ -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
+        obs = q < layout.prompt_start
-        goal = (q >= layout.goal_start) & (q < layout.text_start)
+        prompt = (q >= layout.prompt_start) & (q < layout.goal_start)
-        text = q >= layout.text_start
+        goal = (q >= layout.goal_start) & (q < layout.fast_start)
        fast = q >= layout.fast_start
        structural = (
-            (obs & (k < layout.goal_start))
+            (obs & (k < layout.prompt_start))
-            | (goal & (k < layout.text_start))
+            | (
-            | (text & ((k < layout.text_start) | ((k >= layout.text_start) & (k <= q))))
+                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(
+            supplied = torch.cat((
-                (
+                torch.ones(batch, layout.observation_tokens, dtype=torch.bool, device=valid.device),
-                    torch.ones(
+                supplied[:, : layout.prompt_tokens],
-                        batch,
+                torch.ones(batch, layout.goal_tokens, dtype=torch.bool, device=valid.device),
-                        layout.text_start,
+                supplied[:, layout.prompt_tokens :],
-                        dtype=torch.bool,
+            ), dim=1)
                        device=valid.device,
                    ),
                    supplied,
                ),
                dim=1,
            )
        if supplied.shape != valid.shape:
            raise ValueError("action_prefix_mask must cover text or the full prefix")
        readable = valid & supplied
@@ -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
@@ -9,27 +9,22 @@ MAX_RTC_DELAY_STEPS = 12
@dataclass(frozen=True)
-class RTCAlignment:
+class RTCConditioning:
-    """A fixed-width target aligned to an inference/control delay."""
+    """Hard action-prefix conditioning used by training-time RTC."""
-    delayed_target_actions: torch.Tensor
+    action_prefix: torch.Tensor
-    valid_mask: torch.Tensor
+    prefix_mask: torch.Tensor
-    executed_prefix_actions: 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:
+) -> RTCConditioning:
-    """Shift targets left so prediction step zero starts after control delay.
+    """Pad the committed action prefix to the full prediction horizon."""
    For delay ``d``, target position ``i`` supervises original action ``i+d``.
    Vacated tail positions are zero padded and excluded by ``valid_mask``.
    The already-executed prefix is returned for optional action-prefix input.
    """
    actions = torch.as_tensor(actions, dtype=torch.float32)
    if actions.ndim < 2 or actions.shape[-2] < 1:
@@ -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)
+    prefix_mask = torch.arange(horizon, device=actions.device) < delay_steps
-    remaining = horizon - delay_steps
+    prefix_mask = prefix_mask.expand(actions.shape[:-2] + (horizon,))
-    aligned[..., :remaining, :] = actions[..., delay_steps:, :]
+    action_prefix = torch.zeros_like(actions)
-    valid[..., :remaining] = True
+    action_prefix[..., :delay_steps, :] = actions[..., :delay_steps, :]
-    prefix = actions[..., :delay_steps, :].clone()
+    return RTCConditioning(
-    return RTCAlignment(aligned, valid, prefix, int(delay_steps))
+        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())
@@ -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):
+import numpy as np
-    with Path(path).open() as f:return yaml.safe_load(f)
+import torch
-def seed_everything(seed:int):random.seed(seed);np.random.seed(seed);torch.manual_seed(seed)
+import yaml
-def choose_device(name="auto"):
+
-    if name!="auto":return torch.device(name)
+
-    if torch.cuda.is_available():return torch.device("cuda")
+def load_config(path: str | Path) -> dict:
-    if torch.backends.mps.is_available():return torch.device("mps")
+    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",
    )
@@ -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 = (
+    ACTION_START = CONTROL_EE + 1
        "demonstration",
        "autonomous",
        "intervention",
        "human_video",
        "web",
        "unknown",
    )
    SOURCE_TOKEN_START = CONTROL_EE + 1
    SOURCE_TOKEN_END = SOURCE_TOKEN_START + len(SOURCE_LABELS) - 1
    ACTION_START = SOURCE_TOKEN_END + 1
    ACTION_END = ACTION_START + 1
    ACTION_TOKEN_START = ACTION_END + 1
    ACTION_BINS = 64
@@ -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
@@ -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 .flow import FlowMatching
-from .model import Pi07Config,Pi07Mini,HighLevelPolicy,TinyGoalGenerator
+from .model import HighLevelPolicy, Pi07Config, Pi07Mini, TinyGoalGenerator
-from .runtime import load_config,seed_everything,choose_device,write_json
+from .rtc import sample_rtc_delay
-from .tokenizer import MiniTokenizer,EpisodeMetadata
+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):
+def build_config(cfg: dict, tokenizer: MiniTokenizer) -> Pi07Config:
- cfg=load_config(config_path);root=Path(__file__).resolve().parents[2];seed_everything(cfg['seed']);dev=choose_device(cfg['device'])
+    model = cfg["model"]
- manifest=materialize_subset(cfg['data']['source_manifest'],root/'data/overfit',max_samples=cfg['data']['max_samples'])
+    return Pi07Config(
- 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'])
+        vocab_size=tokenizer.vocab_size,
- 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=[]
+        action_dim=2,
- # Goal generator
+        state_dim=2,
- opt=torch.optim.AdamW(goal.parameters(),lr=cfg['train']['learning_rate'])
+        image_size=cfg["data"]["image_size"],
- for step in range(1,cfg['train']['goal_steps']+1):
+        patch_size=model["patch_size"],
-  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()
+        d_model=model["d_model"],
-  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))
+        n_heads=model["n_heads"],
- # High-level semantic policy
+        n_layers=model["n_layers"],
- opt=torch.optim.AdamW(high.parameters(),lr=cfg['train']['learning_rate'])
+        memory_layers=model.get("memory_layers", 2),
- for step in range(1,cfg['train']['high_level_steps']+1):
+        mlp_ratio=model["mlp_ratio"],
-  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()
+        max_history_frames=cfg["data"]["history_frames"],
-  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))
+        max_text_tokens=model["max_text_tokens"],
- # Full steerable policy, alternating prompt-dropout ablations and RTC targets.
+        max_action_tokens=cfg["data"]["action_horizon"],
- opt=torch.optim.AdamW(policy.parameters(),lr=cfg['train']['learning_rate'])
+        action_start_token_id=tokenizer.ACTION_START,
- for step in range(1,cfg['train']['policy_steps']+1):
+        dropout=model["dropout"],
-  s=ds[(step-1)%len(ds)];meta=EpisodeMetadata(**s['token_metadata']);include_meta=(step%4)!=0;seq=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=include_meta);ids=seq.input_ids[None].to(dev);labels=seq.labels[None].to(dev);target=s['rtc_target_actions'][None].to(dev);noisy,t,vel=flow.training_batch(target);use_goal=(step%4)!=1;use_history=(step%5)!=1
+    )
  o=policy(s['image'][None].to(dev),ids,noisy,t,history_pixel_values=s['history_images'][None].to(dev) if use_history else None,history_states=s['history_states'][None].to(dev) if use_history else None,current_state=s['current_state'][None].to(dev),subgoal_pixel_values=s['subgoal_pixel_values'][None].to(dev) if use_goal else None,action_prefix_mask=seq.action_prefix_mask[None].to(dev),text_labels=labels)
  valid=s['rtc_valid_mask'][None].to(dev)[:,:,None];flow_loss=((o.vector_field-vel).square()*valid).sum()/valid.sum().clamp_min(1);loss=o.text_loss+cfg['train']['flow_loss_weight']*flow_loss;opt.zero_grad();loss.backward();torch.nn.utils.clip_grad_norm_(policy.parameters(),cfg['train']['grad_clip_norm']);opt.step()
  if step==1 or step%cfg['train']['log_every']==0:row={'stage':'policy','step':step,'loss':float(loss.detach().cpu()),'text_loss':float(o.text_loss.detach().cpu()),'flow_loss':float(flow_loss.detach().cpu())};history.append(row);print(json.dumps(row))
 ck={'policy':policy.state_dict(),'high_level':high.state_dict(),'goal':goal.state_dict(),'policy_config':pcfg.__dict__,'config':cfg,'action_statistics':ds.action_normalizer.as_dict()};torch.save(ck,out/'checkpoint.pt');write_json(out/'training_history.json',history);write_json(out/'summary.json',{'device':str(dev),'samples':len(ds),'policy_parameters':sum(p.numel() for p in policy.parameters()),'goal_parameters':sum(p.numel() for p in goal.parameters()),'high_level_parameters':sum(p.numel() for p in high.parameters())});return out/'checkpoint.pt'
-def main():
+
- p=argparse.ArgumentParser();p.add_argument('--config',default='configs/overfit.yaml');train(p.parse_args().config)
+def _metadata_with_paper_dropout(
-if __name__=='__main__':main()
+    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()
@@ -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:
+def test_future_subgoal_has_image_contract_not_coordinate_heatmap() -> None:
-    endpoint = torch.tensor([20.0, 4.0])
+    image = torch.linspace(0.0, 1.0, 3 * 24 * 24).reshape(3, 24, 24)
-    first = endpoint_subgoal_heatmap(endpoint, size=24)
+    future = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1)
-    # Two future paths with radically different interiors produce exactly the
+    subgoal = future_observation_proxy(image, _history(), future)
-    # same condition because only their shared endpoint is passed to the API.
+    assert subgoal.shape == image.shape
-    path_a = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1)
+    assert torch.isfinite(subgoal).all()
-    path_b = path_a.clone()
+    assert not torch.equal(subgoal, image)
    path_b[1:-1, 1] += 100.0
    second = endpoint_subgoal_heatmap(path_b[-1], size=24)
    torch.testing.assert_close(first, second, atol=0.0, rtol=0.0)
    assert first.shape == (1, 24, 24)
    assert float(first.max()) == 1.0
    assert torch.isfinite(first).all()
 def test_history_proxy_is_oldest_to_newest_and_future_free() -> None:
@@ -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_action_prefix"].shape == (16, 2)
-    assert item["rtc_valid_mask"].shape == (16,)
+    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"])
@@ -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 not edges[0, 2:].any()
-    assert edges[2, :4].all()
+    assert edges[2, :3].all() and not edges[2, 3:].any()  # causal prompt
-    assert not edges[2, 4:].any()  # goal reads observation + goal, not text
+    assert edges[4, :6].all() and not edges[4, 6:].any()  # goal reads obs + text + goal
-    assert edges[4, :5].all()
+    assert edges[6, :7].all() and not edges[6, 7]  # FAST is causal after goal
-    assert not edges[4, 5:].any()  # first text token is causal
+    assert edges[7].all()
    assert edges[6].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()
@@ -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)
+    rtc = build_rtc_conditioning(actions, 12)
-    torch.testing.assert_close(aligned.delayed_target_actions, actions)
+    torch.testing.assert_close(rtc.action_prefix[:12], actions[:12])
-    assert aligned.valid_mask.all()
+    assert torch.count_nonzero(rtc.action_prefix[12:]) == 0
-    assert aligned.executed_prefix_actions.shape == (0, 2)
+    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:
+def test_flow_training_sets_prefix_time_to_one_and_masks_its_loss() -> None:
-    actions = _actions()
+    actions = _actions()[None]
-    aligned = align_actions_for_rtc(actions, 12)
+    noisy, times, _, loss_mask = FlowMatching().training_batch(
-    torch.testing.assert_close(aligned.delayed_target_actions[:4], actions[12:])
+        actions, delays=torch.tensor([3]), generator=torch.Generator().manual_seed(4)
-    assert torch.count_nonzero(aligned.delayed_target_actions[4:]) == 0
+    )
-    assert aligned.valid_mask.tolist() == [True] * 4 + [False] * 12
+    torch.testing.assert_close(noisy[:, :3], actions[:, :3])
-    torch.testing.assert_close(aligned.executed_prefix_actions, actions[:12])
+    torch.testing.assert_close(times[:, :3], torch.ones(1, 3))
-
+    assert not loss_mask[:, :3].any()
-
+    assert loss_mask[:, 3:].all()
 def test_rtc_alignment_preserves_batch_dimensions() -> None:
    actions = torch.stack((_actions(), _actions() + 100.0))
    aligned = align_actions_for_rtc(actions, 3)
    assert aligned.delayed_target_actions.shape == (2, 16, 2)
    assert aligned.valid_mask.shape == (2, 16)
    torch.testing.assert_close(aligned.delayed_target_actions[:, :13], actions[:, 3:])
 def test_delay_sampling_covers_closed_zero_to_twelve_range() -> None:
@@ -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)
@@ -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