diff --git a/README.md b/README.md index b40f44c..27c2d69 100644 --- a/README.md +++ b/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。 ## 推理与效果图 diff --git a/outputs/overfit/evaluation.json b/outputs/overfit/evaluation.json index efb31eb..1c3e036 100644 --- a/outputs/overfit/evaluation.json +++ b/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 } ] \ No newline at end of file diff --git a/outputs/overfit/summary.json b/outputs/overfit/summary.json index 887b413..9527da1 100644 --- a/outputs/overfit/summary.json +++ b/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 } } \ No newline at end of file diff --git a/outputs/overfit/training_history.json b/outputs/overfit/training_history.json index fdb3275..5231953 100644 --- a/outputs/overfit/training_history.json +++ b/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 } ] \ No newline at end of file diff --git a/outputs/overfit/trajectory_comparison.png b/outputs/overfit/trajectory_comparison.png index dde41ea..cd30a7c 100644 Binary files a/outputs/overfit/trajectory_comparison.png and b/outputs/overfit/trajectory_comparison.png differ diff --git a/src/pi05mini/data.py b/src/pi05mini/data.py index 6080d5d..c9373b9 100644 --- a/src/pi05mini/data.py +++ b/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, diff --git a/src/pi05mini/infer.py b/src/pi05mini/infer.py index 6dd9b2f..4b1b155 100644 --- a/src/pi05mini/infer.py +++ b/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, diff --git a/src/pi05mini/tokenizer.py b/src/pi05mini/tokenizer.py index 461b644..17d3773 100644 --- a/src/pi05mini/tokenizer.py +++ b/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) diff --git a/src/pi05mini/train.py b/src/pi05mini/train.py index 2524517..bb598e5 100644 --- a/src/pi05mini/train.py +++ b/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)) diff --git a/tests/test_data.py b/tests/test_data.py index 7b23be7..3fb29f9 100644 --- a/tests/test_data.py +++ b/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"