docs: clarify minimal reproduction overview

.gitignore

@@ -5,5 +5,6 @@ __pycache__/
 *.egg-info/
 .venv/
 data/
+notebook_artifacts/
 outputs/**/*.pt
 .DS_Store

README.md

@@ -5,22 +5,23 @@
 ## 实现内容
 
 ```text
-图像 + 本体状态
-        │
-        ▼
-小型视觉/语言主干 ──→ 高层语义 subtask
-        │                    │
-        └──── semantic prefix┘
-                     │
- noisy action + time │
-          └──────────▼
+图像 + 离散本体状态 + 总体任务指令
+                  │
+                  ▼
+         高层语义 subtask 预测
+                  │ predicted subtask
+                  ▼
+图像 + 离散本体状态 + 低层 subtask command
+                  │
+ 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 的对应关系
 
@@ -32,6 +33,8 @@
 - 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 轨迹，与真值三方比较。
 
@@ -39,7 +42,7 @@
 
 - 小型自定义视觉编码器和 Transformer 取代 PaliGemma 2B；
 - 隐藏宽度、层数、图像尺寸和 action horizon 大幅缩小；
-- 使用固定小词表的 subtask 分类/生成，而非开放式语言理解；
+- 使用固定小词表的 task/subtask 分类与生成，而非开放式语言理解；
 - 使用二维驾驶轨迹，取代多机器人、多相机和高维关节控制；
 - 默认使用 1 个 PAI 驾驶样本的 4 个视觉增强版本，没有机器人或网页联合训练数据；
 - 使用正交 DCT-II、低频优先排列和 64-bin 量化实现 FAST-lite；保留全部频率系数，因此只验证表示与训练接口，不复现官方压缩率；
@@ -67,6 +70,13 @@ 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
@@ -95,7 +105,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 与 flow loss。CPU 可以运行；Apple Silicon 可自动选择 MPS，CUDA 可自动选择 GPU。
+配置默认先进行 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。
 
 ## 推理与效果图
 

outputs/overfit/evaluation.json

@@ -2,33 +2,37 @@
   {
     "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.27604275941848755
+    "diffusion_trajectory_mse": 0.14029410481452942
   },
   {
     "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.20049458742141724
+    "diffusion_trajectory_mse": 0.11049529910087585
   },
   {
     "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.03129641339182854
+    "diffusion_trajectory_mse": 0.04682297259569168
   },
   {
     "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.11321215331554413
+    "diffusion_trajectory_mse": 0.1004234328866005
   }
 ]

outputs/overfit/summary.json

@@ -1,19 +1,19 @@
 {
   "device": "mps",
   "samples": 4,
-  "parameters": 816482,
+  "parameters": 817346,
   "initial": {
     "stage": "pretrain",
     "step": 1,
-    "loss": 78.87449645996094,
-    "text_loss": 78.87449645996094,
+    "loss": 80.38571166992188,
+    "text_loss": 80.38571166992188,
     "flow_loss": null
   },
   "final": {
     "stage": "posttrain",
     "step": 500,
-    "loss": 0.12353827804327011,
-    "text_loss": 3.1691649837739533e-06,
-    "flow_loss": 0.06176755577325821
+    "loss": 0.1929057091474533,
+    "text_loss": 2.784368575703411e-07,
+    "flow_loss": 0.09645271301269531
   }
 }

outputs/overfit/training_history.json

@@ -2,295 +2,295 @@
   {
     "stage": "pretrain",
     "step": 1,
-    "loss": 78.87449645996094,
-    "text_loss": 78.87449645996094,
+    "loss": 80.38571166992188,
+    "text_loss": 80.38571166992188,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 20,
-    "loss": 2.475086212158203,
-    "text_loss": 2.475086212158203,
+    "loss": 2.5035173892974854,
+    "text_loss": 2.5035173892974854,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 40,
-    "loss": 1.1025534868240356,
-    "text_loss": 1.1025534868240356,
+    "loss": 1.3843837976455688,
+    "text_loss": 1.3843837976455688,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 60,
-    "loss": 0.601564884185791,
-    "text_loss": 0.601564884185791,
+    "loss": 0.8159888386726379,
+    "text_loss": 0.8159888386726379,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 80,
-    "loss": 0.25571808218955994,
-    "text_loss": 0.25571808218955994,
+    "loss": 0.19992661476135254,
+    "text_loss": 0.19992661476135254,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 100,
-    "loss": 0.03314117714762688,
-    "text_loss": 0.03314117714762688,
+    "loss": 0.05571526661515236,
+    "text_loss": 0.05571526661515236,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 120,
-    "loss": 0.005310459993779659,
-    "text_loss": 0.005310459993779659,
+    "loss": 0.005689455661922693,
+    "text_loss": 0.005689455661922693,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 140,
-    "loss": 3.0251239877543412e-05,
-    "text_loss": 3.0251239877543412e-05,
+    "loss": 0.00010565178672550246,
+    "text_loss": 0.00010565178672550246,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 160,
-    "loss": 1.1354763955750968e-05,
-    "text_loss": 1.1354763955750968e-05,
+    "loss": 4.098461431567557e-05,
+    "text_loss": 4.098461431567557e-05,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 180,
-    "loss": 6.638075319642667e-06,
-    "text_loss": 6.638075319642667e-06,
+    "loss": 2.716351991693955e-05,
+    "text_loss": 2.716351991693955e-05,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 200,
-    "loss": 5.348983449948719e-06,
-    "text_loss": 5.348983449948719e-06,
+    "loss": 2.2414897102862597e-05,
+    "text_loss": 2.2414897102862597e-05,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 220,
-    "loss": 4.736786195280729e-06,
-    "text_loss": 4.736786195280729e-06,
+    "loss": 1.9143404642818496e-05,
+    "text_loss": 1.9143404642818496e-05,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 240,
-    "loss": 4.316162630857434e-06,
-    "text_loss": 4.316162630857434e-06,
+    "loss": 1.6641104593873024e-05,
+    "text_loss": 1.6641104593873024e-05,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 260,
-    "loss": 3.975574600190157e-06,
-    "text_loss": 3.975574600190157e-06,
+    "loss": 1.4626367374148685e-05,
+    "text_loss": 1.4626367374148685e-05,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 280,
-    "loss": 3.6877775073662633e-06,
-    "text_loss": 3.6877775073662633e-06,
+    "loss": 1.2977556252735667e-05,
+    "text_loss": 1.2977556252735667e-05,
     "flow_loss": null
   },
   {
     "stage": "pretrain",
     "step": 300,
-    "loss": 3.4408508327032905e-06,
-    "text_loss": 3.4408508327032905e-06,
+    "loss": 1.16078899736749e-05,
+    "text_loss": 1.16078899736749e-05,
     "flow_loss": null
   },
   {
     "stage": "posttrain",
     "step": 1,
-    "loss": 8.707193374633789,
-    "text_loss": 3.428078571232618e-06,
-    "flow_loss": 4.353594779968262
+    "loss": 3.8964109420776367,
+    "text_loss": 1.1545746929186862e-05,
+    "flow_loss": 1.94819974899292
   },
   {
     "stage": "posttrain",
     "step": 20,
-    "loss": 3.533726453781128,
-    "text_loss": 0.14151602983474731,
-    "flow_loss": 1.6961052417755127
+    "loss": 4.304055690765381,
+    "text_loss": 2.3414058887283318e-05,
+    "flow_loss": 2.1520161628723145
   },
   {
     "stage": "posttrain",
     "step": 40,
-    "loss": 3.0850143432617188,
-    "text_loss": 0.00014649856893811375,
-    "flow_loss": 1.5424339771270752
+    "loss": 2.925170660018921,
+    "text_loss": 2.142304765584413e-06,
+    "flow_loss": 1.4625842571258545
   },
   {
     "stage": "posttrain",
     "step": 60,
-    "loss": 1.3556954860687256,
-    "text_loss": 0.00011281041224719957,
-    "flow_loss": 0.6777913570404053
+    "loss": 1.7996485233306885,
+    "text_loss": 5.1970469030493405e-06,
+    "flow_loss": 0.8998216390609741
   },
   {
     "stage": "posttrain",
     "step": 80,
-    "loss": 1.4165724515914917,
-    "text_loss": 6.599419066333212e-06,
-    "flow_loss": 0.7082829475402832
+    "loss": 2.0647871494293213,
+    "text_loss": 3.251374073443003e-05,
+    "flow_loss": 1.0323773622512817
   },
   {
     "stage": "posttrain",
     "step": 100,
-    "loss": 0.8609042167663574,
-    "text_loss": 8.11900736152893e-06,
-    "flow_loss": 0.430448055267334
+    "loss": 2.3402059078216553,
+    "text_loss": 2.50748803409806e-06,
+    "flow_loss": 1.1701016426086426
   },
   {
     "stage": "posttrain",
     "step": 120,
-    "loss": 0.5642163753509521,
-    "text_loss": 7.576879397674929e-06,
-    "flow_loss": 0.28210440278053284
+    "loss": 1.4503896236419678,
+    "text_loss": 6.828942105130409e-07,
+    "flow_loss": 0.7251944541931152
   },
   {
     "stage": "posttrain",
     "step": 140,
-    "loss": 0.7110214233398438,
-    "text_loss": 6.249698344618082e-06,
-    "flow_loss": 0.35550758242607117
+    "loss": 0.23238907754421234,
+    "text_loss": 1.3709052382182563e-07,
+    "flow_loss": 0.1161944717168808
   },
   {
     "stage": "posttrain",
     "step": 160,
-    "loss": 0.18725578486919403,
-    "text_loss": 2.876289045161684e-06,
-    "flow_loss": 0.09362645447254181
+    "loss": 0.46604833006858826,
+    "text_loss": 2.1542784622852196e-07,
+    "flow_loss": 0.23302406072616577
   },
   {
     "stage": "posttrain",
     "step": 180,
-    "loss": 0.33923929929733276,
-    "text_loss": 6.748496616637567e-06,
-    "flow_loss": 0.16961628198623657
+    "loss": 0.169603168964386,
+    "text_loss": 2.418236704215815e-07,
+    "flow_loss": 0.08480146527290344
   },
   {
     "stage": "posttrain",
     "step": 200,
-    "loss": 0.3003162443637848,
-    "text_loss": 9.495882295595948e-06,
-    "flow_loss": 0.15015336871147156
+    "loss": 0.7555881142616272,
+    "text_loss": 8.855541011598689e-08,
+    "flow_loss": 0.3777940273284912
   },
   {
     "stage": "posttrain",
     "step": 220,
-    "loss": 0.23716998100280762,
-    "text_loss": 5.660519946104614e-06,
-    "flow_loss": 0.11858215928077698
+    "loss": 0.10211063176393509,
+    "text_loss": 1.3623892414216243e-07,
+    "flow_loss": 0.05105524882674217
   },
   {
     "stage": "posttrain",
     "step": 240,
-    "loss": 0.42717865109443665,
-    "text_loss": 8.400202204938978e-06,
-    "flow_loss": 0.21358512341976166
+    "loss": 0.4649242162704468,
+    "text_loss": 2.2819931189133058e-07,
+    "flow_loss": 0.23246198892593384
   },
   {
     "stage": "posttrain",
     "step": 260,
-    "loss": 0.11702287942171097,
-    "text_loss": 4.122769951209193e-06,
-    "flow_loss": 0.058509379625320435
+    "loss": 0.280495822429657,
+    "text_loss": 1.3198149417803506e-07,
+    "flow_loss": 0.14024785161018372
   },
   {
     "stage": "posttrain",
     "step": 280,
-    "loss": 0.19245445728302002,
-    "text_loss": 5.941245035501197e-06,
-    "flow_loss": 0.09622425585985184
+    "loss": 0.5054064989089966,
+    "text_loss": 9.451582627661992e-08,
+    "flow_loss": 0.2527031898498535
   },
   {
     "stage": "posttrain",
     "step": 300,
-    "loss": 0.18475297093391418,
-    "text_loss": 4.744139005197212e-06,
-    "flow_loss": 0.09237411618232727
+    "loss": 0.19181671738624573,
+    "text_loss": 4.453279984772962e-07,
+    "flow_loss": 0.09590813517570496
   },
   {
     "stage": "posttrain",
     "step": 320,
-    "loss": 0.38481318950653076,
-    "text_loss": 3.1461640901397914e-06,
-    "flow_loss": 0.19240501523017883
+    "loss": 0.20589734613895416,
+    "text_loss": 2.8524956974251836e-07,
+    "flow_loss": 0.10294853150844574
   },
   {
     "stage": "posttrain",
     "step": 340,
-    "loss": 0.2163764089345932,
-    "text_loss": 4.1934181354008615e-06,
-    "flow_loss": 0.10818611085414886
+    "loss": 0.09927964955568314,
+    "text_loss": 1.0132762895409542e-07,
+    "flow_loss": 0.04963977262377739
   },
   {
     "stage": "posttrain",
     "step": 360,
-    "loss": 0.23533868789672852,
-    "text_loss": 3.3019753118423978e-06,
-    "flow_loss": 0.11766768991947174
+    "loss": 0.2213846743106842,
+    "text_loss": 1.2942695093443035e-07,
+    "flow_loss": 0.11069227010011673
   },
   {
     "stage": "posttrain",
     "step": 380,
-    "loss": 0.08397223055362701,
-    "text_loss": 8.292082384286914e-06,
-    "flow_loss": 0.04198196902871132
+    "loss": 0.9335335493087769,
+    "text_loss": 1.3964488232431904e-07,
+    "flow_loss": 0.46676671504974365
   },
   {
     "stage": "posttrain",
     "step": 400,
-    "loss": 0.13249839842319489,
-    "text_loss": 1.0871653103095014e-05,
-    "flow_loss": 0.06624376028776169
+    "loss": 0.12127648293972015,
+    "text_loss": 7.578287153364727e-08,
+    "flow_loss": 0.06063820421695709
   },
   {
     "stage": "posttrain",
     "step": 420,
-    "loss": 0.4479130506515503,
-    "text_loss": 1.0128642315976322e-05,
-    "flow_loss": 0.22395145893096924
+    "loss": 0.1503632515668869,
+    "text_loss": 4.0871746875836834e-08,
+    "flow_loss": 0.07518160343170166
   },
   {
     "stage": "posttrain",
     "step": 440,
-    "loss": 0.22416232526302338,
-    "text_loss": 2.6114355478057405e-06,
-    "flow_loss": 0.11207985877990723
+    "loss": 0.15399529039859772,
+    "text_loss": 2.0435869885204738e-08,
+    "flow_loss": 0.07699763774871826
   },
   {
     "stage": "posttrain",
     "step": 460,
-    "loss": 0.06346385926008224,
-    "text_loss": 3.1035808660817565e-06,
-    "flow_loss": 0.031730376183986664
+    "loss": 0.18908466398715973,
+    "text_loss": 4.172323286866231e-08,
+    "flow_loss": 0.09454230964183807
   },
   {
     "stage": "posttrain",
     "step": 480,
-    "loss": 0.385486900806427,
-    "text_loss": 2.184032382501755e-06,
-    "flow_loss": 0.19274236261844635
+    "loss": 0.2329331636428833,
+    "text_loss": 2.8099325533048614e-08,
+    "flow_loss": 0.11646656692028046
   },
   {
     "stage": "posttrain",
     "step": 500,
-    "loss": 0.12353827804327011,
-    "text_loss": 3.1691649837739533e-06,
-    "flow_loss": 0.06176755577325821
+    "loss": 0.1929057091474533,
+    "text_loss": 2.784368575703411e-07,
+    "flow_loss": 0.09645271301269531
   }
 ]

pyproject.toml

@@ -22,6 +22,10 @@ dependencies = [
 dev = [
   "pytest>=8.0",
 ]
+notebook = [
+  "jupyterlab>=4.2",
+  "nbclient>=0.10",
+]
 
 [project.scripts]
 pi05mini-train = "pi05mini.train:main"

src/pi05mini/data.py

@@ -19,6 +19,7 @@ from .tokenizer import MiniTokenizer
 class SampleRecord:
     clip_id: str
     t0_us: int
+    task_instruction: str
     reasoning: str
     image_path: Path
     trajectory_path: Path
@@ -111,6 +112,9 @@ 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"]),
@@ -156,6 +160,7 @@ 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)),
@@ -303,7 +308,13 @@ 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),
+            "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,
             "subtask_target": target,
             "subtask_label": self.tokenizer.subtasks.decode(target),
             "actions": actions,

src/pi05mini/infer.py

@@ -40,13 +40,20 @@ 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 _ in range(expected_tokens):
+        for generation_step in range(expected_tokens):
             semantic = model(
                 image, current_tokens, zero_action,
                 torch.zeros(1, device=device),
                 generated_token_start=tokens.shape[1]
             )
-            next_token = int(semantic.text_logits[0, -1].argmax().cpu())
+            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())
             generated.append(next_token)
             current_tokens = torch.cat(
                 (current_tokens, torch.tensor([[next_token]], device=device)), dim=1
@@ -62,13 +69,12 @@ 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 = torch.cat(
-            (tokens, torch.tensor([[predicted_token]], device=device)), dim=1
-        )
+        conditioned_tokens = dataset.tokenizer.build_low_level_prefix(
+            sample["history_xy"], predicted_token
+        )[None].to(device)
         def velocity(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
             return model(
-                image, conditioned_tokens, x, t,
-                generated_token_start=tokens.shape[1]
+                image, conditioned_tokens, x, t
             ).vector_field
         normalized = flow.sample(
             velocity, (1, dataset.action_horizon, 2), device,
@@ -88,6 +94,7 @@ 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,

src/pi05mini/tokenizer.py

@@ -52,6 +52,61 @@ 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."""
 
@@ -114,7 +169,13 @@ class SubtaskVocabulary:
 
 
 class MiniTokenizer:
-    """Compose image/state, semantic, and FAST-lite action tokens.
+    """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.
 
     Actions are z-score-normalized by the dataset before reaching this
     tokenizer.  A complete ``[horizon, 2]`` chunk is transformed along time
@@ -130,11 +191,19 @@ class MiniTokenizer:
     high-frequency coefficients.
     """
 
+    DEFAULT_TASK_INSTRUCTION = "Follow the visible road."
+
     BOS = 256
     IMAGE = 257
     STATE = 258
-    SUBTASK_QUERY = 259
-    ACTION_START = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels)
+    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
     ACTION_END = ACTION_START + 1
     ACTION_TOKEN_START = ACTION_END + 1
     ACTION_BINS = 64
@@ -145,14 +214,16 @@ 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.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1)
+        self.tasks = tasks or TaskVocabulary(first_token_id=self.TASK_TOKEN_START)
+        self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_TOKEN_START)
 
     @property
     def vocab_size(self) -> int:
-        return max(max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1
+        return max(max(self.tasks.token_ids), max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1
 
     @property
     def action_token_ids(self) -> range:
@@ -160,11 +231,37 @@ class MiniTokenizer:
 
         return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)
 
-    def build_prefix(self, history_xy: torch.Tensor) -> torch.Tensor:
+    def build_observation(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)
-        query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long)
-        return torch.cat((specials, coordinate_tokens.cpu(), query), dim=0)
+        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)
 
     def encode_subtask(self, reasoning: str) -> torch.Tensor:
         return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long)

src/pi05mini/train.py

@@ -15,7 +15,8 @@ from .runtime import choose_device, load_config, seed_everything, write_json
 
 def _collate(rows: list[dict]) -> dict:
     tensor_keys = (
-        "image", "prefix_tokens", "subtask_target", "actions", "action_tokens"
+        "image", "prefix_tokens", "low_level_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]
@@ -99,8 +100,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
-        # Flow tokens may read observation/state and the generated semantic
-        # subtask, but not teacher-forced discrete action tokens.
+        # The action expert is present in the module but receives no loss in
+        # discrete pre-training, and its parameters are excluded from optimizer.
         zero_action = torch.zeros_like(actions)
         output = model(
             image, model_tokens, zero_action,
@@ -136,6 +137,7 @@ 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)
@@ -143,27 +145,35 @@ 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)
-        output = model(
-            image, model_tokens, noisy, timestep,
-            action_prefix_mask=action_prefix_mask,
+        # 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),
             generated_token_start=tokens.shape[1],
-            text_labels=model_labels, target_velocity=velocity,
-            text_loss_weight=cfg["train"]["semantic_loss_weight"],
-            flow_loss_weight=cfg["train"]["flow_loss_weight"],
+            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
         )
         optimizer.zero_grad(set_to_none=True)
-        output.loss.backward()
+        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(output.loss.detach().cpu()),
-                "text_loss": float(output.text_loss.detach().cpu()),
-                "flow_loss": float(output.flow_loss.detach().cpu()),
+                "loss": float(loss.detach().cpu()),
+                "text_loss": float(token_output.text_loss.detach().cpu()),
+                "flow_loss": float(flow_output.flow_loss.detach().cpu()),
             }
             history.append(record)
             print(json.dumps(record))

tests/test_data.py

@@ -8,7 +8,12 @@ import torch
 from PIL import Image
 
 from pi05mini.data import Pi05OverfitDataset, load_manifest, materialize_subset
-from pi05mini.tokenizer import CoordinateTokenizer, MiniTokenizer, SubtaskVocabulary
+from pi05mini.tokenizer import (
+    CoordinateTokenizer,
+    MiniTokenizer,
+    SubtaskVocabulary,
+    TaskVocabulary,
+)
 
 
 def _make_source_project(root: Path) -> Path:
@@ -35,6 +40,7 @@ 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",
@@ -56,17 +62,28 @@ def test_coordinate_tokens_are_literal_bytes_and_round_trip() -> None:
     assert torch.allclose(decoded, coordinates, atol=10.1 / 255.0)
 
 
-def test_subtask_vocabulary_and_prefix_contract() -> None:
+def test_task_subtask_vocabulary_and_hierarchical_prefix_contract() -> None:
     tokenizer = MiniTokenizer()
     history = torch.zeros(8, 2)
-    prefix = tokenizer.build_prefix(history)
+    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)
 
-    assert prefix.shape == (20,)
+    assert prefix.shape == (22,)
     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:-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 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 SubtaskVocabulary().classify("Slow down for a pedestrian") == "slow_down"
 
 
@@ -98,7 +115,9 @@ 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 == (20,)
+    assert original["prefix_tokens"].shape == (22,)
+    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"