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 预测
-                  │ predicted subtask
-                  ▼
-图像 + 离散本体状态 + 低层 subtask command
-                  │
- noisy action + time
-          └───────▼
+图像 + 本体状态
+        │
+        ▼
+小型视觉/语言主干 ──→ 高层语义 subtask
+        │                    │
+        └──── semantic prefix┘
+                     │
+ 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,
-    "text_loss": 80.38571166992188,
+    "loss": 78.87449645996094,
+    "text_loss": 78.87449645996094,
    "flow_loss": null
  },
  "final": {
    "stage": "posttrain",
    "step": 500,
-    "loss": 0.1929057091474533,
-    "text_loss": 2.784368575703411e-07,
-    "flow_loss": 0.09645271301269531
+    "loss": 0.12353827804327011,
+    "text_loss": 3.1691649837739533e-06,
+    "flow_loss": 0.06176755577325821
  }
 }
@@ -2,295 +2,295 @@
  {
    "stage": "pretrain",
    "step": 1,
-    "loss": 80.38571166992188,
-    "text_loss": 80.38571166992188,
+    "loss": 78.87449645996094,
+    "text_loss": 78.87449645996094,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 20,
-    "loss": 2.5035173892974854,
-    "text_loss": 2.5035173892974854,
+    "loss": 2.475086212158203,
+    "text_loss": 2.475086212158203,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 40,
-    "loss": 1.3843837976455688,
-    "text_loss": 1.3843837976455688,
+    "loss": 1.1025534868240356,
+    "text_loss": 1.1025534868240356,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 60,
-    "loss": 0.8159888386726379,
-    "text_loss": 0.8159888386726379,
+    "loss": 0.601564884185791,
+    "text_loss": 0.601564884185791,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 80,
-    "loss": 0.19992661476135254,
-    "text_loss": 0.19992661476135254,
+    "loss": 0.25571808218955994,
+    "text_loss": 0.25571808218955994,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 100,
-    "loss": 0.05571526661515236,
-    "text_loss": 0.05571526661515236,
+    "loss": 0.03314117714762688,
+    "text_loss": 0.03314117714762688,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 120,
-    "loss": 0.005689455661922693,
-    "text_loss": 0.005689455661922693,
+    "loss": 0.005310459993779659,
+    "text_loss": 0.005310459993779659,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 140,
-    "loss": 0.00010565178672550246,
-    "text_loss": 0.00010565178672550246,
+    "loss": 3.0251239877543412e-05,
+    "text_loss": 3.0251239877543412e-05,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 160,
-    "loss": 4.098461431567557e-05,
-    "text_loss": 4.098461431567557e-05,
+    "loss": 1.1354763955750968e-05,
+    "text_loss": 1.1354763955750968e-05,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 180,
-    "loss": 2.716351991693955e-05,
-    "text_loss": 2.716351991693955e-05,
+    "loss": 6.638075319642667e-06,
+    "text_loss": 6.638075319642667e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 200,
-    "loss": 2.2414897102862597e-05,
-    "text_loss": 2.2414897102862597e-05,
+    "loss": 5.348983449948719e-06,
+    "text_loss": 5.348983449948719e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 220,
-    "loss": 1.9143404642818496e-05,
-    "text_loss": 1.9143404642818496e-05,
+    "loss": 4.736786195280729e-06,
+    "text_loss": 4.736786195280729e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 240,
-    "loss": 1.6641104593873024e-05,
-    "text_loss": 1.6641104593873024e-05,
+    "loss": 4.316162630857434e-06,
+    "text_loss": 4.316162630857434e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 260,
-    "loss": 1.4626367374148685e-05,
-    "text_loss": 1.4626367374148685e-05,
+    "loss": 3.975574600190157e-06,
+    "text_loss": 3.975574600190157e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 280,
-    "loss": 1.2977556252735667e-05,
-    "text_loss": 1.2977556252735667e-05,
+    "loss": 3.6877775073662633e-06,
+    "text_loss": 3.6877775073662633e-06,
    "flow_loss": null
  },
  {
    "stage": "pretrain",
    "step": 300,
-    "loss": 1.16078899736749e-05,
-    "text_loss": 1.16078899736749e-05,
+    "loss": 3.4408508327032905e-06,
+    "text_loss": 3.4408508327032905e-06,
    "flow_loss": null
  },
  {
    "stage": "posttrain",
    "step": 1,
-    "loss": 3.8964109420776367,
-    "text_loss": 1.1545746929186862e-05,
-    "flow_loss": 1.94819974899292
+    "loss": 8.707193374633789,
+    "text_loss": 3.428078571232618e-06,
+    "flow_loss": 4.353594779968262
  },
  {
    "stage": "posttrain",
    "step": 20,
-    "loss": 4.304055690765381,
-    "text_loss": 2.3414058887283318e-05,
-    "flow_loss": 2.1520161628723145
+    "loss": 3.533726453781128,
+    "text_loss": 0.14151602983474731,
+    "flow_loss": 1.6961052417755127
  },
  {
    "stage": "posttrain",
    "step": 40,
-    "loss": 2.925170660018921,
-    "text_loss": 2.142304765584413e-06,
-    "flow_loss": 1.4625842571258545
+    "loss": 3.0850143432617188,
+    "text_loss": 0.00014649856893811375,
+    "flow_loss": 1.5424339771270752
  },
  {
    "stage": "posttrain",
    "step": 60,
-    "loss": 1.7996485233306885,
-    "text_loss": 5.1970469030493405e-06,
-    "flow_loss": 0.8998216390609741
+    "loss": 1.3556954860687256,
+    "text_loss": 0.00011281041224719957,
+    "flow_loss": 0.6777913570404053
  },
  {
    "stage": "posttrain",
    "step": 80,
-    "loss": 2.0647871494293213,
-    "text_loss": 3.251374073443003e-05,
-    "flow_loss": 1.0323773622512817
+    "loss": 1.4165724515914917,
+    "text_loss": 6.599419066333212e-06,
+    "flow_loss": 0.7082829475402832
  },
  {
    "stage": "posttrain",
    "step": 100,
-    "loss": 2.3402059078216553,
-    "text_loss": 2.50748803409806e-06,
-    "flow_loss": 1.1701016426086426
+    "loss": 0.8609042167663574,
+    "text_loss": 8.11900736152893e-06,
+    "flow_loss": 0.430448055267334
  },
  {
    "stage": "posttrain",
    "step": 120,
-    "loss": 1.4503896236419678,
-    "text_loss": 6.828942105130409e-07,
-    "flow_loss": 0.7251944541931152
+    "loss": 0.5642163753509521,
+    "text_loss": 7.576879397674929e-06,
+    "flow_loss": 0.28210440278053284
  },
  {
    "stage": "posttrain",
    "step": 140,
-    "loss": 0.23238907754421234,
-    "text_loss": 1.3709052382182563e-07,
-    "flow_loss": 0.1161944717168808
+    "loss": 0.7110214233398438,
+    "text_loss": 6.249698344618082e-06,
+    "flow_loss": 0.35550758242607117
  },
  {
    "stage": "posttrain",
    "step": 160,
-    "loss": 0.46604833006858826,
-    "text_loss": 2.1542784622852196e-07,
-    "flow_loss": 0.23302406072616577
+    "loss": 0.18725578486919403,
+    "text_loss": 2.876289045161684e-06,
+    "flow_loss": 0.09362645447254181
  },
  {
    "stage": "posttrain",
    "step": 180,
-    "loss": 0.169603168964386,
-    "text_loss": 2.418236704215815e-07,
-    "flow_loss": 0.08480146527290344
+    "loss": 0.33923929929733276,
+    "text_loss": 6.748496616637567e-06,
+    "flow_loss": 0.16961628198623657
  },
  {
    "stage": "posttrain",
    "step": 200,
-    "loss": 0.7555881142616272,
-    "text_loss": 8.855541011598689e-08,
-    "flow_loss": 0.3777940273284912
+    "loss": 0.3003162443637848,
+    "text_loss": 9.495882295595948e-06,
+    "flow_loss": 0.15015336871147156
  },
  {
    "stage": "posttrain",
    "step": 220,
-    "loss": 0.10211063176393509,
-    "text_loss": 1.3623892414216243e-07,
-    "flow_loss": 0.05105524882674217
+    "loss": 0.23716998100280762,
+    "text_loss": 5.660519946104614e-06,
+    "flow_loss": 0.11858215928077698
  },
  {
    "stage": "posttrain",
    "step": 240,
-    "loss": 0.4649242162704468,
-    "text_loss": 2.2819931189133058e-07,
-    "flow_loss": 0.23246198892593384
+    "loss": 0.42717865109443665,
+    "text_loss": 8.400202204938978e-06,
+    "flow_loss": 0.21358512341976166
  },
  {
    "stage": "posttrain",
    "step": 260,
-    "loss": 0.280495822429657,
-    "text_loss": 1.3198149417803506e-07,
-    "flow_loss": 0.14024785161018372
+    "loss": 0.11702287942171097,
+    "text_loss": 4.122769951209193e-06,
+    "flow_loss": 0.058509379625320435
  },
  {
    "stage": "posttrain",
    "step": 280,
-    "loss": 0.5054064989089966,
-    "text_loss": 9.451582627661992e-08,
-    "flow_loss": 0.2527031898498535
+    "loss": 0.19245445728302002,
+    "text_loss": 5.941245035501197e-06,
+    "flow_loss": 0.09622425585985184
  },
  {
    "stage": "posttrain",
    "step": 300,
-    "loss": 0.19181671738624573,
-    "text_loss": 4.453279984772962e-07,
-    "flow_loss": 0.09590813517570496
+    "loss": 0.18475297093391418,
+    "text_loss": 4.744139005197212e-06,
+    "flow_loss": 0.09237411618232727
  },
  {
    "stage": "posttrain",
    "step": 320,
-    "loss": 0.20589734613895416,
-    "text_loss": 2.8524956974251836e-07,
-    "flow_loss": 0.10294853150844574
+    "loss": 0.38481318950653076,
+    "text_loss": 3.1461640901397914e-06,
+    "flow_loss": 0.19240501523017883
  },
  {
    "stage": "posttrain",
    "step": 340,
-    "loss": 0.09927964955568314,
-    "text_loss": 1.0132762895409542e-07,
-    "flow_loss": 0.04963977262377739
+    "loss": 0.2163764089345932,
+    "text_loss": 4.1934181354008615e-06,
+    "flow_loss": 0.10818611085414886
  },
  {
    "stage": "posttrain",
    "step": 360,
-    "loss": 0.2213846743106842,
-    "text_loss": 1.2942695093443035e-07,
-    "flow_loss": 0.11069227010011673
+    "loss": 0.23533868789672852,
+    "text_loss": 3.3019753118423978e-06,
+    "flow_loss": 0.11766768991947174
  },
  {
    "stage": "posttrain",
    "step": 380,
-    "loss": 0.9335335493087769,
-    "text_loss": 1.3964488232431904e-07,
-    "flow_loss": 0.46676671504974365
+    "loss": 0.08397223055362701,
+    "text_loss": 8.292082384286914e-06,
+    "flow_loss": 0.04198196902871132
  },
  {
    "stage": "posttrain",
    "step": 400,
-    "loss": 0.12127648293972015,
-    "text_loss": 7.578287153364727e-08,
-    "flow_loss": 0.06063820421695709
+    "loss": 0.13249839842319489,
+    "text_loss": 1.0871653103095014e-05,
+    "flow_loss": 0.06624376028776169
  },
  {
    "stage": "posttrain",
    "step": 420,
-    "loss": 0.1503632515668869,
-    "text_loss": 4.0871746875836834e-08,
-    "flow_loss": 0.07518160343170166
+    "loss": 0.4479130506515503,
+    "text_loss": 1.0128642315976322e-05,
+    "flow_loss": 0.22395145893096924
  },
  {
    "stage": "posttrain",
    "step": 440,
-    "loss": 0.15399529039859772,
-    "text_loss": 2.0435869885204738e-08,
-    "flow_loss": 0.07699763774871826
+    "loss": 0.22416232526302338,
+    "text_loss": 2.6114355478057405e-06,
+    "flow_loss": 0.11207985877990723
  },
  {
    "stage": "posttrain",
    "step": 460,
-    "loss": 0.18908466398715973,
-    "text_loss": 4.172323286866231e-08,
-    "flow_loss": 0.09454230964183807
+    "loss": 0.06346385926008224,
+    "text_loss": 3.1035808660817565e-06,
+    "flow_loss": 0.031730376183986664
  },
  {
    "stage": "posttrain",
    "step": 480,
-    "loss": 0.2329331636428833,
-    "text_loss": 2.8099325533048614e-08,
-    "flow_loss": 0.11646656692028046
+    "loss": 0.385486900806427,
+    "text_loss": 2.184032382501755e-06,
+    "flow_loss": 0.19274236261844635
  },
  {
    "stage": "posttrain",
    "step": 500,
-    "loss": 0.1929057091474533,
-    "text_loss": 2.784368575703411e-07,
-    "flow_loss": 0.09645271301269531
+    "loss": 0.12353827804327011,
+    "text_loss": 3.1691649837739533e-06,
+    "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(
-                history_xy, record.task_instruction
-            ),
-            "low_level_prefix_tokens": self.tokenizer.build_low_level_prefix(
-                history_xy, target
-            ),
-            "task_instruction": record.task_instruction,
+            "prefix_tokens": self.tokenizer.build_prefix(history_xy),
            "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:
-                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())
+            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(
-            sample["history_xy"], predicted_token
-        )[None].to(device)
+        conditioned_tokens = torch.cat(
+            (tokens, torch.tensor([[predicted_token]], device=device)), dim=1
+        )
        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.
-
-    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.
+    """Compose image/state, semantic, and FAST-lite action tokens.

    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
-    TASK_TOKEN_START = TASK + 1
-    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
+    SUBTASK_QUERY = 259
+    ACTION_START = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
    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_TOKEN_START)
+        self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1)

    @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)
-
-    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)
+        query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
+        return torch.cat((specials, coordinate_tokens.cpu(), query), 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",
-        "actions", "action_tokens"
+        "image", "prefix_tokens", "subtask_target", "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
-        # discrete pre-training, and its parameters are excluded from optimizer.
+        # Flow tokens may read observation/state and the generated semantic
+        # 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.
-        # Continuous low-level actions use a second call whose prompt contains
-        # only the observation and teacher-forced subtask. This preserves the
-        # 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),
+        output = model(
+            image, model_tokens, noisy, timestep,
+            action_prefix_mask=action_prefix_mask,
            generated_token_start=tokens.shape[1],
-            text_labels=model_labels,
-        )
-        flow_output = model(
-            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
+            text_labels=model_labels, target_velocity=velocity,
+            text_loss_weight=cfg["train"]["semantic_loss_weight"],
+            flow_loss_weight=cfg["train"]["flow_loss_weight"],
        )
        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()),
-                "text_loss": float(token_output.text_loss.detach().cpu()),
-                "flow_loss": float(flow_output.flow_loss.detach().cpu()),
+                "loss": float(output.loss.detach().cpu()),
+                "text_loss": float(output.text_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 (
-    CoordinateTokenizer,
-    MiniTokenizer,
-    SubtaskVocabulary,
-    TaskVocabulary,
-)
+from pi05mini.tokenizer import CoordinateTokenizer, MiniTokenizer, SubtaskVocabulary


 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.")
-    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)
+    prefix = tokenizer.build_prefix(history)

-    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 other_task[-2] != prefix[-2]
-    assert low_level.shape == (21,)
-    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 prefix[3:-1].min() >= 0 and prefix[3:-1].max() <= 255
+    target = tokenizer.encode_subtask("Steer left around the cones")
+    assert tokenizer.subtasks.decode(target) == "steer_left"
    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["low_level_prefix_tokens"].shape == (21,)
-    assert original["task_instruction"].startswith("Follow the road")
+    assert original["prefix_tokens"].shape == (20,)
    assert original["actions"].shape == (16, 2)
    assert original["subtask_target"].shape == ()
    assert original["subtask_label"] == "steer_left"