docs: clarify minimal reproduction overview

2026-07-12 22:00:11 +08:00
13 changed files with 174 additions and 1522 deletions
@@ -5,6 +5,5 @@ __pycache__/
 *.egg-info/
 .venv/
 data/
 notebook_artifacts/
 outputs/**/*.pt
 .DS_Store
@@ -5,23 +5,22 @@
 ## 实现内容
 ```text
-图像 + 离散本体状态 + 总体任务指令
+图像 + 本体状态
-                  │
+        │
-                  ▼
+        ▼
-         高层语义 subtask 预测
+小型视觉/语言主干 ──→ 高层语义 subtask
-                  │ predicted subtask
+        │                    │
-                  ▼
+        └──── semantic prefix┘
-图像 + 离散本体状态 + 低层 subtask command
+                     │
-                  │
+ noisy action + time │
- noisy action + time
+          └──────────▼
          └───────▼
       action expert（双向 action attention）
-                  │
+                     │
-                  ▼
+                     ▼
-       flow velocity → action chunk
+          flow velocity → action chunk
 ```
-默认任务使用 `64×64` 前视相机图像、离散化历史 XY 状态、总体任务指令和一句 subtask，预测未来 `16×2` 的二维轨迹。样本少且重复训练，因此成功标准是“能够记忆训练集”，而不是泛化。
+默认任务使用 `64×64` 前视相机图像、离散化历史 XY 状态和一句 subtask，预测未来 `16×2` 的二维轨迹。样本少且重复训练，因此成功标准是“能够记忆训练集”，而不是泛化。
 ## 与 π0.5 的对应关系
@@ -33,8 +32,6 @@
 - flow matching 训练连续动作，推理时从高斯噪声做 Euler 积分；
 - timestep 通过 adaptive RMSNorm 调制 action expert；
 - 同时优化高层语义目标和低层连续动作目标。
 - 高层调用显式读取总体任务指令；低层调用只读取 observation 与预测 subtask，
  对应论文分解 `π(action | observation, subtask) · π(subtask | observation, task)`；
 - VLM 自回归输出高层 subtask 以及 64-bin 离散动作 token；
 - 推理时分别还原 VLM token 轨迹和 flow/diffusion 轨迹，与真值三方比较。
@@ -42,7 +39,7 @@
 - 小型自定义视觉编码器和 Transformer 取代 PaliGemma 2B；
 - 隐藏宽度、层数、图像尺寸和 action horizon 大幅缩小；
- 使用固定小词表的 task/subtask 分类与生成，而非开放式语言理解；
+- 使用固定小词表的 subtask 分类/生成，而非开放式语言理解；
 - 使用二维驾驶轨迹，取代多机器人、多相机和高维关节控制；
 - 默认使用 1 个 PAI 驾驶样本的 4 个视觉增强版本，没有机器人或网页联合训练数据；
 - 使用正交 DCT-II、低频优先排列和 64-bin 量化实现 FAST-lite；保留全部频率系数，因此只验证表示与训练接口，不复现官方压缩率；
@@ -70,13 +67,6 @@ python -m pip install -e '.[dev]'
 如果该路径不同，请先修改 `configs/overfit.yaml` 的 `data.source_manifest`。
 ## 一键 Notebook 复现
 [`notebooks/pi05_minimal_reproduction.ipynb`](notebooks/pi05_minimal_reproduction.ipynb)
 提供从环境安装、论文与架构说明、数据/权重导入、测试、完整训练、推理到结果图的
 顺序执行流程。Notebook 支持本地 manifest、数据 zip URL 或自动生成的确定性最小
 样本；默认从零训练，因此不要求预训练权重。
 ## 运行测试
 ```bash
@@ -105,7 +95,7 @@ pi05mini-train --config configs/overfit.yaml
 python -m pi05mini.train --config configs/overfit.yaml
 ```
-配置默认先进行 300 步离散 VLM/FAST-lite 预训练，再进行 500 步联合 post-training。Stage 1 不更新 flow action expert；Stage 2 使用同一模型的两个调用联合优化：总体任务条件下的 token CE，以及预测/teacher-forced subtask 条件下的 flow loss。CPU 可以运行；Apple Silicon 可自动选择 MPS，CUDA 可自动选择 GPU。
+配置默认先进行 300 步离散 VLM/FAST-lite 预训练，再进行 500 步联合 post-training。Stage 1 不更新 flow action expert；Stage 2 联合优化 token CE 与 flow loss。CPU 可以运行；Apple Silicon 可自动选择 MPS，CUDA 可自动选择 GPU。
 ## 推理与效果图
@@ -2,37 +2,33 @@
  {
    "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug0",
    "target_subtask": "steer_left",
    "task_instruction": "Follow the visible road.",
    "predicted_subtask": "steer_left",
    "vlm_action_tokens_valid": true,
    "vlm_trajectory_mse": 0.07854792475700378,
-    "diffusion_trajectory_mse": 0.14029410481452942
+    "diffusion_trajectory_mse": 0.27604275941848755
  },
  {
    "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug1",
    "target_subtask": "steer_left",
    "task_instruction": "Follow the visible road.",
    "predicted_subtask": "steer_left",
    "vlm_action_tokens_valid": true,
    "vlm_trajectory_mse": 0.025470610707998276,
-    "diffusion_trajectory_mse": 0.11049529910087585
+    "diffusion_trajectory_mse": 0.20049458742141724
  },
  {
    "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug2",
    "target_subtask": "steer_left",
    "task_instruction": "Follow the visible road.",
    "predicted_subtask": "steer_left",
    "vlm_action_tokens_valid": true,
    "vlm_trajectory_mse": 0.013774838298559189,
-    "diffusion_trajectory_mse": 0.04682297259569168
+    "diffusion_trajectory_mse": 0.03129641339182854
  },
  {
    "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug3",
    "target_subtask": "steer_left",
    "task_instruction": "Follow the visible road.",
    "predicted_subtask": "steer_left",
    "vlm_action_tokens_valid": true,
    "vlm_trajectory_mse": 0.059195056557655334,
-    "diffusion_trajectory_mse": 0.1004234328866005
+    "diffusion_trajectory_mse": 0.11321215331554413
  }
 ]
@@ -1,19 +1,19 @@
 {
  "device": "mps",
  "samples": 4,
-  "parameters": 817346,
+  "parameters": 816482,
  "initial": {
    "stage": "pretrain",
    "step": 1,
-    "loss": 80.38571166992188,
+    "loss": 78.87449645996094,
-    "text_loss": 80.38571166992188,
+    "text_loss": 78.87449645996094,
    "flow_loss": null
  },
  "final": {
    "stage": "posttrain",
    "step": 500,
-    "loss": 0.1929057091474533,
+    "loss": 0.12353827804327011,
-    "text_loss": 2.784368575703411e-07,
+    "text_loss": 3.1691649837739533e-06,
-    "flow_loss": 0.09645271301269531
+    "flow_loss": 0.06176755577325821
  }
 }
@@ -2,295 +2,295 @@
  {
    "stage": "pretrain",
    "step": 1,
-    "loss": 80.38571166992188,
+    "loss": 78.87449645996094,
-    "text_loss": 80.38571166992188,
+    "text_loss": 78.87449645996094,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 20,
-    "loss": 2.5035173892974854,
+    "loss": 2.475086212158203,
-    "text_loss": 2.5035173892974854,
+    "text_loss": 2.475086212158203,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 40,
-    "loss": 1.3843837976455688,
+    "loss": 1.1025534868240356,
-    "text_loss": 1.3843837976455688,
+    "text_loss": 1.1025534868240356,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 60,
-    "loss": 0.8159888386726379,
+    "loss": 0.601564884185791,
-    "text_loss": 0.8159888386726379,
+    "text_loss": 0.601564884185791,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 80,
-    "loss": 0.19992661476135254,
+    "loss": 0.25571808218955994,
-    "text_loss": 0.19992661476135254,
+    "text_loss": 0.25571808218955994,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 100,
-    "loss": 0.05571526661515236,
+    "loss": 0.03314117714762688,
-    "text_loss": 0.05571526661515236,
+    "text_loss": 0.03314117714762688,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 120,
-    "loss": 0.005689455661922693,
+    "loss": 0.005310459993779659,
-    "text_loss": 0.005689455661922693,
+    "text_loss": 0.005310459993779659,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 140,
-    "loss": 0.00010565178672550246,
+    "loss": 3.0251239877543412e-05,
-    "text_loss": 0.00010565178672550246,
+    "text_loss": 3.0251239877543412e-05,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 160,
-    "loss": 4.098461431567557e-05,
+    "loss": 1.1354763955750968e-05,
-    "text_loss": 4.098461431567557e-05,
+    "text_loss": 1.1354763955750968e-05,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 180,
-    "loss": 2.716351991693955e-05,
+    "loss": 6.638075319642667e-06,
-    "text_loss": 2.716351991693955e-05,
+    "text_loss": 6.638075319642667e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 200,
-    "loss": 2.2414897102862597e-05,
+    "loss": 5.348983449948719e-06,
-    "text_loss": 2.2414897102862597e-05,
+    "text_loss": 5.348983449948719e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 220,
-    "loss": 1.9143404642818496e-05,
+    "loss": 4.736786195280729e-06,
-    "text_loss": 1.9143404642818496e-05,
+    "text_loss": 4.736786195280729e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 240,
-    "loss": 1.6641104593873024e-05,
+    "loss": 4.316162630857434e-06,
-    "text_loss": 1.6641104593873024e-05,
+    "text_loss": 4.316162630857434e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 260,
-    "loss": 1.4626367374148685e-05,
+    "loss": 3.975574600190157e-06,
-    "text_loss": 1.4626367374148685e-05,
+    "text_loss": 3.975574600190157e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 280,
-    "loss": 1.2977556252735667e-05,
+    "loss": 3.6877775073662633e-06,
-    "text_loss": 1.2977556252735667e-05,
+    "text_loss": 3.6877775073662633e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 300,
-    "loss": 1.16078899736749e-05,
+    "loss": 3.4408508327032905e-06,
-    "text_loss": 1.16078899736749e-05,
+    "text_loss": 3.4408508327032905e-06,
    "flow_loss": null
  },
  {
    "stage": "posttrain",
    "step": 1,
-    "loss": 3.8964109420776367,
+    "loss": 8.707193374633789,
-    "text_loss": 1.1545746929186862e-05,
+    "text_loss": 3.428078571232618e-06,
-    "flow_loss": 1.94819974899292
+    "flow_loss": 4.353594779968262
  },
  {
    "stage": "posttrain",
    "step": 20,
-    "loss": 4.304055690765381,
+    "loss": 3.533726453781128,
-    "text_loss": 2.3414058887283318e-05,
+    "text_loss": 0.14151602983474731,
-    "flow_loss": 2.1520161628723145
+    "flow_loss": 1.6961052417755127
  },
  {
    "stage": "posttrain",
    "step": 40,
-    "loss": 2.925170660018921,
+    "loss": 3.0850143432617188,
-    "text_loss": 2.142304765584413e-06,
+    "text_loss": 0.00014649856893811375,
-    "flow_loss": 1.4625842571258545
+    "flow_loss": 1.5424339771270752
  },
  {
    "stage": "posttrain",
    "step": 60,
-    "loss": 1.7996485233306885,
+    "loss": 1.3556954860687256,
-    "text_loss": 5.1970469030493405e-06,
+    "text_loss": 0.00011281041224719957,
-    "flow_loss": 0.8998216390609741
+    "flow_loss": 0.6777913570404053
  },
  {
    "stage": "posttrain",
    "step": 80,
-    "loss": 2.0647871494293213,
+    "loss": 1.4165724515914917,
-    "text_loss": 3.251374073443003e-05,
+    "text_loss": 6.599419066333212e-06,
-    "flow_loss": 1.0323773622512817
+    "flow_loss": 0.7082829475402832
  },
  {
    "stage": "posttrain",
    "step": 100,
-    "loss": 2.3402059078216553,
+    "loss": 0.8609042167663574,
-    "text_loss": 2.50748803409806e-06,
+    "text_loss": 8.11900736152893e-06,
-    "flow_loss": 1.1701016426086426
+    "flow_loss": 0.430448055267334
  },
  {
    "stage": "posttrain",
    "step": 120,
-    "loss": 1.4503896236419678,
+    "loss": 0.5642163753509521,
-    "text_loss": 6.828942105130409e-07,
+    "text_loss": 7.576879397674929e-06,
-    "flow_loss": 0.7251944541931152
+    "flow_loss": 0.28210440278053284
  },
  {
    "stage": "posttrain",
    "step": 140,
-    "loss": 0.23238907754421234,
+    "loss": 0.7110214233398438,
-    "text_loss": 1.3709052382182563e-07,
+    "text_loss": 6.249698344618082e-06,
-    "flow_loss": 0.1161944717168808
+    "flow_loss": 0.35550758242607117
  },
  {
    "stage": "posttrain",
    "step": 160,
-    "loss": 0.46604833006858826,
+    "loss": 0.18725578486919403,
-    "text_loss": 2.1542784622852196e-07,
+    "text_loss": 2.876289045161684e-06,
-    "flow_loss": 0.23302406072616577
+    "flow_loss": 0.09362645447254181
  },
  {
    "stage": "posttrain",
    "step": 180,
-    "loss": 0.169603168964386,
+    "loss": 0.33923929929733276,
-    "text_loss": 2.418236704215815e-07,
+    "text_loss": 6.748496616637567e-06,
-    "flow_loss": 0.08480146527290344
+    "flow_loss": 0.16961628198623657
  },
  {
    "stage": "posttrain",
    "step": 200,
-    "loss": 0.7555881142616272,
+    "loss": 0.3003162443637848,
-    "text_loss": 8.855541011598689e-08,
+    "text_loss": 9.495882295595948e-06,
-    "flow_loss": 0.3777940273284912
+    "flow_loss": 0.15015336871147156
  },
  {
    "stage": "posttrain",
    "step": 220,
-    "loss": 0.10211063176393509,
+    "loss": 0.23716998100280762,
-    "text_loss": 1.3623892414216243e-07,
+    "text_loss": 5.660519946104614e-06,
-    "flow_loss": 0.05105524882674217
+    "flow_loss": 0.11858215928077698
  },
  {
    "stage": "posttrain",
    "step": 240,
-    "loss": 0.4649242162704468,
+    "loss": 0.42717865109443665,
-    "text_loss": 2.2819931189133058e-07,
+    "text_loss": 8.400202204938978e-06,
-    "flow_loss": 0.23246198892593384
+    "flow_loss": 0.21358512341976166
  },
  {
    "stage": "posttrain",
    "step": 260,
-    "loss": 0.280495822429657,
+    "loss": 0.11702287942171097,
-    "text_loss": 1.3198149417803506e-07,
+    "text_loss": 4.122769951209193e-06,
-    "flow_loss": 0.14024785161018372
+    "flow_loss": 0.058509379625320435
  },
  {
    "stage": "posttrain",
    "step": 280,
-    "loss": 0.5054064989089966,
+    "loss": 0.19245445728302002,
-    "text_loss": 9.451582627661992e-08,
+    "text_loss": 5.941245035501197e-06,
-    "flow_loss": 0.2527031898498535
+    "flow_loss": 0.09622425585985184
  },
  {
    "stage": "posttrain",
    "step": 300,
-    "loss": 0.19181671738624573,
+    "loss": 0.18475297093391418,
-    "text_loss": 4.453279984772962e-07,
+    "text_loss": 4.744139005197212e-06,
-    "flow_loss": 0.09590813517570496
+    "flow_loss": 0.09237411618232727
  },
  {
    "stage": "posttrain",
    "step": 320,
-    "loss": 0.20589734613895416,
+    "loss": 0.38481318950653076,
-    "text_loss": 2.8524956974251836e-07,
+    "text_loss": 3.1461640901397914e-06,
-    "flow_loss": 0.10294853150844574
+    "flow_loss": 0.19240501523017883
  },
  {
    "stage": "posttrain",
    "step": 340,
-    "loss": 0.09927964955568314,
+    "loss": 0.2163764089345932,
-    "text_loss": 1.0132762895409542e-07,
+    "text_loss": 4.1934181354008615e-06,
-    "flow_loss": 0.04963977262377739
+    "flow_loss": 0.10818611085414886
  },
  {
    "stage": "posttrain",
    "step": 360,
-    "loss": 0.2213846743106842,
+    "loss": 0.23533868789672852,
-    "text_loss": 1.2942695093443035e-07,
+    "text_loss": 3.3019753118423978e-06,
-    "flow_loss": 0.11069227010011673
+    "flow_loss": 0.11766768991947174
  },
  {
    "stage": "posttrain",
    "step": 380,
-    "loss": 0.9335335493087769,
+    "loss": 0.08397223055362701,
-    "text_loss": 1.3964488232431904e-07,
+    "text_loss": 8.292082384286914e-06,
-    "flow_loss": 0.46676671504974365
+    "flow_loss": 0.04198196902871132
  },
  {
    "stage": "posttrain",
    "step": 400,
-    "loss": 0.12127648293972015,
+    "loss": 0.13249839842319489,
-    "text_loss": 7.578287153364727e-08,
+    "text_loss": 1.0871653103095014e-05,
-    "flow_loss": 0.06063820421695709
+    "flow_loss": 0.06624376028776169
  },
  {
    "stage": "posttrain",
    "step": 420,
-    "loss": 0.1503632515668869,
+    "loss": 0.4479130506515503,
-    "text_loss": 4.0871746875836834e-08,
+    "text_loss": 1.0128642315976322e-05,
-    "flow_loss": 0.07518160343170166
+    "flow_loss": 0.22395145893096924
  },
  {
    "stage": "posttrain",
    "step": 440,
-    "loss": 0.15399529039859772,
+    "loss": 0.22416232526302338,
-    "text_loss": 2.0435869885204738e-08,
+    "text_loss": 2.6114355478057405e-06,
-    "flow_loss": 0.07699763774871826
+    "flow_loss": 0.11207985877990723
  },
  {
    "stage": "posttrain",
    "step": 460,
-    "loss": 0.18908466398715973,
+    "loss": 0.06346385926008224,
-    "text_loss": 4.172323286866231e-08,
+    "text_loss": 3.1035808660817565e-06,
-    "flow_loss": 0.09454230964183807
+    "flow_loss": 0.031730376183986664
  },
  {
    "stage": "posttrain",
    "step": 480,
-    "loss": 0.2329331636428833,
+    "loss": 0.385486900806427,
-    "text_loss": 2.8099325533048614e-08,
+    "text_loss": 2.184032382501755e-06,
-    "flow_loss": 0.11646656692028046
+    "flow_loss": 0.19274236261844635
  },
  {
    "stage": "posttrain",
    "step": 500,
-    "loss": 0.1929057091474533,
+    "loss": 0.12353827804327011,
-    "text_loss": 2.784368575703411e-07,
+    "text_loss": 3.1691649837739533e-06,
-    "flow_loss": 0.09645271301269531
+    "flow_loss": 0.06176755577325821
  }
 ]
@@ -22,10 +22,6 @@ dependencies = [
 dev = [
  "pytest>=8.0",
 ]
 notebook = [
  "jupyterlab>=4.2",
  "nbclient>=0.10",
 ]
 [project.scripts]
 pi05mini-train = "pi05mini.train:main"
@@ -19,7 +19,6 @@ from .tokenizer import MiniTokenizer
 class SampleRecord:
    clip_id: str
    t0_us: int
    task_instruction: str
    reasoning: str
    image_path: Path
    trajectory_path: Path
@@ -112,9 +111,6 @@ def load_manifest(path: str | Path) -> list[SampleRecord]:
            SampleRecord(
                clip_id=str(row["clip_id"]),
                t0_us=int(row["t0_us"]),
                task_instruction=str(
                    row.get("task_instruction", MiniTokenizer.DEFAULT_TASK_INSTRUCTION)
                ),
                reasoning=str(row.get("reasoning", "")),
                image_path=_resolve_data_path(manifest_path, row["image_path"]),
                trajectory_path=_resolve_data_path(manifest_path, row["trajectory_path"]),
@@ -160,7 +156,6 @@ def materialize_subset(
            {
                "clip_id": record.clip_id,
                "t0_us": record.t0_us,
                "task_instruction": record.task_instruction,
                "reasoning": record.reasoning,
                "image_path": str(image_path.relative_to(destination)),
                "trajectory_path": str(trajectory_path.relative_to(destination)),
@@ -308,13 +303,7 @@ class Pi05OverfitDataset(Dataset[dict[str, Any]]):
        actions = self.action_normalizer.encode(history_xy, future_xy)
        return {
            "image": self._load_image(base_index, variant),
-            "prefix_tokens": self.tokenizer.build_prefix(
+            "prefix_tokens": self.tokenizer.build_prefix(history_xy),
                history_xy, record.task_instruction
            ),
            "low_level_prefix_tokens": self.tokenizer.build_low_level_prefix(
                history_xy, target
            ),
            "task_instruction": record.task_instruction,
            "subtask_target": target,
            "subtask_label": self.tokenizer.subtasks.decode(target),
            "actions": actions,
@@ -40,20 +40,13 @@ def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path:
        generated: list[int] = []
        expected_tokens = 1 + 2 + dataset.action_horizon * 2
        current_tokens = tokens
-        for generation_step in range(expected_tokens):
+        for _ in range(expected_tokens):
            semantic = model(
                image, current_tokens, zero_action,
                torch.zeros(1, device=device),
                generated_token_start=tokens.shape[1]
            )
-            if generation_step == 0:
+            next_token = int(semantic.text_logits[0, -1].argmax().cpu())
                candidates = torch.tensor(
                    dataset.tokenizer.subtasks.token_ids, device=device
                )
                candidate_logits = semantic.text_logits[0, -1, candidates]
                next_token = int(candidates[candidate_logits.argmax()].cpu())
            else:
                next_token = int(semantic.text_logits[0, -1].argmax().cpu())
            generated.append(next_token)
            current_tokens = torch.cat(
                (current_tokens, torch.tensor([[next_token]], device=device)), dim=1
@@ -69,12 +62,13 @@ def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path:
        except ValueError:
            vlm_xy = np.full_like(sample["future_xy"].numpy(), np.nan)
            vlm_valid = False
-        conditioned_tokens = dataset.tokenizer.build_low_level_prefix(
+        conditioned_tokens = torch.cat(
-            sample["history_xy"], predicted_token
+            (tokens, torch.tensor([[predicted_token]], device=device)), dim=1
-        )[None].to(device)
+        )
        def velocity(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
            return model(
-                image, conditioned_tokens, x, t
+                image, conditioned_tokens, x, t,
                generated_token_start=tokens.shape[1]
            ).vector_field
        normalized = flow.sample(
            velocity, (1, dataset.action_horizon, 2), device,
@@ -94,7 +88,6 @@ def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path:
            predicted_label = f"token_{predicted_token}"
        rows.append({
            "sample_id": sample["sample_id"], "target_subtask": target_label,
            "task_instruction": sample["task_instruction"],
            "predicted_subtask": predicted_label,
            "vlm_action_tokens_valid": vlm_valid,
            "vlm_trajectory_mse": vlm_mse,
@@ -52,61 +52,6 @@ class CoordinateTokenizer:
        return low + encoded.float() / (self.bins - 1) * (high - low)
 class TaskVocabulary:
    """Small overall-task vocabulary for the high-level pi0.5 inference pass."""
    labels = (
        "follow_road",
        "clean_kitchen",
        "clean_bedroom",
        "pick_place",
        "fold_laundry",
        "open_close",
        "unknown",
    )
    def __init__(self, first_token_id: int = 260) -> None:
        self.first_token_id = int(first_token_id)
        self._label_to_token = {
            label: self.first_token_id + index for index, label in enumerate(self.labels)
        }
        self._token_to_label = {token: label for label, token in self._label_to_token.items()}
    @property
    def token_ids(self) -> tuple[int, ...]:
        return tuple(self._label_to_token[label] for label in self.labels)
    def classify(self, instruction: str) -> str:
        text = instruction.casefold()
        if any(term in text for term in ("road", "lane", "drive", "waypoint", "steer")):
            return "follow_road"
        if any(term in text for term in ("fold", "folding")):
            return "fold_laundry"
        if any(term in text for term in ("kitchen", "dish", "sink", "counter")):
            return "clean_kitchen"
        if any(term in text for term in ("bedroom", "bed", "pillow", "blanket")):
            return "clean_bedroom"
        if any(term in text for term in ("open", "close", "drawer", "cabinet", "fridge")):
            return "open_close"
        if any(term in text for term in ("pick", "place", "put", "move")):
            return "pick_place"
        return instruction if instruction in self._label_to_token else "unknown"
    def encode(self, label_or_instruction: str) -> int:
        label = (
            label_or_instruction
            if label_or_instruction in self._label_to_token
            else self.classify(label_or_instruction)
        )
        return self._label_to_token[label]
    def decode(self, token_id: int | torch.Tensor) -> str:
        try:
            return self._token_to_label[int(token_id)]
        except KeyError as error:
            raise ValueError(f"unknown task token: {int(token_id)}") from error
 class SubtaskVocabulary:
    """Small semantic vocabulary standing in for pi0.5's generated language subtask."""
@@ -169,13 +114,7 @@ class SubtaskVocabulary:
 class MiniTokenizer:
-    """Compose hierarchical task/state, subtask, and FAST-lite action tokens.
+    """Compose image/state, semantic, and FAST-lite action tokens.
    pi0.5 performs two semantically different calls with the same model.  The
    high-level call reads the overall task and predicts a semantic subtask.  The
    low-level call deliberately replaces the overall task with that predicted
    subtask, matching ``pi(a | observation, predicted_subtask)`` rather than
    allowing continuous actions to depend directly on the overall task.
    Actions are z-score-normalized by the dataset before reaching this
    tokenizer.  A complete ``[horizon, 2]`` chunk is transformed along time
@@ -191,19 +130,11 @@ class MiniTokenizer:
    high-frequency coefficients.
    """
    DEFAULT_TASK_INSTRUCTION = "Follow the visible road."
    BOS = 256
    IMAGE = 257
    STATE = 258
-    TASK = 259
+    SUBTASK_QUERY = 259
-    TASK_TOKEN_START = TASK + 1
+    ACTION_START = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
    TASK_TOKEN_END = TASK_TOKEN_START + len(TaskVocabulary.labels) - 1
    SUBTASK_QUERY = TASK_TOKEN_END + 1
    SUBTASK_COMMAND = SUBTASK_QUERY + 1
    SUBTASK_TOKEN_START = SUBTASK_COMMAND + 1
    SUBTASK_TOKEN_END = SUBTASK_TOKEN_START + len(SubtaskVocabulary.labels) - 1
    ACTION_START = SUBTASK_TOKEN_END + 1
    ACTION_END = ACTION_START + 1
    ACTION_TOKEN_START = ACTION_END + 1
    ACTION_BINS = 64
@@ -214,16 +145,14 @@ class MiniTokenizer:
    def __init__(
        self,
        coordinate_tokenizer: CoordinateTokenizer | None = None,
        tasks: TaskVocabulary | None = None,
        subtasks: SubtaskVocabulary | None = None,
    ) -> None:
        self.coordinates = coordinate_tokenizer or CoordinateTokenizer()
-        self.tasks = tasks or TaskVocabulary(first_token_id=self.TASK_TOKEN_START)
+        self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1)
        self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_TOKEN_START)
    @property
    def vocab_size(self) -> int:
-        return max(max(self.tasks.token_ids), max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1
+        return max(max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1
    @property
    def action_token_ids(self) -> range:
@@ -231,37 +160,11 @@ class MiniTokenizer:
        return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)
-    def build_observation(self, history_xy: torch.Tensor) -> torch.Tensor:
+    def build_prefix(self, history_xy: torch.Tensor) -> torch.Tensor:
        coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1)
        specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
-        return torch.cat((specials, coordinate_tokens.cpu()), dim=0)
+        query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
-
+        return torch.cat((specials, coordinate_tokens.cpu(), query), dim=0)
    def build_prefix(
        self,
        history_xy: torch.Tensor,
        task_instruction: str = DEFAULT_TASK_INSTRUCTION,
    ) -> torch.Tensor:
        """Build the high-level ``observation + overall task -> subtask`` prefix."""
        observation = self.build_observation(history_xy)
        task = torch.tensor(
            [self.TASK, self.tasks.encode(task_instruction), self.SUBTASK_QUERY],
            dtype=torch.long,
        )
        return torch.cat((observation, task), dim=0)
    def build_low_level_prefix(
        self,
        history_xy: torch.Tensor,
        subtask: str | int | torch.Tensor,
    ) -> torch.Tensor:
        """Build ``observation + predicted subtask`` for FAST/flow action inference."""
        token = self.subtasks.encode(subtask) if isinstance(subtask, str) else int(subtask)
        if token not in self.subtasks.token_ids:
            raise ValueError(f"low-level command is not a subtask token: {token}")
        command = torch.tensor([self.SUBTASK_COMMAND, token], dtype=torch.long)
        return torch.cat((self.build_observation(history_xy), command), dim=0)
    def encode_subtask(self, reasoning: str) -> torch.Tensor:
        return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long)
@@ -15,8 +15,7 @@ from .runtime import choose_device, load_config, seed_everything, write_json
 def _collate(rows: list[dict]) -> dict:
    tensor_keys = (
-        "image", "prefix_tokens", "low_level_prefix_tokens", "subtask_target",
+        "image", "prefix_tokens", "subtask_target", "actions", "action_tokens"
        "actions", "action_tokens"
    )
    result = {key: torch.stack([row[key] for row in rows]) for key in tensor_keys}
    result["sample_id"] = [row["sample_id"] for row in rows]
@@ -100,8 +99,8 @@ def train(config_path: str) -> Path:
        model_labels = torch.full_like(model_tokens, -100)
        query_index = tokens.shape[1] - 1
        model_labels[:, query_index:] = targets
-        # The action expert is present in the module but receives no loss in
+        # Flow tokens may read observation/state and the generated semantic
-        # discrete pre-training, and its parameters are excluded from optimizer.
+        # subtask, but not teacher-forced discrete action tokens.
        zero_action = torch.zeros_like(actions)
        output = model(
            image, model_tokens, zero_action,
@@ -137,7 +136,6 @@ def train(config_path: str) -> Path:
        batch, iterator = next_batch(iterator)
        image = batch["image"].to(device)
        tokens = batch["prefix_tokens"].to(device)
        low_level_tokens = batch["low_level_prefix_tokens"].to(device)
        actions = batch["actions"].to(device)
        subtask = batch["subtask_target"].to(device)[:, None]
        action_tokens = batch["action_tokens"].to(device)
@@ -145,35 +143,27 @@ def train(config_path: str) -> Path:
        model_tokens = torch.cat((tokens, targets[:, :-1]), dim=1)
        model_labels = torch.full_like(model_tokens, -100)
        model_labels[:, tokens.shape[1] - 1:] = targets
        action_prefix_mask = torch.zeros_like(model_tokens, dtype=torch.bool)
        action_prefix_mask[:, : tokens.shape[1] + 1] = True
        noisy, timestep, velocity = flow.training_batch(actions)
-        # Text/FAST training uses the overall task and predicts the subtask.
+        output = model(
-        # Continuous low-level actions use a second call whose prompt contains
+            image, model_tokens, noisy, timestep,
-        # only the observation and teacher-forced subtask. This preserves the
+            action_prefix_mask=action_prefix_mask,
        # paper's factorization pi(a | observation, predicted_subtask).
        token_output = model(
            image, model_tokens, torch.zeros_like(actions),
            torch.zeros(image.shape[0], device=device),
            generated_token_start=tokens.shape[1],
-            text_labels=model_labels,
+            text_labels=model_labels, target_velocity=velocity,
-        )
+            text_loss_weight=cfg["train"]["semantic_loss_weight"],
-        flow_output = model(
+            flow_loss_weight=cfg["train"]["flow_loss_weight"],
            image, low_level_tokens, noisy, timestep,
            target_velocity=velocity,
        )
        loss = (
            cfg["train"]["semantic_loss_weight"] * token_output.text_loss
            + cfg["train"]["flow_loss_weight"] * flow_output.flow_loss
        )
        optimizer.zero_grad(set_to_none=True)
-        loss.backward()
+        output.loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"])
        optimizer.step()
        if step == 1 or step % cfg["train"]["log_every"] == 0:
            record = {
                "stage": "posttrain", "step": step,
-                "loss": float(loss.detach().cpu()),
+                "loss": float(output.loss.detach().cpu()),
-                "text_loss": float(token_output.text_loss.detach().cpu()),
+                "text_loss": float(output.text_loss.detach().cpu()),
-                "flow_loss": float(flow_output.flow_loss.detach().cpu()),
+                "flow_loss": float(output.flow_loss.detach().cpu()),
            }
            history.append(record)
            print(json.dumps(record))
@@ -8,12 +8,7 @@ import torch
 from PIL import Image
 from pi05mini.data import Pi05OverfitDataset, load_manifest, materialize_subset
-from pi05mini.tokenizer import (
+from pi05mini.tokenizer import CoordinateTokenizer, MiniTokenizer, SubtaskVocabulary
    CoordinateTokenizer,
    MiniTokenizer,
    SubtaskVocabulary,
    TaskVocabulary,
 )
 def _make_source_project(root: Path) -> Path:
@@ -40,7 +35,6 @@ def _make_source_project(root: Path) -> Path:
    row = {
        "clip_id": "clip-a",
        "t0_us": 1234,
        "task_instruction": "Follow the road through the construction zone.",
        "reasoning": "Steer left through the construction lane.",
        "image_path": "data/processed/images/sample.jpg",
        "trajectory_path": "data/processed/trajectories/sample.npz",
@@ -62,28 +56,17 @@ def test_coordinate_tokens_are_literal_bytes_and_round_trip() -> None:
    assert torch.allclose(decoded, coordinates, atol=10.1 / 255.0)
-def test_task_subtask_vocabulary_and_hierarchical_prefix_contract() -> None:
+def test_subtask_vocabulary_and_prefix_contract() -> None:
    tokenizer = MiniTokenizer()
    history = torch.zeros(8, 2)
-    prefix = tokenizer.build_prefix(history, "Follow the road.")
+    prefix = tokenizer.build_prefix(history)
    other_task = tokenizer.build_prefix(history, "Clean the kitchen.")
    subtask = tokenizer.encode_subtask("Steer left around the cones")
    low_level = tokenizer.build_low_level_prefix(history, subtask)
-    assert prefix.shape == (22,)
+    assert prefix.shape == (20,)
    assert prefix[:3].tolist() == [tokenizer.BOS, tokenizer.IMAGE, tokenizer.STATE]
    assert prefix[-3].item() == tokenizer.TASK
    assert tokenizer.tasks.decode(prefix[-2]) == "follow_road"
    assert prefix[-1].item() == tokenizer.SUBTASK_QUERY
-    assert prefix[3:19].min() >= 0 and prefix[3:19].max() <= 255
+    assert prefix[3:-1].min() >= 0 and prefix[3:-1].max() <= 255
-    assert other_task[-2] != prefix[-2]
+    target = tokenizer.encode_subtask("Steer left around the cones")
-    assert low_level.shape == (21,)
+    assert tokenizer.subtasks.decode(target) == "steer_left"
    torch.testing.assert_close(low_level[:19], prefix[:19])
    assert low_level[-2].item() == tokenizer.SUBTASK_COMMAND
    assert low_level[-1].item() == subtask
    assert tokenizer.TASK not in low_level
    assert tokenizer.subtasks.decode(subtask) == "steer_left"
    assert TaskVocabulary().classify("Clean all dishes in the kitchen") == "clean_kitchen"
    assert SubtaskVocabulary().classify("Slow down for a pedestrian") == "slow_down"
@@ -115,9 +98,7 @@ def test_materialize_and_dataset_are_self_contained_and_deterministic(tmp_path:
    augmented = dataset[1]
    repeated = same_dataset[1]
    assert original["image"].shape == (3, 32, 48)
-    assert original["prefix_tokens"].shape == (22,)
+    assert original["prefix_tokens"].shape == (20,)
    assert original["low_level_prefix_tokens"].shape == (21,)
    assert original["task_instruction"].startswith("Follow the road")
    assert original["actions"].shape == (16, 2)
    assert original["subtask_target"].shape == ()
    assert original["subtask_label"] == "steer_left"