add task-conditioned pi0.5 hierarchy

README.md

@@ -7,22 +7,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 的对应关系
 
@@ -34,6 +35,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 轨迹，与真值三方比较。
 
@@ -41,7 +44,7 @@
 
 - 小型自定义视觉编码器和 Transformer 取代 PaliGemma 2B；
 - 隐藏宽度、层数、图像尺寸和 action horizon 大幅缩小；
-- 使用固定小词表的 subtask 分类/生成，而非开放式语言理解；
+- 使用固定小词表的 task/subtask 分类与生成，而非开放式语言理解；
 - 使用二维驾驶轨迹，取代多机器人、多相机和高维关节控制；
 - 默认使用 1 个 PAI 驾驶样本的 4 个视觉增强版本，没有机器人或网页联合训练数据；
 - 使用正交 DCT-II、低频优先排列和 64-bin 量化实现 FAST-lite；保留全部频率系数，因此只验证表示与训练接口，不复现官方压缩率；
@@ -97,7 +100,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。
 
 ## 推理与效果图