Add 0.5B Alpamayo validation

.gitignore

@@ -0,0 +1,22 @@
+.venv/
+.cache/
+.ruff_cache/
+models/
+data/raw/
+data/processed/
+outputs/
+references/
+*.egg-info/
+*.pyc
+__pycache__/
+.pytest_cache/
+.DS_Store
+.env
+.env.*
+wandb/
+*.safetensors
+*.pt
+*.pth
+*.parquet
+*.npz
+*.log

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,341 @@
+# alpamayo-0.5B
+
+`alpamayo-0.5B` 是一个单机、单样本的 Alpamayo 接口验证工程。它把一个带
+Qwen2.5-0.5B 文本骨干的 VLM 接到缩小版 diffusion expert，验证以下链路能够实际
+训练、保存、重新加载和推理：
+
+1. 图像、自然语言和 ego-relative 历史坐标进入 VLM；
+2. VLM 生成驾驶文字，也可以在未截断验证分支中生成结构化 CoT、meta-action 和
+   离散轨迹 token；
+3. VLM 的 24 层 past-KV 作为 expert 的前缀；
+4. flow-matching diffusion 预测连续 action，再积分回未来 XY 位置。
+
+这是针对连接方式和输出契约的最小过拟合实验，不是 NVIDIA Alpamayo-R1 的完整复现，
+也不是可部署的驾驶模型。
+
+## 架构
+
+```mermaid
+flowchart LR
+    A["前视图像"] --> V["LLaVA VLM"]
+    B["驾驶问题与历史 XY"] --> V
+    Q["CLIP vision tower"] --> V
+    L["Qwen2.5-0.5B text backbone"] --> V
+    V --> T["驾驶文字 / structured tokens"]
+    V --> K["24-layer past-KV"]
+    N["noisy normalized XY-delta"] --> E["105.6M Qwen2 expert"]
+    S["diffusion timestep"] --> E
+    K --> E
+    E --> F["flow-matching vector field"]
+    F --> D["10-step denoising"]
+    D --> X["16 x 2 continuous action"]
+    X --> P["ego-relative future XY"]
+```
+
+连接契约不是按 hidden size 复制 VLM。expert 采用自己的 `hidden_size=512`，但保持与
+VLM KV cache 相同的 24 层、2 个 KV heads 和 `head_dim=64`，因此可以直接消费缓存。
+expert 没有词嵌入层；noisy action 和 timestep 经 Fourier/MLP 投影成 token embeddings，
+最后由线性层输出二维 vector field。
+
+| 组件 | 本地配置 |
+| --- | --- |
+| VLM checkpoint | `thkim0305/qwen2.5_0.5B_vl`，revision `c318c4b464753d86ddffe065ba6df255ae354f41` |
+| VLM 类型 | `LlavaForConditionalGeneration`；CLIP vision tower + Qwen2.5-0.5B text backbone |
+| VLM 参数量 | 799,262,336（“0.5B”指文本骨干；加入视觉塔后完整 VLM 约 0.80B） |
+| VLM text cache | 24 层，2 KV heads，`head_dim=64` |
+| Expert | 105,575,170 参数；24 层，hidden 512，intermediate 2432，8 query heads，2 KV heads，non-causal |
+| Diffusion 输出 | 16 个二维 normalized XY-delta action；10 个 flow-matching inference steps |
+| 运行精度 | 当前产物为 MPS + float16；代码也可选择可用的 CUDA 或 CPU |
+
+主配置位于 [`configs/local_overfit.yaml`](configs/local_overfit.yaml)。
+
+## 单样本 PAI 数据
+
+实验只使用 `nvidia/PhysicalAI-Autonomous-Vehicles` 的一个 `train` 事件。PAI 是 gated
+dataset，使用前需要在 Hugging Face 接受其许可。
+
+| 字段 | 值 |
+| --- | --- |
+| Clip ID | `314f59c2-e420-4fcd-9df2-a51d572e0c20` |
+| `t0_us` | `1503899` |
+| Camera | `camera_front_wide_120fov` |
+| Event cluster | `WORK_ZONES_TEMP_TRAFFIC_CONTROL` |
+| Reasoning source | `reasoning/ood_reasoning.parquet` 中的 `coc` 字段 |
+| Reasoning | “Steer left following the temporary lane delineated by traffic cones through the construction zone.” |
+| History | 8 个点，最后一个点为当前 ego-local 原点 |
+| Future source | egomotion 的 64 个未来点，时间步 0.1 s |
+| Expert target | 从 64 点以索引 `[0,4,8,12,16,21,25,29,33,37,42,46,50,54,58,63]` 取 16 点 |
+
+[`scripts/materialize_pai_subset.py`](scripts/materialize_pai_subset.py) 将世界坐标平移、旋转到
+`t0` ego-relative frame，并把单样本 manifest、图像和轨迹写入本地 `data/processed/`。
+原始/处理后的 PAI 数据都被 `.gitignore` 排除，不应提交到仓库。
+
+## 训练阶段
+
+### Stage 1：自然语言驾驶输出
+
+[`scripts/train_stage1_vlm.py`](scripts/train_stage1_vlm.py) 冻结 vision tower、multimodal
+projector 和 LM head，只训练最后 8 个 Qwen2 language layers 及 final norm。当前 checkpoint
+训练 80 steps，共 119,299,968 个可训练参数：loss 从 `3.9338069` 降至
+`0.0000216427`，greedy 输出与数据集 reasoning 逐字一致。
+
+Stage 1 输出到本地 `outputs/stage1_vlm/`，供后续两个分支使用：
+
+- plain-text KV 分支训练现有 diffusion expert；
+- structured VLM 分支继续训练轨迹 logits。
+
+### VLM structured 未截断轨迹-token 分支
+
+[`scripts/train_vlm_trajectory_logits.py`](scripts/train_vlm_trajectory_logits.py) 从 Stage 1
+checkpoint 继续训练，并添加：
+
+- 128 个连续 token：`<i0>` … `<i127>`；
+- `<|cot_start|>`、`<|cot_end|>`；
+- `<|meta_action_start|>`、`<|meta_action_end|>`；
+- `<|traj_future_start|>`、`<|traj_future_end|>`。
+
+当前 Qwen LM head 已预留 151,936 行，而原 tokenizer 长度为 151,666；扩展后的 151,800
+个 token 不需要增大权重矩阵。每个 16 x 2 normalized XY-delta action 被均匀量化到
+`[-4, 4]` 的 128 bins，形成 32 个监督 token。
+
+当前结构化 checkpoint 训练 160 steps，loss 从 `14.17349` 降至 `0.0000347825`，最终
+训练 accuracy 为 1.0。未截断 greedy generation 生成 32/32 个正确轨迹 token，其完整
+目标/输出格式为：
+
+```text
+<|cot_start|>Steer left following the temporary lane delineated by traffic cones through the construction zone.<|cot_end|>
+<|meta_action_start|>STEER_LEFT<|meta_action_end|>
+<|traj_future_start|><i9><i37><i71><i38><i70><i41><i68><i46><i66><i49><i84><i63><i63><i63><i65><i67><i66><i69><i65><i67><i82><i80><i61><i77><i60><i79><i59><i78><i56><i76><i72><i86><|traj_future_end|><|im_end|>
+```
+
+另有 source-style handoff 验证：generation 到 `<|traj_future_start|>` 后屏蔽全部 128 个
+trajectory logits，并把 exact generation cache 交给 diffusion。这个模式验证了 KV 接口，
+但不会把 VLM 离散轨迹 token 作为 diffusion 分支的可见输出。
+
+### Stage 2：KV-conditioned diffusion expert
+
+[`scripts/train_stage2_expert.py`](scripts/train_stage2_expert.py) 使用 Stage 1 plain-text prefix
+的 past-KV，训练 105.6M expert 预测 flow-matching vector field。目标 action 是按单样本统计量
+标准化的相邻 waypoint XY delta，预测 action 经反标准化和累积求和恢复未来轨迹。
+
+当前产物记录为 1,000 total steps：早期 250-step 权重缺少 optimizer state，之后以 legacy
+warm start 运行了剩余 750 steps。因此它是 1,000-step 参数训练记录，但不是严格连续、
+optimizer 完整恢复的 1,000-step run。新 checkpoint 已保存 expert、optimizer、step 和 RNG，
+后续中断可用 `--resume` 精确恢复。
+
+普通 reload evaluation 仍使用训练 expert 时的旧 plain-text prefix（sequence length 770）。
+structured exact-cache handoff 的 prefix 长度为 777；接口验证成功，但旧 expert 没有在这个
+新 prefix 分布上联合训练，所以该分支的轨迹质量只作诊断。
+
+## 目录
+
+```text
+alpamayo-0.5B/
+├── configs/                     # 单样本实验配置
+├── scripts/                     # 下载、数据准备、训练、评估和比较入口
+├── src/alpamayo_lite/           # VLM、action space、expert、diffusion 实现
+├── tests/                       # 单元测试和产物一致性测试
+├── results/complete_comparison/ # 仓库中保留的轻量比较 JSON、CSV 和三方图
+├── references/alpamayo-upstream/# 可选的本地只读上游 checkout；gitignored
+├── models/                      # 本地模型；gitignored
+├── data/                        # gated/raw/processed 数据；gitignored
+├── outputs/                     # checkpoint 和评估产物；gitignored
+└── .venv/                       # Python 3.12 虚拟环境；gitignored
+```
+
+## 安装
+
+需要 Python 3.12。以下命令把环境和所有运行产物留在仓库目录内：
+
+```bash
+git clone <repository-url> alpamayo-0.5B
+cd alpamayo-0.5B
+python3.12 -m venv .venv
+source .venv/bin/activate
+python -m pip install --upgrade pip
+python -m pip install -e '.[dev]'
+```
+
+登录 Hugging Face 时使用交互式认证，不要把 access token 写进配置、脚本、README、日志或
+shell history：
+
+```bash
+hf auth login
+```
+
+## 运行
+
+所有命令均从仓库根目录执行。
+
+下载 0.5B VLM，并记录 checkpoint provenance：
+
+```bash
+python scripts/download_model.py --config configs/local_overfit.yaml
+```
+
+下载并物化一个 PAI 样本，然后验证 VLM KV 与 expert 的形状契约：
+
+```bash
+python scripts/materialize_pai_subset.py --config configs/local_overfit.yaml
+python scripts/verify_connection.py --config configs/local_overfit.yaml
+```
+
+训练 Stage 1 自然语言 VLM：
+
+```bash
+python scripts/train_stage1_vlm.py \
+  --config configs/local_overfit.yaml \
+  --steps 80
+```
+
+训练并验证 VLM 的 structured/untruncated trajectory logits：
+
+```bash
+python scripts/train_vlm_trajectory_logits.py \
+  --config configs/local_overfit.yaml \
+  --output-dir outputs/stage1_vlm_trajectory \
+  --steps 160 \
+  --trajectory-vocab-size 128 \
+  --action-min -4 \
+  --action-max 4 \
+  --max-new-tokens 96
+```
+
+从头训练 Stage 2 到 1,000 total steps：
+
+```bash
+python scripts/train_stage2_expert.py \
+  --config configs/local_overfit.yaml \
+  --steps 1000
+```
+
+若 `outputs/stage2_expert/checkpoints/latest/` 中已有完整 checkpoint，则精确续训到总计
+1,000 steps：
+
+```bash
+python scripts/train_stage2_expert.py \
+  --config configs/local_overfit.yaml \
+  --steps 1000 \
+  --resume
+```
+
+重新加载 VLM/expert，固定从 seed 12 开始做 16 个 diffusion rollouts：
+
+```bash
+python scripts/evaluate_pipeline.py \
+  --config configs/local_overfit.yaml \
+  --stage2-dir outputs/stage2_expert \
+  --num-rollouts 16 \
+  --base-seed 12 \
+  --output outputs/reload_evaluation.json
+```
+
+验证 source-style structured VLM exact cache 能直接进入旧 expert：
+
+```bash
+python scripts/verify_structured_vlm_handoff.py \
+  --config configs/local_overfit.yaml \
+  --vlm-dir outputs/stage1_vlm_trajectory \
+  --stage2-dir outputs/stage2_expert \
+  --output-dir outputs/structured_vlm_handoff
+```
+
+汇总 GT、VLM direct trajectory 和 diffusion seed 12：
+
+```bash
+python scripts/build_complete_comparison.py \
+  --diffusion-seed 12 \
+  --output-dir outputs/complete_comparison
+```
+
+## 当前结果
+
+两条预测使用逐字节相同的 GT action 和 GT XY。下表由保存的 NPZ 重新计算；action MSE 的
+单位是 normalized-action squared，轨迹指标单位是米。
+
+| 输出 | 解码协议 | Action MSE | ADE | FDE | Max error |
+| --- | --- | ---: | ---: | ---: | ---: |
+| Ground truth | GT 对自身 | 0 | 0 m | 0 m | 0 m |
+| VLM trajectory tokens | greedy、未截断、128-bin | 0.0003439452 | 0.026104664 m | 0.045444489 m | 0.054956669 m |
+| Diffusion | flow matching、seed 12、10 steps | 0.7453363983 | 1.125946622 m | 0.258275141 m | 2.925430649 m |
+
+VLM 的 32 个生成 token 与量化 GT 全部匹配，因此其非零轨迹误差完全来自 128-bin
+反量化。Diffusion 是随机采样；seed 12 只代表一个 rollout。16 个 seed 的平均 ADE 为
+`1.530302633 m`，不应把 seed 12 当作总体均值。
+
+![GT、VLM 与 diffusion 三方比较](results/complete_comparison/trajectory_comparison.png)
+
+### 文字对比
+
+| 字段 | Ground truth | 模型输出 | Exact match |
+| --- | --- | --- | ---: |
+| Reasoning / CoT | Steer left following the temporary lane delineated by traffic cones through the construction zone. | Steer left following the temporary lane delineated by traffic cones through the construction zone. | 是 |
+| Meta-action | `STEER_LEFT`（由 reasoning 规则派生） | `STEER_LEFT` | 是 |
+| Diffusion conditioning text | 同上方 reasoning | 同上方 reasoning | 是 |
+
+## 本地产物
+
+仓库保留以下不含权重或 gated 样本的轻量结果，可直接审计完整三方对比：
+
+- [GT/VLM/diffusion 完整比较 JSON](results/complete_comparison/comparison.json)
+- [16 个 waypoint 的完整 CSV](results/complete_comparison/waypoints.csv)
+- [GT/VLM/diffusion 三方比较图](results/complete_comparison/trajectory_comparison.png)
+
+`models/`、`data/` 和 `outputs/` 默认不进入 Git。运行完成后，以下链接仅在本地工作区存在；
+不要为了让它们在 GitHub 页面可用而提交 gated data 或 checkpoint 权重。
+
+- [Stage 1 文字训练指标](outputs/stage1_vlm/stage1_metrics.json)
+- [Structured VLM 训练、logits、cache 与生成指标](outputs/stage1_vlm_trajectory/trajectory_logits_metrics.json)
+- [Structured VLM trajectory vocabulary](outputs/stage1_vlm_trajectory/trajectory_vocabulary.json)
+- [Structured VLM 轨迹图](outputs/stage1_vlm_trajectory/vlm_discrete_trajectory.png)
+- [Stage 2 训练指标](outputs/stage2_expert/stage2_metrics.json)
+- [Stage 2 expert 配置](outputs/stage2_expert/expert_config.json)
+- [Stage 2 seed-12 reload 轨迹图](outputs/stage2_expert/reload_trajectory.png)
+- [16-seed reload 评估](outputs/reload_evaluation.json)
+- [Structured exact-cache handoff 指标](outputs/structured_vlm_handoff/handoff_metrics.json)
+- [Structured exact-cache handoff 轨迹图](outputs/structured_vlm_handoff/handoff_trajectory.png)
+- [本地 GT/VLM/diffusion 完整比较 JSON](outputs/complete_comparison/comparison.json)
+- [本地 16 个 waypoint CSV](outputs/complete_comparison/waypoints.csv)
+- [本地 GT/VLM/diffusion 比较图](outputs/complete_comparison/trajectory_comparison.png)
+
+没有列出或链接任何 checkpoint 权重文件。
+
+## 测试状态
+
+发布或提交结果前运行：
+
+```bash
+pytest -q
+ruff check src scripts tests
+```
+
+| 检查 | 最近结果 |
+| --- | --- |
+| `pytest -q` | `11 passed` |
+| `ruff check src scripts tests` | `All checks passed` |
+
+## 已知限制与安全边界
+
+- 这是一个 PAI `train` 样本的显式过拟合实验，没有验证跨 clip、跨场景或真实驾驶泛化。
+- 本地 VLM 轨迹词表只有 128 bins，action 是 normalized per-waypoint XY delta；它不是上游
+  使用的 768-token vocabulary 或约 4,000-bin delta-tokenizer 配置，也不是上游
+  unicycle acceleration/curvature action space。
+- action mean/scale 由同一个单样本目标轨迹拟合；VLM direct trajectory 的 32 个监督 token
+  也来自该目标，因此其近零误差只证明记忆和输出链路。
+- `STEER_LEFT` 是 [`_infer_meta_action`](scripts/train_vlm_trajectory_logits.py) 从 reasoning
+  文字规则派生的标签，不是 PAI 中独立提供的 meta-action ground truth。
+- 当前 expert 用旧 plain-text `"Continuous action:"` prefix 训练，并未与 structured VLM
+  prefix 或 exact generation cache 联合训练。structured handoff 证明 24-layer KV 兼容，
+  不代表该 prefix 下的轨迹质量已经收敛。
+- source-style handoff 会在 `<|traj_future_start|>` 屏蔽 VLM 离散轨迹 logits；未截断 VLM
+  轨迹输出是一个单独的能力验证分支。
+- 当前 1,000-step expert 包含一次缺少 optimizer state 的 legacy warm start；需要全新、严格
+  连续的训练曲线时应删除本地旧产物并从头运行。
+- 模型输出和本项目都不能用于控制真实车辆或做安全决策。
+- 不要提交或发布 Hugging Face access token、认证缓存、gated PAI 原始/处理数据、下载的模型
+  权重、训练 checkpoint 或可能包含私有路径/凭据的日志。公开前应再次检查 `git status` 和
+  Git 历史；仅提交有权再分发的代码与文档。
+
+本项目代码按 [Apache-2.0](LICENSE) 发布。上游归属、修改说明及不随仓库分发的模型/数据
+声明见 [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md)。本地工作区可另行检出
+`references/alpamayo-upstream/` 进行接口研究，但该目录不会随本仓库分发。

THIRD_PARTY_NOTICES.md

@@ -0,0 +1,26 @@
+# Third-party notices
+
+This repository is an unofficial, Alpamayo-inspired local validation project. It is not
+affiliated with, endorsed by, or a release of NVIDIA Corporation.
+
+## NVlabs/alpamayo
+
+Parts of the local expert connection pattern and flow-matching implementation were derived
+from or informed by [`NVlabs/alpamayo`](https://github.com/NVlabs/alpamayo), inspected at
+commit `4b96c546ce3416e7b1e4d86e7649662852c3c1ac`.
+
+- Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+- Licensed under the Apache License, Version 2.0.
+- Locally modified files retain their SPDX copyright and license headers.
+
+The upstream source checkout is not redistributed in this repository.
+
+## Models and datasets
+
+The code references the Hugging Face checkpoint `thkim0305/qwen2.5_0.5B_vl` and the gated
+dataset `nvidia/PhysicalAI-Autonomous-Vehicles`. Neither downloaded model weights nor raw or
+processed dataset samples are redistributed here. Their respective repository terms and
+licenses apply independently from this project's Apache-2.0 code license.
+
+Users must obtain their own authorized access, review the applicable terms, and avoid using
+this research smoke test for real-vehicle control or safety decisions.

configs/local_overfit.yaml

@@ -0,0 +1,56 @@
+seed: 11
+
+paths:
+  model_dir: models/qwen2.5_0.5B_vl
+  pai_raw_dir: data/raw/pai
+  manifest: data/processed/overfit_manifest.jsonl
+  stage1_dir: outputs/stage1_vlm
+  stage2_dir: outputs/stage2_expert
+
+vlm:
+  repo_id: thkim0305/qwen2.5_0.5B_vl
+  dtype: float16
+  attn_implementation: eager
+  max_reasoning_tokens: 48
+  num_additional_image_tokens: 1
+
+data:
+  repo_id: nvidia/PhysicalAI-Autonomous-Vehicles
+  num_clips: 1
+  camera: camera_front_wide_120fov
+  num_frames: 1
+  history_steps: 8
+  source_future_steps: 64
+  expert_waypoints: 16
+  time_step: 0.1
+
+stage1:
+  steps: 80
+  learning_rate: 0.0001
+  weight_decay: 0.0
+  adam_epsilon: 0.0001
+  freeze_vision_tower: true
+  freeze_multimodal_projector: true
+  train_last_n_language_layers: 8
+  train_final_norm: true
+  train_lm_head: false
+  gradient_checkpointing: true
+  log_every: 5
+
+stage2:
+  steps: 1000
+  learning_rate: 0.0001
+  weight_decay: 0.0
+  adam_epsilon: 0.0001
+  expert_hidden_size: 512
+  expert_intermediate_size: 2432
+  expert_num_layers: 24
+  expert_num_attention_heads: 8
+  expert_num_key_value_heads: 2
+  expert_head_dim: 64
+  expert_non_causal_attention: true
+  action_projection_hidden_size: 256
+  action_fourier_features: 20
+  inference_steps: 10
+  log_every: 50
+  checkpoint_every: 50

pyproject.toml

@@ -0,0 +1,47 @@
+[build-system]
+requires = ["setuptools>=75", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "alpamayo-qwen05b-validation"
+version = "0.1.0"
+description = "Two-stage local Alpamayo validation with a 0.5B Qwen VLM and a KV-conditioned flow-matching expert."
+readme = "README.md"
+license = "Apache-2.0"
+requires-python = ">=3.12,<3.13"
+dependencies = [
+  "accelerate==1.12.0",
+  "av>=16.0.1",
+  "einops>=0.8.1",
+  "huggingface-hub==0.36.0",
+  "matplotlib>=3.9",
+  "numpy>=2.1,<3",
+  "pandas>=2.2",
+  "physical-ai-av>=0.2.0",
+  "pillow>=11",
+  "psutil>=6",
+  "pyarrow>=18",
+  "pyyaml>=6",
+  "safetensors>=0.5",
+  "scipy>=1.14",
+  "torch==2.8.0",
+  "torchvision==0.23.0",
+  "tqdm>=4.67",
+  "transformers==4.57.1",
+]
+
+[project.optional-dependencies]
+dev = [
+  "pytest>=8.3",
+  "ruff>=0.9",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 100
+extend-exclude = ["references"]

results/complete_comparison/waypoints.csv

@@ -0,0 +1,17 @@
+waypoint_index,source_future_index,time_s,gt_action_x,gt_action_y,vlm_action_x,vlm_action_y,diffusion_action_x,diffusion_action_y,vlm_action_error_l2,diffusion_action_error_l2,gt_x_m,gt_y_m,vlm_x_m,vlm_y_m,diffusion_x_m,diffusion_y_m,vlm_position_error_m,diffusion_position_error_m
+1,0,0.1,-3.4397106170654297,-1.6662895679473877,-3.4330708980560303,-1.6692912578582764,-3.84375,-0.443359375,0.00728670095756034,1.2879464584308167,0.9489696025848389,-0.00039592807297594845,0.9537601470947266,-0.0006286799907684326,0.657459020614624,0.09442976862192154,0.004796195383473365,0.30654580759535116
+2,4,0.5,0.46647655963897705,-1.6230756044387817,0.4724407196044922,-1.6062991619110107,-0.68017578125,-1.453125,0.01780506186402779,1.1591785017050988,4.71621561050415,0.002558943582698703,4.725309371948242,0.0036270394921302795,3.5974063873291016,0.11056257039308548,0.009156272498882664,1.1240101695561981
+3,8,0.9,0.3931064307689667,-1.3917449712753296,0.40944862365722656,-1.4173228740692139,-0.204345703125,-2.953125,0.03035286444028776,1.6717824758008162,8.430525779724121,0.023451128974556923,8.451410293579102,0.02253592759370804,6.880660533905029,0.010385759174823761,0.02090455721909409,1.5499203153987877
+4,12,1.3,0.26794740557670593,-1.1304054260253906,0.2834644317626953,-1.1023621559143066,1.0224609375,-0.82275390625,0.032050009363179546,0.8148252128373736,12.054534912109375,0.06460751593112946,12.086614608764648,0.06586679071187973,11.049043655395508,0.07539737224578857,0.03210440328783258,1.0055491476438747
+5,16,1.7000000000000002,0.1750280112028122,-0.931678295135498,0.15748023986816406,-0.9133858680725098,-0.126953125,-0.8857421875,0.025348316840923476,0.3054549600301843,15.611503601074219,0.12117315828800201,15.630922317504883,0.12385082244873047,14.38813591003418,0.1355248987674713,0.019602459875540927,1.2234518707065785
+6,21,2.2,1.3138797283172607,-0.017029158771038055,1.2913384437561035,-0.03149604797363281,2.5078125,-0.9423828125,0.02678433110733247,1.5105478634480827,19.990142822265625,0.24866046011447906,19.993297576904297,0.25021636486053467,19.628185272216797,0.19126051664352417,0.003517572516526143,0.36648058822783336
+7,25,2.6,-0.0019493416184559464,-0.05992039293050766,-0.03149604797363281,-0.03149604797363281,0.2435302734375,-0.045867919921875,0.04099940539405987,0.24588150277253198,23.419424057006836,0.37282198667526245,23.401260375976562,0.37658190727233887,23.234577178955078,0.31651169061660767,0.018548754984252572,0.19323358343645852
+8,29,3.0,0.07472366094589233,0.23706550896167755,0.09448814392089844,0.2204723358154297,1.005859375,0.17138671875,0.025806359338943326,0.9334492066902851,26.904024124145508,0.5200117230415344,26.90011978149414,0.5224850177764893,27.39098358154297,0.4586087167263031,0.004621804667577402,0.49081548705534744
+9,33,3.4000000000000004,0.13519389927387238,0.3722502589225769,0.15748023986816406,0.34645652770996094,0.12469482421875,0.52001953125,0.034088085117141044,0.14814178485892907,30.432252883911133,0.6776836514472961,30.444427490234375,0.6781569123268127,30.911636352539062,0.6277386546134949,0.01218380133562631,0.4819782284527732
+10,37,3.8000000000000003,0.06780000776052475,0.21779732406139374,0.09448814392089844,0.2204723358154297,0.432373046875,-0.6748046875,0.026821862343970004,0.9641845528178778,33.91185760498047,0.8233793377876282,33.94328689575195,0.8240600228309631,34.65427780151367,0.7042222023010254,0.03143666093157371,0.7519216522735415
+11,42,4.3,1.1512386798858643,1.035483717918396,1.1653542518615723,1.039370059967041,1.95703125,0.78125,0.014640800071050068,0.84494760155749,38.173152923583984,1.0324782133102417,38.21476745605469,1.033460259437561,39.496944427490234,0.8936078548431396,0.04162611833154021,1.331055566937441
+12,46,4.7,-0.15684714913368225,0.8499115109443665,-0.15748023986816406,0.8503937721252441,2.060546875,1.314453125,0.0007958515720084112,2.2655319837630143,41.49067687988281,1.2271878719329834,41.53183364868164,1.2282073497772217,44.41429901123047,1.124337911605835,0.04116939340134956,2.925430648852423
+13,50,5.1000000000000005,-0.2395608276128769,0.9665660858154297,-0.2204723358154297,0.9763779640197754,-0.9697265625,2.46875,0.021462606388743272,1.6702390584698135,44.7485237121582,1.4309428930282593,44.80345153808594,1.4327231645584106,47.14533615112305,1.4445719718933105,0.054956668638742615,2.3968511884068464
+14,54,5.5,-0.2766667306423187,0.9057282209396362,-0.2834646701812744,0.9133858680725098,-0.52783203125,1.48046875,0.010239704174774942,0.6272245881452143,47.97959899902344,1.6299805641174316,48.029624938964844,1.6323546171188354,50.19519805908203,1.688174843788147,0.050082240312057325,2.2163631852922747
+15,58,5.9,-0.44114238023757935,0.7915734648704529,-0.4724409580230713,0.7874016761779785,-1.5810546875,0.258056640625,0.03157538269426274,1.2585866160107093,51.09200668334961,1.8201667070388794,51.11945343017578,1.8222172260284424,52.48516845703125,1.8369922637939453,0.027523236355242352,1.393263373166715
+16,63,6.4,0.5104826092720032,1.4437676668167114,0.535433292388916,1.4173226356506348,-1.0625,1.236328125,0.03635761627740968,1.5866018569828055,54.8910026550293,2.0609238147735596,54.93644714355469,2.0609238147735596,55.14927673339844,2.061664581298828,0.045444488525390625,0.2582751406784516

scripts/build_complete_comparison.py

@@ -0,0 +1,722 @@
+#!/usr/bin/env python3
+"""Build one artifact set comparing ground truth, VLM, and diffusion outputs."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import textwrap
+from pathlib import Path
+from typing import Any
+
+import matplotlib
+import numpy as np
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+METRIC_NAMES = ("action_mse", "ade_m", "fde_m", "max_error_m")
+
+
+def _resolve(path: Path) -> Path:
+    return path.expanduser().resolve() if path.is_absolute() else (PROJECT_ROOT / path).resolve()
+
+
+def _require_file(path: Path) -> Path:
+    resolved = _resolve(path)
+    if not resolved.is_file():
+        raise FileNotFoundError(f"Required input does not exist: {resolved}")
+    return resolved
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    value = json.loads(path.read_text(encoding="utf-8"))
+    if not isinstance(value, dict):
+        raise ValueError(f"Expected a JSON object in {path}")
+    return value
+
+
+def _read_manifest_row(path: Path, index: int) -> dict[str, Any]:
+    rows = [line for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
+    if not 0 <= index < len(rows):
+        raise IndexError(f"Manifest index {index} is outside 0..{len(rows) - 1}: {path}")
+    value = json.loads(rows[index])
+    if not isinstance(value, dict):
+        raise ValueError(f"Expected a JSON object at manifest row {index}: {path}")
+    return value
+
+
+def _load_npz(path: Path) -> dict[str, np.ndarray]:
+    with np.load(path, allow_pickle=False) as archive:
+        return {key: np.asarray(archive[key]) for key in archive.files}
+
+
+def _require_array(
+    arrays: dict[str, np.ndarray], key: str, source: Path, *, allow_batch: bool = True
+) -> np.ndarray:
+    if key not in arrays:
+        raise KeyError(f"Missing NPZ key {key!r} in {source}")
+    value = arrays[key]
+    if allow_batch and value.ndim == 3 and value.shape[0] == 1:
+        value = value[0]
+    if value.ndim != 2 or value.shape[1] != 2:
+        raise ValueError(f"{source}:{key} must have shape (N, 2) or (1, N, 2), got {value.shape}")
+    if not np.issubdtype(value.dtype, np.number) or not np.isfinite(value).all():
+        raise ValueError(f"{source}:{key} must contain only finite numeric values")
+    return value.astype(np.float64, copy=False)
+
+
+def _require_xy_prefix(arrays: dict[str, np.ndarray], key: str, source: Path) -> np.ndarray:
+    if key not in arrays:
+        raise KeyError(f"Missing NPZ key {key!r} in {source}")
+    value = arrays[key]
+    if value.ndim != 2 or value.shape[1] < 2:
+        raise ValueError(f"{source}:{key} must have shape (N, >=2), got {value.shape}")
+    if not np.issubdtype(value.dtype, np.number) or not np.isfinite(value).all():
+        raise ValueError(f"{source}:{key} must contain only finite numeric values")
+    return value[:, :2].astype(np.float64, copy=False)
+
+
+def _require_token_values(arrays: dict[str, np.ndarray], key: str, source: Path) -> np.ndarray:
+    if key not in arrays:
+        raise KeyError(f"Missing NPZ key {key!r} in {source}")
+    value = np.asarray(arrays[key]).reshape(-1)
+    if not np.issubdtype(value.dtype, np.integer):
+        raise ValueError(f"{source}:{key} must contain integer token values")
+    return value.astype(np.int64, copy=False)
+
+
+def _assert_same_truth(
+    name: str,
+    vlm_truth: np.ndarray,
+    diffusion_truth: np.ndarray,
+    vlm_path: Path,
+    diffusion_path: Path,
+) -> float:
+    if vlm_truth.shape != diffusion_truth.shape:
+        raise ValueError(
+            f"Ground-truth {name} shape mismatch: {vlm_path} has {vlm_truth.shape}, "
+            f"{diffusion_path} has {diffusion_truth.shape}"
+        )
+    max_abs_difference = float(np.max(np.abs(vlm_truth - diffusion_truth), initial=0.0))
+    if not np.array_equal(vlm_truth, diffusion_truth):
+        raise ValueError(
+            f"Ground-truth {name} mismatch between {vlm_path} and {diffusion_path}; "
+            f"max_abs_difference={max_abs_difference:.9g}"
+        )
+    return max_abs_difference
+
+
+def _extract_component(text: str, component: str) -> str:
+    start = f"<|{component}_start|>"
+    end = f"<|{component}_end|>"
+    if text.count(start) != 1 or text.count(end) != 1:
+        raise ValueError(f"Structured VLM text must contain exactly one {start} and one {end}")
+    start_index = text.find(start)
+    content_start = start_index + len(start)
+    end_index = text.find(end, content_start)
+    if start_index < 0 or end_index < content_start:
+        raise ValueError(f"Structured VLM markers are out of order for {component}")
+    value = text[content_start:end_index].strip()
+    if not value:
+        raise ValueError(f"Structured VLM component {component!r} is empty")
+    return value
+
+
+def _prediction_metrics(
+    predicted_action: np.ndarray,
+    predicted_xy: np.ndarray,
+    target_action: np.ndarray,
+    target_xy: np.ndarray,
+) -> dict[str, float]:
+    position_errors = np.linalg.norm(predicted_xy - target_xy, axis=-1)
+    return {
+        "action_mse": float(np.mean(np.square(predicted_action - target_action))),
+        "ade_m": float(np.mean(position_errors)),
+        "fde_m": float(position_errors[-1]),
+        "max_error_m": float(np.max(position_errors)),
+    }
+
+
+def _reported_metrics(metrics: dict[str, Any], section: str | None = None) -> dict[str, float]:
+    source: Any = metrics if section is None else metrics.get(section)
+    if not isinstance(source, dict):
+        raise KeyError(f"Missing metrics section {section!r}")
+    result: dict[str, float] = {}
+    for name in METRIC_NAMES:
+        if name not in source:
+            raise KeyError(f"Missing reported metric {name!r} in section {section!r}")
+        result[name] = float(source[name])
+    return result
+
+
+def _metric_validation(
+    computed: dict[str, float], reported: dict[str, float], tolerance: float = 1e-5
+) -> dict[str, Any]:
+    absolute_differences = {name: abs(computed[name] - reported[name]) for name in METRIC_NAMES}
+    consistent = all(difference <= tolerance for difference in absolute_differences.values())
+    if not consistent:
+        raise ValueError(
+            "Computed metrics do not match the metrics JSON: "
+            f"absolute_differences={absolute_differences}, tolerance={tolerance}"
+        )
+    return {
+        "consistent": True,
+        "absolute_tolerance": tolerance,
+        "absolute_differences": absolute_differences,
+    }
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as stream:
+        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _display_path(path: Path) -> str:
+    try:
+        return str(path.relative_to(PROJECT_ROOT))
+    except ValueError:
+        return str(path)
+
+
+def _source(path: Path, *, keys: list[str] | None = None) -> dict[str, Any]:
+    result: dict[str, Any] = {"path": _display_path(path), "sha256": _sha256(path)}
+    if keys is not None:
+        result["keys"] = keys
+    return result
+
+
+def _decode_action(
+    action: np.ndarray,
+    history_xy: np.ndarray,
+    mean: np.ndarray,
+    scale: np.ndarray,
+) -> np.ndarray:
+    deltas = action * scale + mean
+    return history_xy[-1] + np.cumsum(deltas, axis=0)
+
+
+def _validate_decoded_xy(
+    name: str,
+    decoded_xy: np.ndarray,
+    stored_xy: np.ndarray,
+    tolerance: float = 1e-5,
+) -> float:
+    max_abs_difference = float(np.max(np.abs(decoded_xy - stored_xy), initial=0.0))
+    if not np.allclose(decoded_xy, stored_xy, rtol=0.0, atol=tolerance):
+        raise ValueError(
+            f"Decoded {name} action does not reproduce stored XY; "
+            f"max_abs_difference={max_abs_difference:.9g}, tolerance={tolerance}"
+        )
+    return max_abs_difference
+
+
+def _write_waypoints(path: Path, waypoints: list[dict[str, Any]]) -> None:
+    fieldnames = [
+        "waypoint_index",
+        "source_future_index",
+        "time_s",
+        "gt_action_x",
+        "gt_action_y",
+        "vlm_action_x",
+        "vlm_action_y",
+        "diffusion_action_x",
+        "diffusion_action_y",
+        "vlm_action_error_l2",
+        "diffusion_action_error_l2",
+        "gt_x_m",
+        "gt_y_m",
+        "vlm_x_m",
+        "vlm_y_m",
+        "diffusion_x_m",
+        "diffusion_y_m",
+        "vlm_position_error_m",
+        "diffusion_position_error_m",
+    ]
+    with path.open("w", encoding="utf-8", newline="") as stream:
+        writer = csv.DictWriter(stream, fieldnames=fieldnames, lineterminator="\n")
+        writer.writeheader()
+        for waypoint in waypoints:
+            writer.writerow(
+                {
+                    "waypoint_index": waypoint["index"],
+                    "source_future_index": waypoint["source_future_index"],
+                    "time_s": waypoint["time_s"],
+                    "gt_action_x": waypoint["action"]["ground_truth"][0],
+                    "gt_action_y": waypoint["action"]["ground_truth"][1],
+                    "vlm_action_x": waypoint["action"]["vlm"][0],
+                    "vlm_action_y": waypoint["action"]["vlm"][1],
+                    "diffusion_action_x": waypoint["action"]["diffusion"][0],
+                    "diffusion_action_y": waypoint["action"]["diffusion"][1],
+                    "vlm_action_error_l2": waypoint["action_error_l2"]["vlm"],
+                    "diffusion_action_error_l2": waypoint["action_error_l2"]["diffusion"],
+                    "gt_x_m": waypoint["xy_m"]["ground_truth"][0],
+                    "gt_y_m": waypoint["xy_m"]["ground_truth"][1],
+                    "vlm_x_m": waypoint["xy_m"]["vlm"][0],
+                    "vlm_y_m": waypoint["xy_m"]["vlm"][1],
+                    "diffusion_x_m": waypoint["xy_m"]["diffusion"][0],
+                    "diffusion_y_m": waypoint["xy_m"]["diffusion"][1],
+                    "vlm_position_error_m": waypoint["position_error_m"]["vlm"],
+                    "diffusion_position_error_m": waypoint["position_error_m"]["diffusion"],
+                }
+            )
+
+
+def _save_plot(
+    path: Path,
+    history_xy: np.ndarray,
+    target_xy: np.ndarray,
+    vlm_xy: np.ndarray,
+    diffusion_xy: np.ndarray,
+    target_action: np.ndarray,
+    vlm_action: np.ndarray,
+    diffusion_action: np.ndarray,
+    waypoint_times_s: np.ndarray,
+    vlm_errors: np.ndarray,
+    diffusion_errors: np.ndarray,
+    ground_truth_reasoning: str,
+    vlm_reasoning: str,
+) -> None:
+    figure, axes = plt.subplots(2, 2, figsize=(14, 10))
+    trajectory_axis, action_x_axis, action_y_axis, error_axis = axes.reshape(-1)
+    trajectory_axis.plot(history_xy[:, 0], history_xy[:, 1], "o-", color="#555555", label="history")
+    trajectory_axis.plot(
+        target_xy[:, 0], target_xy[:, 1], "o-", color="#d62728", label="ground truth"
+    )
+    trajectory_axis.plot(
+        vlm_xy[:, 0],
+        vlm_xy[:, 1],
+        "o--",
+        color="#1f77b4",
+        label=f"VLM discrete (ADE {vlm_errors.mean():.3f} m)",
+    )
+    trajectory_axis.plot(
+        diffusion_xy[:, 0],
+        diffusion_xy[:, 1],
+        "o-.",
+        color="#ff7f0e",
+        label=f"diffusion (ADE {diffusion_errors.mean():.3f} m)",
+    )
+    trajectory_axis.scatter([0.0], [0.0], marker="x", color="black", label="current ego")
+    trajectory_axis.set_xlabel("ego-relative x (m)")
+    trajectory_axis.set_ylabel("ego-relative y (m)")
+    trajectory_axis.set_title("XY trajectory (axes independently scaled)")
+    trajectory_axis.grid(alpha=0.25)
+    trajectory_axis.legend()
+
+    action_series = (
+        ("ground truth", target_action, "#d62728", "o-"),
+        ("VLM discrete", vlm_action, "#1f77b4", "o--"),
+        ("diffusion", diffusion_action, "#ff7f0e", "o-."),
+    )
+    for label, action, color, style in action_series:
+        action_x_axis.plot(waypoint_times_s, action[:, 0], style, color=color, label=label)
+        action_y_axis.plot(waypoint_times_s, action[:, 1], style, color=color, label=label)
+    for axis, component in ((action_x_axis, "x"), (action_y_axis, "y")):
+        axis.set_xlabel("future time (s)")
+        axis.set_ylabel(f"normalized action {component}")
+        axis.set_title(f"Normalized action {component}")
+        axis.grid(alpha=0.25)
+    action_x_axis.legend()
+
+    error_axis.plot(waypoint_times_s, vlm_errors, "o--", color="#1f77b4", label="VLM discrete")
+    error_axis.plot(waypoint_times_s, diffusion_errors, "o-.", color="#ff7f0e", label="diffusion")
+    error_axis.set_xlabel("future time (s)")
+    error_axis.set_ylabel("position error (m)")
+    error_axis.set_title("Per-waypoint position error")
+    error_axis.grid(alpha=0.25)
+    error_axis.legend()
+
+    match = "yes" if vlm_reasoning == ground_truth_reasoning else "no"
+    text_block = (
+        f"Ground truth text: {ground_truth_reasoning}\n"
+        f"VLM text: {vlm_reasoning}\n"
+        f"Exact match: {match}"
+    )
+    figure.suptitle("Ground truth vs VLM logits vs diffusion")
+    figure.text(0.02, 0.015, textwrap.fill(text_block, width=150), va="bottom")
+    figure.tight_layout(rect=(0.0, 0.10, 1.0, 0.97))
+    figure.savefig(path, dpi=180)
+    plt.close(figure)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--vlm-npz",
+        type=Path,
+        default=Path("outputs/stage1_vlm_trajectory/trajectory_logits.npz"),
+    )
+    parser.add_argument(
+        "--vlm-metrics",
+        type=Path,
+        default=Path("outputs/stage1_vlm_trajectory/trajectory_logits_metrics.json"),
+    )
+    parser.add_argument(
+        "--diffusion-npz",
+        type=Path,
+        default=Path("outputs/stage2_expert/reload_prediction.npz"),
+    )
+    parser.add_argument(
+        "--diffusion-metrics",
+        type=Path,
+        default=Path("outputs/stage2_expert/stage2_metrics.json"),
+    )
+    parser.add_argument(
+        "--expert-config",
+        type=Path,
+        default=Path("outputs/stage2_expert/expert_config.json"),
+    )
+    parser.add_argument(
+        "--manifest", type=Path, default=Path("data/processed/overfit_manifest.jsonl")
+    )
+    parser.add_argument("--manifest-index", type=int, default=0)
+    parser.add_argument("--diffusion-seed", type=int, default=12)
+    parser.add_argument("--time-step-s", type=float, default=0.1)
+    parser.add_argument("--output-dir", type=Path, default=Path("outputs/complete_comparison"))
+    args = parser.parse_args()
+    if args.time_step_s <= 0:
+        raise ValueError("--time-step-s must be positive")
+
+    vlm_npz_path = _require_file(args.vlm_npz)
+    vlm_metrics_path = _require_file(args.vlm_metrics)
+    diffusion_npz_path = _require_file(args.diffusion_npz)
+    diffusion_metrics_path = _require_file(args.diffusion_metrics)
+    expert_config_path = _require_file(args.expert_config)
+    manifest_path = _require_file(args.manifest)
+    output_dir = _resolve(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    vlm_arrays = _load_npz(vlm_npz_path)
+    diffusion_arrays = _load_npz(diffusion_npz_path)
+    vlm_metrics = _read_json(vlm_metrics_path)
+    diffusion_metrics = _read_json(diffusion_metrics_path)
+    expert_config = _read_json(expert_config_path)
+    manifest = _read_manifest_row(manifest_path, args.manifest_index)
+    trajectory_value = manifest.get("trajectory_path")
+    if not isinstance(trajectory_value, str):
+        raise KeyError("Manifest row lacks string field 'trajectory_path'")
+    trajectory_path = _require_file(Path(trajectory_value))
+    trajectory_arrays = _load_npz(trajectory_path)
+
+    vlm_target_action = _require_array(vlm_arrays, "target_action", vlm_npz_path)
+    vlm_target_xy = _require_array(vlm_arrays, "target_xy", vlm_npz_path)
+    vlm_action = _require_array(vlm_arrays, "generated_action", vlm_npz_path)
+    vlm_xy = _require_array(vlm_arrays, "generated_xy", vlm_npz_path)
+    diffusion_target_action = _require_array(diffusion_arrays, "target_action", diffusion_npz_path)
+    diffusion_target_xy = _require_array(diffusion_arrays, "target_xy", diffusion_npz_path)
+    diffusion_action = _require_array(diffusion_arrays, "predicted_action", diffusion_npz_path)
+    diffusion_xy = _require_array(diffusion_arrays, "predicted_xy", diffusion_npz_path)
+    history_xy = _require_array(
+        diffusion_arrays, "history_xy", diffusion_npz_path, allow_batch=False
+    )
+    target_token_values = _require_token_values(vlm_arrays, "quantized_target_values", vlm_npz_path)
+    generated_token_values = _require_token_values(vlm_arrays, "generated_values", vlm_npz_path)
+    dataset_history_xy = _require_xy_prefix(trajectory_arrays, "history_xyz", trajectory_path)
+    dataset_target_xy = _require_xy_prefix(trajectory_arrays, "future_xyz", trajectory_path)
+    if "waypoint_indices" not in trajectory_arrays:
+        raise KeyError(f"Missing NPZ key 'waypoint_indices' in {trajectory_path}")
+    source_future_indices = np.asarray(trajectory_arrays["waypoint_indices"]).reshape(-1)
+    if not np.issubdtype(source_future_indices.dtype, np.integer):
+        raise ValueError(f"{trajectory_path}:waypoint_indices must contain integers")
+
+    action_truth_difference = _assert_same_truth(
+        "action", vlm_target_action, diffusion_target_action, vlm_npz_path, diffusion_npz_path
+    )
+    xy_truth_difference = _assert_same_truth(
+        "XY", vlm_target_xy, diffusion_target_xy, vlm_npz_path, diffusion_npz_path
+    )
+    dataset_xy_difference = _assert_same_truth(
+        "dataset XY", vlm_target_xy, dataset_target_xy, vlm_npz_path, trajectory_path
+    )
+    history_xy_difference = _assert_same_truth(
+        "history XY", history_xy, dataset_history_xy, diffusion_npz_path, trajectory_path
+    )
+    target_action = vlm_target_action
+    target_xy = vlm_target_xy
+    expected_shape = target_action.shape
+    for name, value in {
+        "target_xy": target_xy,
+        "vlm_action": vlm_action,
+        "vlm_xy": vlm_xy,
+        "diffusion_action": diffusion_action,
+        "diffusion_xy": diffusion_xy,
+    }.items():
+        if value.shape != expected_shape:
+            raise ValueError(f"{name} shape {value.shape} does not match target {expected_shape}")
+    if source_future_indices.shape != (target_xy.shape[0],):
+        raise ValueError(
+            "waypoint_indices length does not match the compared trajectory: "
+            f"{source_future_indices.shape} vs {target_xy.shape[0]}"
+        )
+    if target_token_values.shape != generated_token_values.shape:
+        raise ValueError(
+            "VLM generated and target trajectory token shapes differ: "
+            f"{generated_token_values.shape} vs {target_token_values.shape}"
+        )
+    if target_token_values.size != target_action.size:
+        raise ValueError(
+            f"Expected {target_action.size} trajectory tokens, got {target_token_values.size}"
+        )
+    token_matches = generated_token_values == target_token_values
+    waypoint_times_s = (source_future_indices.astype(np.float64) + 1.0) * args.time_step_s
+
+    action_mean = np.asarray(expert_config.get("action_mean"), dtype=np.float64)
+    action_scale = np.asarray(expert_config.get("action_scale"), dtype=np.float64)
+    if action_mean.shape != (2,) or action_scale.shape != (2,):
+        raise ValueError("expert_config action_mean and action_scale must both have shape (2,)")
+    if not np.isfinite(action_mean).all() or not np.isfinite(action_scale).all():
+        raise ValueError("expert_config action statistics must be finite")
+    if np.any(action_scale <= 0):
+        raise ValueError("expert_config action_scale must be positive")
+    action_decode_differences = {
+        "ground_truth": _validate_decoded_xy(
+            "ground-truth",
+            _decode_action(target_action, history_xy, action_mean, action_scale),
+            target_xy,
+        ),
+        "vlm": _validate_decoded_xy(
+            "VLM",
+            _decode_action(vlm_action, history_xy, action_mean, action_scale),
+            vlm_xy,
+        ),
+        "diffusion": _validate_decoded_xy(
+            "diffusion",
+            _decode_action(diffusion_action, history_xy, action_mean, action_scale),
+            diffusion_xy,
+        ),
+    }
+
+    untruncated = vlm_metrics.get("untruncated_vlm_generation")
+    if not isinstance(untruncated, dict) or not isinstance(untruncated.get("raw_text"), str):
+        raise KeyError("VLM metrics lack untruncated_vlm_generation.raw_text")
+    vlm_structured_text = untruncated["raw_text"]
+    target_structured_text = vlm_metrics.get("structured_target")
+    if not isinstance(target_structured_text, str):
+        raise KeyError("VLM metrics lack structured_target")
+    vlm_reasoning = _extract_component(vlm_structured_text, "cot")
+    vlm_meta_action = _extract_component(vlm_structured_text, "meta_action")
+    vlm_target_reasoning = _extract_component(target_structured_text, "cot")
+    vlm_target_meta_action = _extract_component(target_structured_text, "meta_action")
+    ground_truth_reasoning = manifest.get("reasoning")
+    if not isinstance(ground_truth_reasoning, str):
+        raise KeyError("Manifest row lacks string field 'reasoning'")
+    diffusion_target_reasoning = diffusion_metrics.get("target_reasoning")
+    if vlm_target_reasoning != ground_truth_reasoning:
+        raise ValueError("VLM target reasoning does not match the manifest ground truth")
+    if diffusion_target_reasoning != ground_truth_reasoning:
+        raise ValueError("Diffusion target reasoning does not match the manifest ground truth")
+    if vlm_metrics.get("target_meta_action") != vlm_target_meta_action:
+        raise ValueError("Structured target meta-action does not match VLM metrics")
+
+    vlm_position_errors = np.linalg.norm(vlm_xy - target_xy, axis=-1)
+    diffusion_position_errors = np.linalg.norm(diffusion_xy - target_xy, axis=-1)
+    vlm_action_errors = np.linalg.norm(vlm_action - target_action, axis=-1)
+    diffusion_action_errors = np.linalg.norm(diffusion_action - target_action, axis=-1)
+
+    vlm_computed_metrics = _prediction_metrics(vlm_action, vlm_xy, target_action, target_xy)
+    diffusion_computed_metrics = _prediction_metrics(
+        diffusion_action, diffusion_xy, target_action, target_xy
+    )
+    vlm_reported_metrics = _reported_metrics(vlm_metrics, section="untruncated_vlm_generation")
+    diffusion_reported_metrics = _reported_metrics(diffusion_metrics)
+    vlm_metric_validation = _metric_validation(vlm_computed_metrics, vlm_reported_metrics)
+    diffusion_metric_validation = _metric_validation(
+        diffusion_computed_metrics, diffusion_reported_metrics
+    )
+
+    waypoints: list[dict[str, Any]] = []
+    for index in range(target_xy.shape[0]):
+        waypoints.append(
+            {
+                "index": index + 1,
+                "source_future_index": int(source_future_indices[index]),
+                "time_s": float(waypoint_times_s[index]),
+                "action": {
+                    "ground_truth": target_action[index].tolist(),
+                    "vlm": vlm_action[index].tolist(),
+                    "diffusion": diffusion_action[index].tolist(),
+                },
+                "action_error_l2": {
+                    "vlm": float(vlm_action_errors[index]),
+                    "diffusion": float(diffusion_action_errors[index]),
+                },
+                "xy_m": {
+                    "ground_truth": target_xy[index].tolist(),
+                    "vlm": vlm_xy[index].tolist(),
+                    "diffusion": diffusion_xy[index].tolist(),
+                },
+                "position_error_m": {
+                    "vlm": float(vlm_position_errors[index]),
+                    "diffusion": float(diffusion_position_errors[index]),
+                },
+            }
+        )
+
+    comparison = {
+        "schema_version": 1,
+        "sample": {
+            "manifest_index": args.manifest_index,
+            "clip_id": manifest.get("clip_id"),
+            "t0_us": manifest.get("t0_us"),
+            "camera": manifest.get("camera"),
+            "history_steps": manifest.get("history_steps"),
+            "expert_waypoints": manifest.get("expert_waypoints"),
+            "source_time_step_s": args.time_step_s,
+            "image_path": _display_path(_resolve(Path(str(manifest.get("image_path"))))),
+            "trajectory_path": _display_path(trajectory_path),
+        },
+        "prediction_protocol": {
+            "vlm": {
+                "decoding": "greedy untruncated structured generation",
+                "trajectory_representation": "32 discrete normalized-XY-delta tokens",
+                "trajectory_vocabulary_bins": int(vlm_metrics["vocabulary"]["size"]),
+            },
+            "diffusion": {
+                "decoding": "flow-matching sample",
+                "seed": args.diffusion_seed,
+                "inference_steps": expert_config.get("inference_steps"),
+            },
+            "units": {
+                "action_mse": "normalized_action_squared",
+                "trajectory": "metres",
+            },
+        },
+        "text": {
+            "ground_truth_reasoning": ground_truth_reasoning,
+            "vlm_reasoning": vlm_reasoning,
+            "vlm_reasoning_exact_match": vlm_reasoning == ground_truth_reasoning,
+            "vlm_meta_action": vlm_meta_action,
+            "target_meta_action": vlm_target_meta_action,
+            "vlm_meta_action_exact_match": vlm_meta_action == vlm_target_meta_action,
+            "vlm_structured_output": vlm_structured_text,
+            "target_structured_output": target_structured_text,
+            "vlm_structured_exact_match": vlm_structured_text == target_structured_text,
+            "diffusion_conditioning_reasoning": diffusion_metrics.get("generated_reasoning"),
+            "diffusion_reasoning_exact_match": diffusion_metrics.get("generated_reasoning")
+            == ground_truth_reasoning,
+        },
+        "arrays": {
+            "ground_truth": {
+                "action": target_action.tolist(),
+                "xy_m": target_xy.tolist(),
+                "discrete_token_values": target_token_values.tolist(),
+            },
+            "vlm": {
+                "action": vlm_action.tolist(),
+                "xy_m": vlm_xy.tolist(),
+                "discrete_token_values": generated_token_values.tolist(),
+            },
+            "diffusion": {
+                "action": diffusion_action.tolist(),
+                "xy_m": diffusion_xy.tolist(),
+                "seed": args.diffusion_seed,
+            },
+            "history_xy_m": history_xy.tolist(),
+            "source_future_indices": source_future_indices.tolist(),
+            "waypoint_times_s": waypoint_times_s.tolist(),
+        },
+        "waypoints": waypoints,
+        "metrics": {
+            "vlm": {
+                "computed_from_npz": vlm_computed_metrics,
+                "reported_in_metrics_json": vlm_reported_metrics,
+                "trajectory_token_matches": int(token_matches.sum()),
+                "trajectory_token_count": int(token_matches.size),
+                "trajectory_token_accuracy": float(token_matches.mean()),
+            },
+            "diffusion": {
+                "computed_from_npz": diffusion_computed_metrics,
+                "reported_in_metrics_json": diffusion_reported_metrics,
+                "seed": args.diffusion_seed,
+            },
+            "diffusion_minus_vlm": {
+                name: diffusion_computed_metrics[name] - vlm_computed_metrics[name]
+                for name in METRIC_NAMES
+            },
+        },
+        "validation": {
+            "ground_truth_action_exact_match": True,
+            "ground_truth_action_max_abs_difference": action_truth_difference,
+            "ground_truth_xy_exact_match": True,
+            "ground_truth_xy_max_abs_difference": xy_truth_difference,
+            "dataset_target_xy_exact_match": True,
+            "dataset_target_xy_max_abs_difference": dataset_xy_difference,
+            "dataset_history_xy_exact_match": True,
+            "dataset_history_xy_max_abs_difference": history_xy_difference,
+            "ground_truth_text_matches_manifest": True,
+            "vlm_discrete_tokens_exact_match": bool(token_matches.all()),
+            "action_decode_max_abs_differences": action_decode_differences,
+            "vlm_reported_metrics_match_npz": vlm_metric_validation,
+            "diffusion_reported_metrics_match_npz": diffusion_metric_validation,
+        },
+        "data_sources": {
+            "manifest": _source(manifest_path),
+            "vlm_npz": _source(vlm_npz_path, keys=sorted(vlm_arrays)),
+            "vlm_metrics": _source(vlm_metrics_path),
+            "diffusion_npz": _source(diffusion_npz_path, keys=sorted(diffusion_arrays)),
+            "diffusion_metrics": _source(diffusion_metrics_path),
+            "expert_config": _source(expert_config_path),
+            "trajectory_npz": _source(trajectory_path, keys=sorted(trajectory_arrays)),
+        },
+        "artifacts": {
+            "comparison_json": _display_path(output_dir / "comparison.json"),
+            "waypoints_csv": _display_path(output_dir / "waypoints.csv"),
+            "trajectory_plot": _display_path(output_dir / "trajectory_comparison.png"),
+        },
+    }
+
+    comparison_path = output_dir / "comparison.json"
+    waypoints_path = output_dir / "waypoints.csv"
+    plot_path = output_dir / "trajectory_comparison.png"
+    comparison_path.write_text(
+        json.dumps(comparison, indent=2, ensure_ascii=False, allow_nan=False) + "\n",
+        encoding="utf-8",
+    )
+    _write_waypoints(waypoints_path, waypoints)
+    _save_plot(
+        plot_path,
+        history_xy,
+        target_xy,
+        vlm_xy,
+        diffusion_xy,
+        target_action,
+        vlm_action,
+        diffusion_action,
+        waypoint_times_s,
+        vlm_position_errors,
+        diffusion_position_errors,
+        ground_truth_reasoning,
+        vlm_reasoning,
+    )
+    print(
+        json.dumps(
+            {
+                "comparison_json": str(comparison_path),
+                "waypoints_csv": str(waypoints_path),
+                "trajectory_plot": str(plot_path),
+                "vlm_metrics": vlm_computed_metrics,
+                "diffusion_metrics": diffusion_computed_metrics,
+                "diffusion_seed": args.diffusion_seed,
+                "vlm_trajectory_tokens": {
+                    "matches": int(token_matches.sum()),
+                    "count": int(token_matches.size),
+                },
+                "vlm_reasoning_exact_match": vlm_reasoning == ground_truth_reasoning,
+                "ground_truth_consistent": True,
+            },
+            indent=2,
+            ensure_ascii=False,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/download_model.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+
+from huggingface_hub import model_info, snapshot_download
+
+from alpamayo_lite.config import load_config
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/local_overfit.yaml")
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    output_dir = cfg.resolve(cfg.paths.model_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    info = model_info(cfg.vlm.repo_id, files_metadata=True)
+    total = info.safetensors.total if info.safetensors is not None else None
+    if total is not None and not (700_000_000 <= total <= 900_000_000):
+        raise RuntimeError(f"Unexpected model size for {cfg.vlm.repo_id}: {total:,} parameters")
+
+    snapshot_download(
+        repo_id=cfg.vlm.repo_id,
+        local_dir=output_dir,
+        allow_patterns=[
+            "*.json",
+            "*.txt",
+            "*.safetensors",
+        ],
+    )
+    provenance = {
+        "repo_id": cfg.vlm.repo_id,
+        "revision": info.sha,
+        "safetensors_parameters": total,
+        "pipeline_tag": info.pipeline_tag,
+        "library_name": info.library_name,
+    }
+    (output_dir / "download_provenance.json").write_text(
+        json.dumps(provenance, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
+    )
+    print(json.dumps(provenance, indent=2, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()

scripts/evaluate_pipeline.py

@@ -0,0 +1,244 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+from functools import partial
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import load_file
+
+from alpamayo_lite.action_space import RelativeDeltaActionSpace
+from alpamayo_lite.config import ExperimentConfig, load_config
+from alpamayo_lite.data import OverfitSample, load_manifest
+from alpamayo_lite.diffusion import FlowMatching
+from alpamayo_lite.evaluation import (
+    save_prediction_arrays,
+    save_trajectory_plot,
+    trajectory_metrics,
+)
+from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert, PrefixKV
+from alpamayo_lite.prompts import build_expert_prefix, build_vlm_prefix
+from alpamayo_lite.runtime import (
+    empty_device_cache,
+    parameter_counts,
+    resolve_dtype,
+    seed_everything,
+    select_device,
+    write_json,
+)
+from alpamayo_lite.vlm import (
+    extract_prefix_kv,
+    generate_reasoning,
+    load_vlm,
+    vlm_text_config,
+)
+
+
+def _resolve(cfg: ExperimentConfig, value: str) -> Path:
+    path = Path(value).expanduser()
+    return path.resolve() if path.is_absolute() else (cfg.root / path).resolve()
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _summary(values: list[float]) -> dict[str, float]:
+    return {
+        "mean": statistics.fmean(values),
+        "std": statistics.pstdev(values),
+        "min": min(values),
+        "max": max(values),
+    }
+
+
+def _evaluate_expert(
+    stage2_dir: Path,
+    vlm: Any,
+    prefix: PrefixKV,
+    sample: OverfitSample,
+    device: torch.device,
+    dtype: torch.dtype,
+    num_rollouts: int,
+    base_seed: int,
+) -> dict[str, Any]:
+    expert_config = _read_json(stage2_dir / "expert_config.json")
+    metrics_path = stage2_dir / "stage2_metrics.json"
+    saved_metrics = _read_json(metrics_path) if metrics_path.exists() else {}
+    architecture = ExpertArchitecture(**expert_config["architecture"])
+    expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
+        device=device, dtype=dtype
+    )
+    expert.load_state_dict(load_file(stage2_dir / "expert.safetensors"), strict=True)
+    expert.requires_grad_(False)
+    expert.eval()
+
+    action_shape = tuple(int(value) for value in expert_config["action_shape"])
+    if len(action_shape) != 2:
+        raise ValueError(f"Invalid action shape in {stage2_dir}: {action_shape}")
+    flow = FlowMatching(
+        action_shape=action_shape,
+        inference_steps=int(expert_config["inference_steps"]),
+    )
+    action_space = RelativeDeltaActionSpace(
+        torch.tensor(expert_config["action_mean"], device=device),
+        torch.tensor(expert_config["action_scale"], device=device),
+    ).to(device)
+    history, future = sample.load_trajectory()
+    history_xy = history[None, :, :2].to(device)
+    target_xy = future[None, :, :2].to(device)
+    target_action = action_space.trajectory_to_action(history_xy, target_xy).float()
+
+    rollouts: list[dict[str, float | int]] = []
+    first_prediction: tuple[torch.Tensor, torch.Tensor] | None = None
+    for index in range(num_rollouts):
+        seed = base_seed + index
+        seed_everything(seed)
+        predicted_action = flow.sample(
+            batch_size=1,
+            step_fn=partial(expert, prefix=prefix),
+            device=device,
+            dtype=dtype,
+        ).float()
+        predicted_xy = action_space.action_to_trajectory(predicted_action, history_xy)
+        record: dict[str, float | int] = {
+            "seed": seed,
+            "action_mse": float(
+                torch.nn.functional.mse_loss(predicted_action, target_action).cpu()
+            ),
+            **trajectory_metrics(predicted_xy, target_xy),
+        }
+        if not torch.isfinite(predicted_action).all() or not torch.isfinite(predicted_xy).all():
+            raise RuntimeError(f"Non-finite reload prediction for seed {seed}")
+        rollouts.append(record)
+        if first_prediction is None:
+            first_prediction = (predicted_action, predicted_xy)
+
+    assert first_prediction is not None
+    first_action, first_xy = first_prediction
+    save_prediction_arrays(
+        stage2_dir / "reload_prediction.npz",
+        history[:, :2],
+        target_action[0],
+        first_action[0],
+        target_xy[0],
+        first_xy[0],
+    )
+    save_trajectory_plot(
+        stage2_dir / "reload_trajectory.png",
+        history[:, :2],
+        target_xy[0],
+        first_xy[0],
+    )
+
+    metric_names = ("action_mse", "ade_m", "fde_m", "max_error_m")
+    result = {
+        "stage2_dir": str(stage2_dir),
+        "checkpoint_steps": int(
+            expert_config.get("completed_steps", saved_metrics.get("steps", 0))
+        ),
+        "expert_parameters": parameter_counts(expert),
+        "architecture": expert_config["architecture"],
+        "action_mean": expert_config["action_mean"],
+        "action_scale": expert_config["action_scale"],
+        "action_shape": list(action_shape),
+        "inference_steps": flow.inference_steps,
+        "num_rollouts": num_rollouts,
+        "rollouts": rollouts,
+        "summary": {
+            name: _summary([float(record[name]) for record in rollouts])
+            for name in metric_names
+        },
+    }
+    del expert
+    empty_device_cache(device)
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/local_overfit.yaml")
+    parser.add_argument(
+        "--stage2-dir",
+        action="append",
+        dest="stage2_dirs",
+        help="Stage-2 directory to reload; repeat to compare checkpoints.",
+    )
+    parser.add_argument("--num-rollouts", type=int, default=16)
+    parser.add_argument("--base-seed", type=int, default=None)
+    parser.add_argument("--output", default="outputs/reload_evaluation.json")
+    args = parser.parse_args()
+    if args.num_rollouts <= 0:
+        raise ValueError("--num-rollouts must be positive")
+
+    cfg = load_config(args.config)
+    seed_everything(cfg.seed)
+    device = select_device()
+    dtype = resolve_dtype(cfg.vlm.dtype, device)
+    stage2_values = args.stage2_dirs or [cfg.paths.stage2_dir]
+    stage2_dirs = [_resolve(cfg, value) for value in stage2_values]
+    sample = load_manifest(cfg.resolve(cfg.paths.manifest))[0]
+    history, _ = sample.load_trajectory()
+
+    stage1_dir = cfg.resolve(cfg.paths.stage1_dir)
+    processor, vlm = load_vlm(stage1_dir, cfg.vlm, device, dtype)
+    vlm.requires_grad_(False)
+    vlm.eval()
+    generated_reasoning = generate_reasoning(
+        vlm,
+        processor,
+        sample.load_image(),
+        build_vlm_prefix(history),
+        device,
+        dtype,
+        cfg.vlm.max_reasoning_tokens,
+    )
+    prefix = extract_prefix_kv(
+        vlm,
+        processor,
+        sample.load_image(),
+        build_expert_prefix(history, generated_reasoning),
+        device,
+        dtype,
+    )
+    base_seed = args.base_seed if args.base_seed is not None else cfg.seed + 1
+    models = [
+        _evaluate_expert(
+            stage2_dir,
+            vlm,
+            prefix,
+            sample,
+            device,
+            dtype,
+            args.num_rollouts,
+            base_seed,
+        )
+        for stage2_dir in stage2_dirs
+    ]
+    result = {
+        "stage": "reloaded_pipeline_evaluation",
+        "device": str(device),
+        "dtype": str(dtype),
+        "stage1_dir": str(stage1_dir),
+        "vlm_parameters": parameter_counts(vlm),
+        "generated_reasoning": generated_reasoning,
+        "target_reasoning": sample.reasoning,
+        "reasoning_exact_match": generated_reasoning == sample.reasoning,
+        "prefix_layers": prefix.num_layers,
+        "prefix_sequence_length": prefix.sequence_length,
+        "prefix_key_shape": list(prefix.keys[0].shape),
+        "models": models,
+    }
+    output = _resolve(cfg, args.output)
+    write_json(output, result)
+    print(json.dumps(result, indent=2, ensure_ascii=False))
+    del vlm
+    empty_device_cache(device)
+
+
+if __name__ == "__main__":
+    main()

scripts/materialize_pai_subset.py

@@ -0,0 +1,204 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from huggingface_hub import get_token, hf_hub_download
+from physical_ai_av import PhysicalAIAVDatasetInterface
+from PIL import Image
+from scipy.spatial.transform import Rotation
+
+from alpamayo_lite.config import load_config
+from alpamayo_lite.prompts import clean_reasoning_text
+
+
+CLIP_DURATION_US = 20_000_000
+
+
+def parse_events(cell: object) -> list[dict[str, Any]]:
+    if cell is None or (np.isscalar(cell) and pd.isna(cell)):
+        return []
+    parsed = json.loads(cell) if isinstance(cell, str) else cell
+    if not isinstance(parsed, (list, tuple)):
+        return []
+    return [event for event in parsed if isinstance(event, dict)]
+
+
+def collect_candidates(
+    reasoning_path: Path,
+    history_seconds: float,
+    future_seconds: float,
+) -> list[dict[str, Any]]:
+    frame = pd.read_parquet(reasoning_path)
+    candidates: list[dict[str, Any]] = []
+    for clip_id, events_cell in frame["events"].items():
+        for event in parse_events(events_cell):
+            if "event_start_timestamp" not in event:
+                continue
+            t0_us = int(event["event_start_timestamp"])
+            if t0_us < int(history_seconds * 1_000_000):
+                continue
+            if t0_us + int(future_seconds * 1_000_000) > CLIP_DURATION_US:
+                continue
+            # Current PAI releases call the driving rationale/command `coc`.
+            # Keep `cot` as a compatibility fallback for older metadata shards.
+            reasoning = clean_reasoning_text(str(event.get("coc") or event.get("cot") or ""))
+            candidates.append({"clip_id": str(clip_id), "t0_us": t0_us, "reasoning": reasoning})
+    if not candidates:
+        raise RuntimeError("No valid reasoning events remain after trajectory safety filtering")
+    return candidates
+
+
+def load_local_trajectory(
+    interface: PhysicalAIAVDatasetInterface,
+    clip_id: str,
+    t0_us: int,
+    history_steps: int,
+    future_steps: int,
+    time_step: float,
+) -> tuple[np.ndarray, np.ndarray]:
+    egomotion = interface.get_clip_feature(
+        clip_id, interface.features.LABELS.EGOMOTION, maybe_stream=True
+    )
+    history_offsets = np.arange(-(history_steps - 1), 1, dtype=np.int64)
+    history_timestamps = t0_us + history_offsets * int(time_step * 1_000_000)
+    future_offsets = np.arange(1, future_steps + 1, dtype=np.int64)
+    future_timestamps = t0_us + future_offsets * int(time_step * 1_000_000)
+
+    history_state = egomotion(history_timestamps)
+    future_state = egomotion(future_timestamps)
+    history_xyz = history_state.pose.translation
+    future_xyz = future_state.pose.translation
+    t0_rotation = Rotation.from_quat(history_state.pose.rotation.as_quat()[-1])
+    inverse = t0_rotation.inv()
+    origin = history_xyz[-1].copy()
+    history_local = inverse.apply(history_xyz - origin).astype(np.float32)
+    future_local = inverse.apply(future_xyz - origin).astype(np.float32)
+    return history_local, future_local
+
+
+def load_camera_frame(
+    interface: PhysicalAIAVDatasetInterface,
+    clip_id: str,
+    t0_us: int,
+    camera_name: str,
+) -> np.ndarray:
+    camera_feature = getattr(interface.features.CAMERA, camera_name.upper())
+    reader = interface.get_clip_feature(clip_id, camera_feature, maybe_stream=True)
+    frames, _ = reader.decode_images_from_timestamps(np.asarray([t0_us], dtype=np.int64))
+    return frames[0]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/local_overfit.yaml")
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    token = os.environ.get("HF_TOKEN") or get_token()
+    if not token:
+        raise SystemExit(
+            "PAI is gated. Set HF_TOKEN for an account that accepted the dataset license."
+        )
+
+    raw_dir = cfg.resolve(cfg.paths.pai_raw_dir)
+    raw_dir.mkdir(parents=True, exist_ok=True)
+    cache_dir = cfg.root / ".cache" / "huggingface"
+    reasoning_file = Path(
+        hf_hub_download(
+            repo_id=cfg.data.repo_id,
+            repo_type="dataset",
+            filename="reasoning/ood_reasoning.parquet",
+            token=token,
+            cache_dir=cache_dir,
+            local_dir=raw_dir,
+        )
+    )
+    candidates = collect_candidates(
+        reasoning_file,
+        history_seconds=(cfg.data.history_steps - 1) * cfg.data.time_step,
+        future_seconds=cfg.data.source_future_steps * cfg.data.time_step,
+    )
+    rng = np.random.default_rng(cfg.seed)
+    selected_indices = rng.choice(
+        len(candidates), size=min(cfg.data.num_clips, len(candidates)), replace=False
+    )
+    selected = [candidates[int(index)] for index in selected_indices]
+
+    interface = PhysicalAIAVDatasetInterface(
+        token=token,
+        cache_dir=cache_dir,
+        local_dir=raw_dir,
+        confirm_download_threshold_gb=1000.0,
+    )
+    processed_dir = cfg.resolve(cfg.paths.manifest).parent
+    image_dir = processed_dir / "images"
+    trajectory_dir = processed_dir / "trajectories"
+    image_dir.mkdir(parents=True, exist_ok=True)
+    trajectory_dir.mkdir(parents=True, exist_ok=True)
+
+    rows: list[dict[str, Any]] = []
+    for sample_index, event in enumerate(selected):
+        clip_id = event["clip_id"]
+        t0_us = event["t0_us"]
+        history, future_full = load_local_trajectory(
+            interface=interface,
+            clip_id=clip_id,
+            t0_us=t0_us,
+            history_steps=cfg.data.history_steps,
+            future_steps=cfg.data.source_future_steps,
+            time_step=cfg.data.time_step,
+        )
+        waypoint_indices = np.linspace(
+            0,
+            cfg.data.source_future_steps - 1,
+            cfg.data.expert_waypoints,
+            dtype=np.int64,
+        )
+        future = future_full[waypoint_indices]
+        frame = load_camera_frame(
+            interface=interface,
+            clip_id=clip_id,
+            t0_us=t0_us,
+            camera_name=cfg.data.camera,
+        )
+
+        stem = f"sample_{sample_index:02d}_{clip_id}_{t0_us}"
+        image_path = image_dir / f"{stem}.jpg"
+        trajectory_path = trajectory_dir / f"{stem}.npz"
+        Image.fromarray(frame).save(image_path, quality=95)
+        np.savez_compressed(
+            trajectory_path,
+            history_xyz=history,
+            future_xyz=future,
+            source_future_xyz=future_full,
+            waypoint_indices=waypoint_indices,
+        )
+        rows.append(
+            {
+                "clip_id": clip_id,
+                "t0_us": t0_us,
+                "reasoning": event["reasoning"],
+                "image_path": str(image_path.relative_to(cfg.root)),
+                "trajectory_path": str(trajectory_path.relative_to(cfg.root)),
+                "camera": cfg.data.camera,
+                "history_steps": cfg.data.history_steps,
+                "expert_waypoints": cfg.data.expert_waypoints,
+            }
+        )
+
+    manifest_path = cfg.resolve(cfg.paths.manifest)
+    manifest_path.write_text(
+        "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
+        encoding="utf-8",
+    )
+    print(json.dumps({"manifest": str(manifest_path), "samples": rows}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/train_stage1_vlm.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import time
+
+import torch
+from tqdm import trange
+
+from alpamayo_lite.config import load_config
+from alpamayo_lite.data import load_manifest
+from alpamayo_lite.prompts import build_stage1_text, build_vlm_prefix
+from alpamayo_lite.runtime import (
+    empty_device_cache,
+    resolve_dtype,
+    seed_everything,
+    select_device,
+    write_json,
+)
+from alpamayo_lite.vlm import (
+    configure_stage1_trainability,
+    enforce_frozen_module_eval,
+    generate_reasoning,
+    load_vlm,
+    prepare_stage1_batch,
+)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/local_overfit.yaml")
+    parser.add_argument("--steps", type=int, default=None)
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    seed_everything(cfg.seed)
+    device = select_device()
+    dtype = resolve_dtype(cfg.vlm.dtype, device)
+    samples = load_manifest(cfg.resolve(cfg.paths.manifest))
+
+    processor, model = load_vlm(cfg.resolve(cfg.paths.model_dir), cfg.vlm, device, dtype)
+    first = samples[0]
+    first_history, _ = first.load_trajectory()
+    baseline = generate_reasoning(
+        model,
+        processor,
+        first.load_image(),
+        build_vlm_prefix(first_history),
+        device,
+        dtype,
+        cfg.vlm.max_reasoning_tokens,
+    )
+    trainability = configure_stage1_trainability(model, cfg.stage1)
+    print({"device": str(device), "dtype": str(dtype), "trainability": trainability})
+
+    batches = []
+    for sample in samples:
+        history, _ = sample.load_trajectory()
+        prefix, full_text = build_stage1_text(history, sample.reasoning)
+        inputs, labels, prefix_length = prepare_stage1_batch(
+            processor,
+            sample.load_image(),
+            prefix,
+            full_text,
+            device,
+            dtype,
+        )
+        batches.append((inputs, labels, prefix_length))
+
+    optimizer = torch.optim.AdamW(
+        [parameter for parameter in model.parameters() if parameter.requires_grad],
+        lr=cfg.stage1.learning_rate,
+        weight_decay=cfg.stage1.weight_decay,
+        eps=cfg.stage1.adam_epsilon,
+    )
+    steps = args.steps if args.steps is not None else cfg.stage1.steps
+    losses: list[float] = []
+    started = time.perf_counter()
+    model.train()
+    enforce_frozen_module_eval(model)
+    for step in trange(steps, desc="stage1-vlm"):
+        inputs, labels, _ = batches[step % len(batches)]
+        optimizer.zero_grad(set_to_none=True)
+        outputs = model(**inputs, labels=labels, use_cache=False, return_dict=True)
+        loss = outputs.loss
+        if loss is None or not torch.isfinite(loss):
+            raise RuntimeError(f"Non-finite Stage-1 loss at step {step}: {loss}")
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(
+            [parameter for parameter in model.parameters() if parameter.requires_grad], 1.0
+        )
+        optimizer.step()
+        loss_value = float(loss.detach().cpu())
+        losses.append(loss_value)
+        if step == 0 or (step + 1) % cfg.stage1.log_every == 0:
+            print(f"stage1 step={step + 1} loss={loss_value:.6f}")
+
+    elapsed = time.perf_counter() - started
+    model.gradient_checkpointing_disable()
+    model.config.use_cache = True
+    model.model.language_model.config.use_cache = True
+    output_dir = cfg.resolve(cfg.paths.stage1_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(output_dir, safe_serialization=True, max_shard_size="2GB")
+    processor.save_pretrained(output_dir)
+    trained = generate_reasoning(
+        model,
+        processor,
+        first.load_image(),
+        build_vlm_prefix(first_history),
+        device,
+        dtype,
+        cfg.vlm.max_reasoning_tokens,
+    )
+    metrics = {
+        "stage": "stage1_vlm",
+        "device": str(device),
+        "dtype": str(dtype),
+        "trainability": trainability,
+        "steps": steps,
+        "elapsed_seconds": elapsed,
+        "initial_loss": losses[0],
+        "final_loss": losses[-1],
+        "minimum_loss": min(losses),
+        "losses": losses,
+        "baseline_text": baseline,
+        "trained_text": trained,
+        "target_text": first.reasoning,
+    }
+    write_json(output_dir / "stage1_metrics.json", metrics)
+    print(metrics)
+    del optimizer, model
+    empty_device_cache(device)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_stage2_expert.py

@@ -0,0 +1,451 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from dataclasses import asdict
+from functools import partial
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import load_file, save_file
+from tqdm import trange
+
+from alpamayo_lite.action_space import RelativeDeltaActionSpace
+from alpamayo_lite.config import load_config
+from alpamayo_lite.data import load_manifest
+from alpamayo_lite.diffusion import FlowMatching
+from alpamayo_lite.evaluation import (
+    save_prediction_arrays,
+    save_trajectory_plot,
+    trajectory_metrics,
+)
+from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert
+from alpamayo_lite.prompts import build_expert_prefix, build_vlm_prefix
+from alpamayo_lite.runtime import (
+    empty_device_cache,
+    parameter_counts,
+    resolve_dtype,
+    seed_everything,
+    select_device,
+    write_json,
+)
+from alpamayo_lite.vlm import (
+    extract_prefix_kv,
+    generate_reasoning,
+    load_vlm,
+    vlm_text_config,
+)
+
+
+def _cpu_copy(value: Any) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    if isinstance(value, dict):
+        return {key: _cpu_copy(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_cpu_copy(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_cpu_copy(item) for item in value)
+    return value
+
+
+def _optimizer_to(optimizer: torch.optim.Optimizer, device: torch.device) -> None:
+    for state in optimizer.state.values():
+        for key, value in state.items():
+            if isinstance(value, torch.Tensor):
+                state[key] = value.to(device)
+
+
+def _atomic_save_safetensors(module: torch.nn.Module, path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.tmp")
+    save_file(
+        {key: value.detach().cpu().contiguous() for key, value in module.state_dict().items()},
+        temporary,
+    )
+    os.replace(temporary, path)
+
+
+def _atomic_torch_save(value: Any, path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.tmp")
+    torch.save(value, temporary)
+    os.replace(temporary, path)
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _save_training_checkpoint(
+    checkpoint_dir: Path,
+    expert: KVCachedDiffusionExpert,
+    optimizer: torch.optim.Optimizer,
+    architecture: ExpertArchitecture,
+    completed_steps: int,
+    losses: list[float],
+    elapsed_seconds: float,
+    device: torch.device,
+) -> None:
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    _atomic_save_safetensors(expert, checkpoint_dir / "expert.safetensors")
+    state: dict[str, Any] = {
+        "completed_steps": completed_steps,
+        "losses": losses,
+        "elapsed_seconds": elapsed_seconds,
+        "architecture": asdict(architecture),
+        "optimizer": _cpu_copy(optimizer.state_dict()),
+        "torch_rng_state": torch.get_rng_state(),
+    }
+    if device.type == "mps" and hasattr(torch.mps, "get_rng_state"):
+        state["device_rng_state"] = torch.mps.get_rng_state().cpu()
+    elif device.type == "cuda":
+        state["device_rng_state"] = torch.cuda.get_rng_state(device).cpu()
+    _atomic_torch_save(state, checkpoint_dir / "training_state.pt")
+    write_json(
+        checkpoint_dir / "progress.json",
+        {
+            "completed_steps": completed_steps,
+            "latest_loss": losses[-1],
+            "minimum_loss": min(losses),
+            "elapsed_seconds": elapsed_seconds,
+            "architecture": asdict(architecture),
+        },
+    )
+
+
+def _load_resume_state(
+    output_dir: Path,
+    expert: KVCachedDiffusionExpert,
+    optimizer: torch.optim.Optimizer,
+    architecture: ExpertArchitecture,
+    device: torch.device,
+) -> tuple[int, list[float], float, bool, str]:
+    checkpoint_dir = output_dir / "checkpoints" / "latest"
+    checkpoint_weights = checkpoint_dir / "expert.safetensors"
+    checkpoint_state = checkpoint_dir / "training_state.pt"
+    if checkpoint_weights.exists() and checkpoint_state.exists():
+        state = torch.load(checkpoint_state, map_location="cpu", weights_only=True)
+        if state.get("architecture") != asdict(architecture):
+            raise RuntimeError("Checkpoint architecture does not match the current config")
+        expert.load_state_dict(load_file(checkpoint_weights), strict=True)
+        optimizer.load_state_dict(state["optimizer"])
+        _optimizer_to(optimizer, device)
+        torch.set_rng_state(state["torch_rng_state"])
+        device_rng_state = state.get("device_rng_state")
+        if device_rng_state is not None and device.type == "mps":
+            torch.mps.set_rng_state(device_rng_state)
+        elif device_rng_state is not None and device.type == "cuda":
+            torch.cuda.set_rng_state(device_rng_state, device)
+        return (
+            int(state["completed_steps"]),
+            [float(value) for value in state["losses"]],
+            float(state.get("elapsed_seconds", 0.0)),
+            True,
+            str(checkpoint_dir),
+        )
+
+    final_weights = output_dir / "expert.safetensors"
+    final_metrics = output_dir / "stage2_metrics.json"
+    final_config = output_dir / "expert_config.json"
+    if not (final_weights.exists() and final_metrics.exists() and final_config.exists()):
+        raise FileNotFoundError(
+            "No resumable Stage-2 checkpoint found. Expected either checkpoints/latest or "
+            "the legacy expert.safetensors + stage2_metrics.json + expert_config.json files."
+        )
+    saved_config = _read_json(final_config)
+    if saved_config.get("architecture") != asdict(architecture):
+        raise RuntimeError("Legacy expert architecture does not match the current config")
+    metrics = _read_json(final_metrics)
+    expert.load_state_dict(load_file(final_weights), strict=True)
+    completed_steps = int(metrics["steps"])
+    losses = [float(value) for value in metrics["losses"]]
+    if completed_steps != len(losses):
+        raise RuntimeError(
+            f"Legacy metrics are inconsistent: steps={completed_steps}, losses={len(losses)}"
+        )
+    return (
+        completed_steps,
+        losses,
+        float(metrics.get("elapsed_seconds", 0.0)),
+        False,
+        str(final_weights),
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/local_overfit.yaml")
+    parser.add_argument(
+        "--steps",
+        type=int,
+        default=None,
+        help="Total target steps. With --resume, only the missing steps are run.",
+    )
+    parser.add_argument(
+        "--resume",
+        action="store_true",
+        help="Resume from checkpoints/latest, or warm-start from the legacy final expert.",
+    )
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    seed_everything(cfg.seed)
+    device = select_device()
+    dtype = resolve_dtype(cfg.vlm.dtype, device)
+    samples = load_manifest(cfg.resolve(cfg.paths.manifest))
+
+    stage1_dir = cfg.resolve(cfg.paths.stage1_dir)
+    if not (stage1_dir / "stage1_metrics.json").exists():
+        raise SystemExit("Stage-1 checkpoint is missing; run train_stage1_vlm.py first")
+    processor, vlm = load_vlm(stage1_dir, cfg.vlm, device, dtype)
+    vlm.requires_grad_(False)
+    vlm.eval()
+    vlm_counts = parameter_counts(vlm)
+    if vlm_counts["trainable"] != 0:
+        raise RuntimeError(f"Stage-2 VLM must be fully frozen: {vlm_counts}")
+
+    histories = []
+    futures = []
+    images = []
+    prefixes = []
+    for sample in samples:
+        history, future = sample.load_trajectory()
+        histories.append(history[:, :2])
+        futures.append(future[:, :2])
+        images.append(sample.load_image())
+        prefixes.append(
+            extract_prefix_kv(
+                vlm,
+                processor,
+                images[-1],
+                build_expert_prefix(history, sample.reasoning),
+                device,
+                dtype,
+            )
+        )
+
+    history_tensor = torch.stack(histories).to(device=device, dtype=torch.float32)
+    future_tensor = torch.stack(futures).to(device=device, dtype=torch.float32)
+    statistics = RelativeDeltaActionSpace.fit(history_tensor, future_tensor)
+    action_space = RelativeDeltaActionSpace(statistics.mean, statistics.scale).to(device)
+    target_actions = action_space.trajectory_to_action(history_tensor, future_tensor).to(dtype)
+
+    stage2 = cfg.stage2
+    architecture = ExpertArchitecture(
+        hidden_size=stage2.expert_hidden_size,
+        intermediate_size=stage2.expert_intermediate_size,
+        num_layers=stage2.expert_num_layers,
+        num_attention_heads=stage2.expert_num_attention_heads,
+        num_key_value_heads=stage2.expert_num_key_value_heads,
+        head_dim=stage2.expert_head_dim,
+        action_projection_hidden_size=stage2.action_projection_hidden_size,
+        action_fourier_features=stage2.action_fourier_features,
+        non_causal_attention=stage2.expert_non_causal_attention,
+    )
+    expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
+        device=device, dtype=dtype
+    )
+    expert_counts = parameter_counts(expert)
+    if not (70_000_000 <= expert_counts["total"] <= 120_000_000):
+        raise RuntimeError(
+            f"Scaled expert parameter count is outside the expected range: {expert_counts}"
+        )
+    flow = FlowMatching(
+        action_shape=(cfg.data.expert_waypoints, 2),
+        inference_steps=stage2.inference_steps,
+    )
+    optimizer = torch.optim.AdamW(
+        expert.parameters(),
+        lr=stage2.learning_rate,
+        weight_decay=stage2.weight_decay,
+        eps=stage2.adam_epsilon,
+    )
+    target_steps = args.steps if args.steps is not None else stage2.steps
+    if target_steps <= 0:
+        raise ValueError("Stage-2 target steps must be positive")
+    if stage2.checkpoint_every <= 0:
+        raise ValueError("stage2.checkpoint_every must be positive")
+
+    output_dir = cfg.resolve(cfg.paths.stage2_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_dir = output_dir / "checkpoints" / "latest"
+    completed_steps = 0
+    losses: list[float] = []
+    previous_elapsed = 0.0
+    optimizer_state_restored = False
+    resume_source: str | None = None
+    if args.resume:
+        (
+            completed_steps,
+            losses,
+            previous_elapsed,
+            optimizer_state_restored,
+            resume_source,
+        ) = _load_resume_state(output_dir, expert, optimizer, architecture, device)
+        if completed_steps > target_steps:
+            raise ValueError(
+                f"Checkpoint already has {completed_steps} steps, beyond target {target_steps}"
+            )
+        mode = "exact checkpoint resume" if optimizer_state_restored else "legacy warm start"
+        print(
+            f"stage2 {mode}: source={resume_source} completed={completed_steps} "
+            f"remaining={target_steps - completed_steps}"
+        )
+
+    new_losses: list[float] = []
+    started = time.perf_counter()
+    expert.train()
+    progress = trange(
+        completed_steps,
+        target_steps,
+        desc="stage2-expert",
+        initial=completed_steps,
+        total=target_steps,
+    )
+    for step in progress:
+        sample_index = step % len(samples)
+        target = target_actions[sample_index : sample_index + 1]
+        training_data = flow.construct_training_data(target)
+        optimizer.zero_grad(set_to_none=True)
+        prediction = expert(
+            noisy_action=training_data["noisy_action"],
+            timestep=training_data["timestep"],
+            prefix=prefixes[sample_index],
+        )
+        loss = flow.loss(training_data, prediction)
+        if not torch.isfinite(loss):
+            raise RuntimeError(f"Non-finite Stage-2 loss at step {step}: {loss}")
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(expert.parameters(), 1.0)
+        optimizer.step()
+        value = float(loss.detach().cpu())
+        losses.append(value)
+        new_losses.append(value)
+        if step == 0 or (step + 1) % stage2.log_every == 0:
+            print(f"stage2 step={step + 1} loss={value:.6f}")
+        if (step + 1) % stage2.checkpoint_every == 0:
+            elapsed_so_far = previous_elapsed + (time.perf_counter() - started)
+            _save_training_checkpoint(
+                checkpoint_dir,
+                expert,
+                optimizer,
+                architecture,
+                step + 1,
+                losses,
+                elapsed_so_far,
+                device,
+            )
+    elapsed_this_run = time.perf_counter() - started
+    elapsed = previous_elapsed + elapsed_this_run
+
+    _save_training_checkpoint(
+        checkpoint_dir,
+        expert,
+        optimizer,
+        architecture,
+        target_steps,
+        losses,
+        elapsed,
+        device,
+    )
+    _atomic_save_safetensors(expert, output_dir / "expert.safetensors")
+    write_json(
+        output_dir / "expert_config.json",
+        {
+            "architecture": asdict(architecture),
+            "action_mean": action_space.mean.detach().cpu().tolist(),
+            "action_scale": action_space.scale.detach().cpu().tolist(),
+            "action_shape": [cfg.data.expert_waypoints, 2],
+            "inference_steps": stage2.inference_steps,
+            "stage1_checkpoint": str(stage1_dir),
+            "completed_steps": target_steps,
+        },
+    )
+
+    first = samples[0]
+    first_history, first_future = first.load_trajectory()
+    generated_reasoning = generate_reasoning(
+        vlm,
+        processor,
+        first.load_image(),
+        build_vlm_prefix(first_history),
+        device,
+        dtype,
+        cfg.vlm.max_reasoning_tokens,
+    )
+    generated_prefix = extract_prefix_kv(
+        vlm,
+        processor,
+        first.load_image(),
+        build_expert_prefix(first_history, generated_reasoning),
+        device,
+        dtype,
+    )
+    expert.eval()
+    seed_everything(cfg.seed + 1)
+    predicted_action = flow.sample(
+        batch_size=1,
+        step_fn=partial(expert, prefix=generated_prefix),
+        device=device,
+        dtype=dtype,
+    ).float()
+    predicted_xy = action_space.action_to_trajectory(
+        predicted_action, first_history[None, :, :2].to(device)
+    )
+    target_xy = first_future[None, :, :2].to(device)
+    target_action = action_space.trajectory_to_action(
+        first_history[None, :, :2].to(device), target_xy
+    )
+    metrics = {
+        "stage": "stage2_expert",
+        "device": str(device),
+        "dtype": str(dtype),
+        "vlm_parameters": vlm_counts,
+        "expert_parameters": expert_counts,
+        "architecture": asdict(architecture),
+        "steps": target_steps,
+        "additional_steps": len(new_losses),
+        "resumed_from_step": completed_steps,
+        "resume_source": resume_source,
+        "optimizer_state_restored": optimizer_state_restored,
+        "learning_rate": stage2.learning_rate,
+        "elapsed_seconds": elapsed,
+        "elapsed_seconds_this_run": elapsed_this_run,
+        "initial_loss": losses[0],
+        "resume_initial_loss": new_losses[0] if new_losses else None,
+        "final_loss": losses[-1],
+        "minimum_loss": min(losses),
+        "losses": losses,
+        "generated_reasoning": generated_reasoning,
+        "target_reasoning": first.reasoning,
+        "action_mse": float(
+            torch.nn.functional.mse_loss(predicted_action, target_action.float()).cpu()
+        ),
+        **trajectory_metrics(predicted_xy, target_xy),
+    }
+    write_json(output_dir / "stage2_metrics.json", metrics)
+    save_prediction_arrays(
+        output_dir / "prediction.npz",
+        first_history[:, :2],
+        target_action[0],
+        predicted_action[0],
+        target_xy[0],
+        predicted_xy[0],
+    )
+    save_trajectory_plot(
+        output_dir / "trajectory.png",
+        first_history[:, :2],
+        target_xy[0],
+        predicted_xy[0],
+    )
+    print(metrics)
+    del optimizer, expert, vlm
+    empty_device_cache(device)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_vlm_trajectory_logits.py

@@ -0,0 +1,662 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from transformers import LogitsProcessorList, StoppingCriteriaList
+from tqdm import trange
+
+from alpamayo_lite.action_space import RelativeDeltaActionSpace
+from alpamayo_lite.config import load_config
+from alpamayo_lite.data import load_manifest
+from alpamayo_lite.evaluation import save_trajectory_plot, trajectory_metrics
+from alpamayo_lite.expert import PrefixKV
+from alpamayo_lite.prompts import build_vlm_prefix
+from alpamayo_lite.runtime import (
+    empty_device_cache,
+    parameter_counts,
+    resolve_dtype,
+    seed_everything,
+    select_device,
+    write_json,
+)
+from alpamayo_lite.trajectory_tokens import (
+    DiscreteTrajectoryMask,
+    StopOneTokenAfterBoundary,
+    TrajectoryVocabulary,
+    add_trajectory_vocabulary,
+    build_structured_target,
+    dequantize_action,
+    extract_trajectory_values,
+    mask_after_boundary,
+    quantize_action,
+)
+from alpamayo_lite.vlm import (
+    configure_stage1_trainability,
+    enforce_frozen_module_eval,
+    load_vlm,
+    prepare_stage1_batch,
+    process_image_text,
+)
+
+
+def _resolve(root: Path, value: str) -> Path:
+    path = Path(value).expanduser()
+    return path.resolve() if path.is_absolute() else (root / path).resolve()
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _infer_meta_action(reasoning: str) -> str:
+    lowered = reasoning.lower()
+    for phrase, label in (
+        ("steer left", "STEER_LEFT"),
+        ("turn left", "TURN_LEFT"),
+        ("steer right", "STEER_RIGHT"),
+        ("turn right", "TURN_RIGHT"),
+        ("stop", "STOP"),
+        ("slow", "SLOW_DOWN"),
+    ):
+        if phrase in lowered:
+            return label
+    return "CONTINUE"
+
+
+def _vocabulary_from_tokenizer(tokenizer: Any, size: int) -> TrajectoryVocabulary:
+    start_id = int(tokenizer.convert_tokens_to_ids("<i0>"))
+    end_id = int(tokenizer.convert_tokens_to_ids(f"<i{size - 1}>"))
+    if end_id - start_id + 1 != size:
+        raise RuntimeError("Reloaded trajectory vocabulary is not contiguous")
+    special_tokens = (
+        "<|cot_start|>",
+        "<|cot_end|>",
+        "<|meta_action_start|>",
+        "<|meta_action_end|>",
+        "<|traj_future_start|>",
+        "<|traj_future_end|>",
+    )
+    special_token_ids = {
+        token: int(tokenizer.convert_tokens_to_ids(token)) for token in special_tokens
+    }
+    return TrajectoryVocabulary(size, start_id, end_id, special_token_ids)
+
+
+def _cache_summary(cache: object) -> dict[str, Any]:
+    prefix = PrefixKV.from_transformers_cache(cache)
+    first_key = prefix.keys[0]
+    sample = first_key[..., : min(8, first_key.shape[-2]), : min(8, first_key.shape[-1])]
+    return {
+        "layers": prefix.num_layers,
+        "sequence_length": prefix.sequence_length,
+        "first_key_shape": list(first_key.shape),
+        "first_key_sample_sum": float(sample.float().sum().cpu()),
+        "all_finite": bool(
+            all(torch.isfinite(key).all() and torch.isfinite(value).all() for key, value in zip(prefix.keys, prefix.values))
+        ),
+    }
+
+
+def _top_tokens(logits: torch.Tensor, tokenizer: Any, count: int = 5) -> list[dict[str, Any]]:
+    values, ids = torch.topk(logits.float(), k=count)
+    return [
+        {
+            "id": int(token_id),
+            "token": tokenizer.convert_ids_to_tokens(int(token_id)),
+            "logit": float(value),
+        }
+        for value, token_id in zip(values.cpu(), ids.cpu())
+    ]
+
+
+@torch.no_grad()
+def _generate(
+    model: Any,
+    processor: Any,
+    image: Any,
+    prefix_text: str,
+    device: torch.device,
+    dtype: torch.dtype,
+    max_new_tokens: int,
+    logits_processor: LogitsProcessorList | None = None,
+    stopping_criteria: StoppingCriteriaList | None = None,
+) -> tuple[dict[str, Any], torch.Tensor, Any]:
+    inputs = process_image_text(processor, image, prefix_text, device, dtype)
+    prompt_length = inputs["input_ids"].shape[1]
+    generated = model.generate(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        do_sample=False,
+        use_cache=True,
+        return_dict_in_generate=True,
+        output_logits=True,
+        output_scores=True,
+        logits_processor=logits_processor,
+        stopping_criteria=stopping_criteria,
+        pad_token_id=processor.tokenizer.pad_token_id,
+        eos_token_id=processor.tokenizer.eos_token_id,
+    )
+    new_ids = generated.sequences[:, prompt_length:]
+    step_records = []
+    returned_logits = generated.logits or ()
+    for index, token_id in enumerate(new_ids[0].tolist()):
+        if index >= len(returned_logits):
+            break
+        logits = returned_logits[index][0]
+        log_probability = torch.log_softmax(logits.float(), dim=-1)[token_id]
+        step_records.append(
+            {
+                "index": index,
+                "id": int(token_id),
+                "token": processor.tokenizer.convert_ids_to_tokens(int(token_id)),
+                "selected_logit": float(logits[token_id].float().cpu()),
+                "selected_log_probability": float(log_probability.cpu()),
+                "top5": _top_tokens(logits, processor.tokenizer),
+            }
+        )
+    result = {
+        "prompt_length": prompt_length,
+        "new_token_ids": new_ids[0].detach().cpu().tolist(),
+        "raw_text": processor.tokenizer.decode(new_ids[0], skip_special_tokens=False),
+        "clean_text": processor.tokenizer.decode(new_ids[0], skip_special_tokens=True),
+        "steps": step_records,
+        "cache": _cache_summary(generated.past_key_values),
+    }
+    return result, new_ids, generated
+
+
+@torch.no_grad()
+def _teacher_forcing_probe(
+    model: Any,
+    inputs: dict[str, torch.Tensor],
+    labels: torch.Tensor,
+    vocabulary: TrajectoryVocabulary,
+    tokenizer: Any,
+) -> tuple[dict[str, Any], np.ndarray, np.ndarray]:
+    outputs = model(**inputs, use_cache=False, return_dict=True)
+    logits = outputs.logits[:, :-1]
+    next_ids = inputs["input_ids"][:, 1:]
+    supervised = labels[:, 1:] != -100
+    trajectory = (
+        (next_ids >= vocabulary.start_id)
+        & (next_ids <= vocabulary.end_id)
+        & supervised
+    )
+    supervised_logits = logits[supervised].float()
+    supervised_ids = next_ids[supervised]
+    trajectory_logits = logits[trajectory].float()
+    trajectory_ids = next_ids[trajectory]
+    if trajectory_logits.shape[0] == 0:
+        raise RuntimeError("No supervised trajectory tokens were found")
+
+    target_logits = trajectory_logits.gather(1, trajectory_ids[:, None]).squeeze(1)
+    ranks = (trajectory_logits > target_logits[:, None]).sum(dim=1) + 1
+    probabilities = torch.softmax(trajectory_logits, dim=-1).gather(
+        1, trajectory_ids[:, None]
+    ).squeeze(1)
+    top_values, top_ids = torch.topk(trajectory_logits, k=5, dim=-1)
+    records = []
+    for index in range(trajectory_logits.shape[0]):
+        records.append(
+            {
+                "index": index,
+                "target_id": int(trajectory_ids[index]),
+                "target_token": tokenizer.convert_ids_to_tokens(int(trajectory_ids[index])),
+                "target_logit": float(target_logits[index].cpu()),
+                "target_probability": float(probabilities[index].cpu()),
+                "target_rank": int(ranks[index].cpu()),
+                "top5": [
+                    {
+                        "id": int(token_id),
+                        "token": tokenizer.convert_ids_to_tokens(int(token_id)),
+                        "logit": float(value),
+                    }
+                    for value, token_id in zip(top_values[index].cpu(), top_ids[index].cpu())
+                ],
+            }
+        )
+
+    discrete_slice = trajectory_logits[:, vocabulary.start_id : vocabulary.end_id + 1]
+    result = {
+        "supervised_token_count": int(supervised.sum()),
+        "trajectory_supervised_token_count": int(trajectory.sum()),
+        "next_token_accuracy": float(
+            (supervised_logits.argmax(dim=-1) == supervised_ids).float().mean().cpu()
+        ),
+        "trajectory_next_token_accuracy": float(
+            (trajectory_logits.argmax(dim=-1) == trajectory_ids).float().mean().cpu()
+        ),
+        "trajectory_mean_target_probability": float(probabilities.mean().cpu()),
+        "trajectory_max_target_rank": int(ranks.max().cpu()),
+        "trajectory_records": records,
+    }
+    return (
+        result,
+        discrete_slice.detach().cpu().numpy(),
+        (trajectory_ids - vocabulary.start_id).detach().cpu().numpy(),
+    )
+
+
+@torch.no_grad()
+def _handoff_probe(
+    model: Any,
+    processor: Any,
+    image: Any,
+    handoff_text: str,
+    first_target_id: int,
+    vocabulary: TrajectoryVocabulary,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[dict[str, Any], np.ndarray, np.ndarray]:
+    inputs = process_image_text(processor, image, handoff_text, device, dtype)
+    outputs = model(
+        **inputs,
+        use_cache=True,
+        return_dict=True,
+        logits_to_keep=1,
+    )
+    raw = outputs.logits[0, -1].float()
+    masked = DiscreteTrajectoryMask(vocabulary.start_id, vocabulary.size)(
+        inputs["input_ids"], raw[None].clone()
+    )[0]
+    target_logit = raw[first_target_id]
+    target_rank = int((raw > target_logit).sum().cpu()) + 1
+    outside = torch.ones_like(raw, dtype=torch.bool)
+    outside[vocabulary.start_id : vocabulary.end_id + 1] = False
+    result = {
+        "first_target_id": first_target_id,
+        "first_target_token": processor.tokenizer.convert_ids_to_tokens(first_target_id),
+        "raw_first_target_logit": float(target_logit.cpu()),
+        "raw_first_target_rank": target_rank,
+        "raw_top5": _top_tokens(raw, processor.tokenizer),
+        "processed_top5": _top_tokens(masked, processor.tokenizer),
+        "raw_trajectory_finite_count": int(
+            torch.isfinite(raw[vocabulary.start_id : vocabulary.end_id + 1]).sum().cpu()
+        ),
+        "processed_trajectory_negative_infinity_count": int(
+            torch.isneginf(masked[vocabulary.start_id : vocabulary.end_id + 1]).sum().cpu()
+        ),
+        "outside_range_max_absolute_change": float((masked[outside] - raw[outside]).abs().max().cpu()),
+        "cache": _cache_summary(outputs.past_key_values),
+    }
+    return (
+        result,
+        raw[vocabulary.start_id : vocabulary.end_id + 1].detach().cpu().numpy(),
+        masked[vocabulary.start_id : vocabulary.end_id + 1].detach().cpu().numpy(),
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/local_overfit.yaml")
+    parser.add_argument("--output-dir", default="outputs/stage1_vlm_trajectory")
+    parser.add_argument("--steps", type=int, default=120)
+    parser.add_argument("--learning-rate", type=float, default=1e-4)
+    parser.add_argument("--trajectory-vocab-size", type=int, default=128)
+    parser.add_argument("--action-min", type=float, default=-4.0)
+    parser.add_argument("--action-max", type=float, default=4.0)
+    parser.add_argument("--max-new-tokens", type=int, default=96)
+    parser.add_argument("--early-stop-loss", type=float, default=1e-4)
+    parser.add_argument("--early-stop-patience", type=int, default=5)
+    args = parser.parse_args()
+    if args.steps <= 0:
+        raise ValueError("--steps must be positive")
+    if args.early_stop_loss < 0:
+        raise ValueError("--early-stop-loss must be non-negative")
+    if args.early_stop_patience <= 0:
+        raise ValueError("--early-stop-patience must be positive")
+
+    cfg = load_config(args.config)
+    seed_everything(cfg.seed)
+    device = select_device()
+    dtype = resolve_dtype(cfg.vlm.dtype, device)
+    samples = load_manifest(cfg.resolve(cfg.paths.manifest))
+    if len(samples) != 1:
+        raise ValueError("This local trajectory-logit validation expects exactly one sample")
+    sample = samples[0]
+    history, future = sample.load_trajectory()
+    image = sample.load_image()
+
+    processor, model = load_vlm(cfg.resolve(cfg.paths.stage1_dir), cfg.vlm, device, dtype)
+    tokenizer = processor.tokenizer
+    original_tokenizer_length = len(tokenizer)
+    vocabulary = add_trajectory_vocabulary(tokenizer, args.trajectory_vocab_size)
+    model_vocab_size = model.get_output_embeddings().weight.shape[0]
+    if len(tokenizer) > model_vocab_size:
+        raise RuntimeError(
+            f"Expanded tokenizer ({len(tokenizer)}) exceeds preallocated LM head ({model_vocab_size})"
+        )
+    if model.get_input_embeddings().weight.data_ptr() != model.get_output_embeddings().weight.data_ptr():
+        raise RuntimeError("Trajectory validation expects tied input/output embeddings")
+
+    expert_config = _read_json(cfg.resolve(cfg.paths.stage2_dir) / "expert_config.json")
+    action_space = RelativeDeltaActionSpace(
+        torch.tensor(expert_config["action_mean"], device=device),
+        torch.tensor(expert_config["action_scale"], device=device),
+    ).to(device)
+    target_xy = future[None, :, :2].to(device)
+    history_xy = history[None, :, :2].to(device)
+    target_action = action_space.trajectory_to_action(history_xy, target_xy)
+    quantized = quantize_action(
+        target_action,
+        args.trajectory_vocab_size,
+        args.action_min,
+        args.action_max,
+    )
+    meta_action = _infer_meta_action(sample.reasoning)
+    structured_target = build_structured_target(sample.reasoning, meta_action, quantized)
+    prefix_text = build_vlm_prefix(history)
+    full_text = prefix_text + structured_target
+
+    baseline, _, baseline_output = _generate(
+        model,
+        processor,
+        image,
+        prefix_text,
+        device,
+        dtype,
+        args.max_new_tokens,
+    )
+    del baseline_output
+
+    trainability = configure_stage1_trainability(model, cfg.stage1)
+    embedding_weight = model.get_input_embeddings().weight
+    embedding_weight.requires_grad_(True)
+    trainable_rows = list(range(vocabulary.start_id, vocabulary.end_id + 1)) + list(
+        vocabulary.special_token_ids.values()
+    )
+    gradient_mask = torch.zeros(
+        embedding_weight.shape[0], 1, device=embedding_weight.device, dtype=embedding_weight.dtype
+    )
+    gradient_mask[trainable_rows] = 1
+    gradient_hook = embedding_weight.register_hook(lambda gradient: gradient * gradient_mask)
+
+    inputs, labels, prefix_length = prepare_stage1_batch(
+        processor,
+        image,
+        prefix_text,
+        full_text,
+        device,
+        dtype,
+    )
+    target_ids = inputs["input_ids"][:, prefix_length:]
+    trajectory_supervision_count = int(
+        (
+            (target_ids >= vocabulary.start_id)
+            & (target_ids <= vocabulary.end_id)
+        ).sum().cpu()
+    )
+    if trajectory_supervision_count != quantized.numel():
+        raise RuntimeError(
+            f"Expected {quantized.numel()} trajectory labels, found {trajectory_supervision_count}"
+        )
+
+    optimizer = torch.optim.AdamW(
+        [parameter for parameter in model.parameters() if parameter.requires_grad],
+        lr=args.learning_rate,
+        weight_decay=cfg.stage1.weight_decay,
+        eps=cfg.stage1.adam_epsilon,
+    )
+    losses: list[float] = []
+    accuracies: list[float] = []
+    converged_steps = 0
+    started = time.perf_counter()
+    model.train()
+    enforce_frozen_module_eval(model)
+    for step in trange(args.steps, desc="vlm-trajectory-logits"):
+        optimizer.zero_grad(set_to_none=True)
+        outputs = model(**inputs, labels=labels, use_cache=False, return_dict=True)
+        loss = outputs.loss
+        if loss is None or not torch.isfinite(loss):
+            raise RuntimeError(f"Non-finite trajectory-logit loss at step {step}: {loss}")
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(
+            [parameter for parameter in model.parameters() if parameter.requires_grad], 1.0
+        )
+        optimizer.step()
+        value = float(loss.detach().cpu())
+        next_predictions = outputs.logits[:, :-1].argmax(dim=-1)
+        shifted_labels = labels[:, 1:]
+        supervised = shifted_labels != -100
+        accuracy = float(
+            (next_predictions[supervised] == shifted_labels[supervised]).float().mean().cpu()
+        )
+        losses.append(value)
+        accuracies.append(accuracy)
+        if step == 0 or (step + 1) % cfg.stage1.log_every == 0:
+            print(f"trajectory-vlm step={step + 1} loss={value:.6f} accuracy={accuracy:.4f}")
+        del outputs, next_predictions
+        if value <= args.early_stop_loss and accuracy == 1.0:
+            converged_steps += 1
+        else:
+            converged_steps = 0
+        if converged_steps >= args.early_stop_patience:
+            print(
+                f"trajectory-vlm converged at step={step + 1}: "
+                f"loss={value:.6f} accuracy={accuracy:.4f}"
+            )
+            break
+    elapsed = time.perf_counter() - started
+
+    model.gradient_checkpointing_disable()
+    model.config.use_cache = True
+    model.model.language_model.config.use_cache = True
+    output_dir = _resolve(cfg.root, args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(output_dir, safe_serialization=True, max_shard_size="2GB")
+    processor.save_pretrained(output_dir)
+    write_json(
+        output_dir / "trajectory_vocabulary.json",
+        {
+            **asdict(vocabulary),
+            "original_tokenizer_length": original_tokenizer_length,
+            "expanded_tokenizer_length": len(tokenizer),
+            "model_vocab_size": model_vocab_size,
+            "preallocated_rows_used": True,
+            "action_min": args.action_min,
+            "action_max": args.action_max,
+            "action_mean": expert_config["action_mean"],
+            "action_scale": expert_config["action_scale"],
+            "action_shape": expert_config["action_shape"],
+        },
+    )
+    gradient_hook.remove()
+    del optimizer, model, processor
+    empty_device_cache(device)
+
+    processor, model = load_vlm(output_dir, cfg.vlm, device, dtype)
+    model.requires_grad_(False)
+    model.eval()
+    vocabulary = _vocabulary_from_tokenizer(processor.tokenizer, args.trajectory_vocab_size)
+    inputs, labels, prefix_length = prepare_stage1_batch(
+        processor,
+        image,
+        prefix_text,
+        full_text,
+        device,
+        dtype,
+    )
+    teacher, teacher_logits, teacher_values = _teacher_forcing_probe(
+        model, inputs, labels, vocabulary, processor.tokenizer
+    )
+
+    cot_and_meta = (
+        f"<|cot_start|>{sample.reasoning}<|cot_end|>\n"
+        f"<|meta_action_start|>{meta_action}<|meta_action_end|>\n"
+    )
+    handoff_text = prefix_text + cot_and_meta + "<|traj_future_start|>"
+    first_target_id = vocabulary.start_id + int(quantized.reshape(-1)[0].cpu())
+    handoff, boundary_raw, boundary_masked = _handoff_probe(
+        model,
+        processor,
+        image,
+        handoff_text,
+        first_target_id,
+        vocabulary,
+        device,
+        dtype,
+    )
+
+    full_generation, full_ids, full_output = _generate(
+        model,
+        processor,
+        image,
+        prefix_text,
+        device,
+        dtype,
+        args.max_new_tokens,
+    )
+    generated_values = extract_trajectory_values(full_ids[0], vocabulary)
+    expected_values = quantized.detach().cpu().reshape(-1)
+    generated_token_accuracy = (
+        float((generated_values.cpu() == expected_values).float().mean())
+        if generated_values.numel() == expected_values.numel()
+        else 0.0
+    )
+    if generated_values.numel() != target_action.numel():
+        raise RuntimeError(
+            f"Generated {generated_values.numel()} trajectory tokens, expected {target_action.numel()}"
+        )
+    generated_action = dequantize_action(
+        generated_values.reshape_as(target_action).to(device),
+        args.trajectory_vocab_size,
+        args.action_min,
+        args.action_max,
+    )
+    generated_xy = action_space.action_to_trajectory(generated_action, history_xy)
+    quantized_target_action = dequantize_action(
+        quantized,
+        args.trajectory_vocab_size,
+        args.action_min,
+        args.action_max,
+    )
+    quantized_target_xy = action_space.action_to_trajectory(quantized_target_action, history_xy)
+    del full_output
+
+    boundary_id = vocabulary.special_token_ids["<|traj_future_start|>"]
+    source_mask = DiscreteTrajectoryMask(vocabulary.start_id, vocabulary.size)
+    source_stop = StopOneTokenAfterBoundary(boundary_id)
+    truncated_generation, truncated_ids, truncated_output = _generate(
+        model,
+        processor,
+        image,
+        prefix_text,
+        device,
+        dtype,
+        args.max_new_tokens,
+        logits_processor=LogitsProcessorList([source_mask]),
+        stopping_criteria=StoppingCriteriaList([source_stop]),
+    )
+    visible_ids = mask_after_boundary(
+        truncated_output.sequences,
+        boundary_id,
+        processor.tokenizer.pad_token_id,
+    )
+    truncated_generation["visible_new_token_ids"] = visible_ids[
+        :, truncated_generation["prompt_length"] :
+    ][0].cpu().tolist()
+    truncated_generation["visible_raw_text"] = processor.tokenizer.decode(
+        visible_ids[:, truncated_generation["prompt_length"] :][0],
+        skip_special_tokens=False,
+    )
+    truncated_generation["boundary_generated"] = bool((truncated_ids == boundary_id).any())
+    truncated_generation["trajectory_tokens_visible"] = int(
+        (
+            (visible_ids >= vocabulary.start_id)
+            & (visible_ids <= vocabulary.end_id)
+        ).sum().cpu()
+    )
+    boundary_positions = torch.nonzero(truncated_ids[0] == boundary_id, as_tuple=False).flatten()
+    post_boundary_logit_index = (
+        int(boundary_positions[0]) + 1 if boundary_positions.numel() else -1
+    )
+    if 0 <= post_boundary_logit_index < len(truncated_output.logits or ()):
+        post_boundary_logits = truncated_output.logits[post_boundary_logit_index][0]
+        truncated_generation["returned_post_boundary_trajectory_negative_infinity_count"] = int(
+            torch.isneginf(
+                post_boundary_logits[vocabulary.start_id : vocabulary.end_id + 1]
+            ).sum().cpu()
+        )
+    del truncated_output
+
+    save_trajectory_plot(
+        output_dir / "vlm_discrete_trajectory.png",
+        history[:, :2],
+        target_xy[0],
+        generated_xy[0],
+        prediction_label="VLM discrete tokens",
+        title="VLM discrete-trajectory logits",
+    )
+    np.savez_compressed(
+        output_dir / "trajectory_logits.npz",
+        teacher_trajectory_logits=teacher_logits,
+        teacher_target_values=teacher_values,
+        boundary_raw_trajectory_logits=boundary_raw,
+        boundary_masked_trajectory_logits=boundary_masked,
+        target_action=target_action.detach().cpu().numpy(),
+        quantized_target_values=quantized.detach().cpu().numpy(),
+        generated_values=generated_values.detach().cpu().numpy(),
+        generated_action=generated_action.detach().cpu().numpy(),
+        target_xy=target_xy.detach().cpu().numpy(),
+        generated_xy=generated_xy.detach().cpu().numpy(),
+    )
+
+    metrics = {
+        "stage": "vlm_discrete_trajectory_logits",
+        "device": str(device),
+        "dtype": str(dtype),
+        "model_parameters": parameter_counts(model),
+        "vocabulary": asdict(vocabulary),
+        "original_tokenizer_length": original_tokenizer_length,
+        "expanded_tokenizer_length": len(processor.tokenizer),
+        "model_vocab_size": model_vocab_size,
+        "semantic_trajectory_supervision_count": trajectory_supervision_count,
+        "structured_target": structured_target,
+        "target_meta_action": meta_action,
+        "training": {
+            "steps": len(losses),
+            "requested_steps": args.steps,
+            "learning_rate": args.learning_rate,
+            "elapsed_seconds": elapsed,
+            "initial_loss": losses[0],
+            "final_loss": losses[-1],
+            "minimum_loss": min(losses),
+            "final_training_accuracy": accuracies[-1],
+            "losses": losses,
+            "accuracies": accuracies,
+            "base_trainability": trainability,
+            "trainable_token_rows": len(trainable_rows),
+        },
+        "baseline_generation_before_structured_training": baseline,
+        "teacher_forcing": teacher,
+        "raw_handoff_logits_and_cache": handoff,
+        "untruncated_vlm_generation": {
+            **full_generation,
+            "trajectory_token_count": int(generated_values.numel()),
+            "trajectory_token_accuracy": generated_token_accuracy,
+            "generated_values": generated_values.cpu().tolist(),
+            "target_values": expected_values.tolist(),
+            "action_mse": float(
+                torch.nn.functional.mse_loss(generated_action, target_action).cpu()
+            ),
+            **trajectory_metrics(generated_xy, target_xy),
+        },
+        "quantization_only": trajectory_metrics(quantized_target_xy, target_xy),
+        "source_style_diffusion_handoff_generation": truncated_generation,
+    }
+    write_json(output_dir / "trajectory_logits_metrics.json", metrics)
+    print(json.dumps(metrics, indent=2, ensure_ascii=False))
+    del model, processor
+    empty_device_cache(device)
+
+
+if __name__ == "__main__":
+    main()

scripts/verify_connection.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+from dataclasses import asdict
+
+import torch
+
+from alpamayo_lite.config import load_config
+from alpamayo_lite.data import load_manifest
+from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert
+from alpamayo_lite.prompts import build_expert_prefix, build_vlm_prefix
+from alpamayo_lite.runtime import (
+    empty_device_cache,
+    parameter_counts,
+    resolve_dtype,
+    seed_everything,
+    select_device,
+    write_json,
+)
+from alpamayo_lite.vlm import (
+    configure_stage1_trainability,
+    extract_prefix_kv,
+    generate_reasoning,
+    load_vlm,
+    vlm_text_config,
+)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/local_overfit.yaml")
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    seed_everything(cfg.seed)
+    device = select_device()
+    dtype = resolve_dtype(cfg.vlm.dtype, device)
+    sample = load_manifest(cfg.resolve(cfg.paths.manifest))[0]
+    history, _ = sample.load_trajectory()
+
+    processor, vlm = load_vlm(cfg.resolve(cfg.paths.model_dir), cfg.vlm, device, dtype)
+    vlm.eval()
+    baseline_text = generate_reasoning(
+        vlm,
+        processor,
+        sample.load_image(),
+        build_vlm_prefix(history),
+        device,
+        dtype,
+        cfg.vlm.max_reasoning_tokens,
+    )
+    prefix = extract_prefix_kv(
+        vlm,
+        processor,
+        sample.load_image(),
+        build_expert_prefix(history, sample.reasoning),
+        device,
+        dtype,
+    )
+
+    stage2 = cfg.stage2
+    architecture = ExpertArchitecture(
+        hidden_size=stage2.expert_hidden_size,
+        intermediate_size=stage2.expert_intermediate_size,
+        num_layers=stage2.expert_num_layers,
+        num_attention_heads=stage2.expert_num_attention_heads,
+        num_key_value_heads=stage2.expert_num_key_value_heads,
+        head_dim=stage2.expert_head_dim,
+        action_projection_hidden_size=stage2.action_projection_hidden_size,
+        action_fourier_features=stage2.action_fourier_features,
+        non_causal_attention=stage2.expert_non_causal_attention,
+    )
+    expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
+        device=device, dtype=dtype
+    )
+    with torch.no_grad():
+        predicted_velocity = expert(
+            noisy_action=torch.randn(1, cfg.data.expert_waypoints, 2, device=device, dtype=dtype),
+            timestep=torch.full((1, 1, 1), 0.5, device=device, dtype=dtype),
+            prefix=prefix,
+        )
+
+    trainability = configure_stage1_trainability(vlm, cfg.stage1)
+    result = {
+        "device": str(device),
+        "dtype": str(dtype),
+        "vlm_parameters": parameter_counts(vlm),
+        "stage1_trainability": trainability,
+        "expert_parameters": parameter_counts(expert),
+        "expert_architecture": asdict(architecture),
+        "prefix_layers": prefix.num_layers,
+        "prefix_sequence_length": prefix.sequence_length,
+        "prefix_key_shape": list(prefix.keys[0].shape),
+        "expert_output_shape": list(predicted_velocity.shape),
+        "expert_output_finite": bool(torch.isfinite(predicted_velocity).all().item()),
+        "baseline_text": baseline_text,
+        "target_text": sample.reasoning,
+    }
+    write_json(cfg.root / "outputs" / "connection_test.json", result)
+    print(result)
+    del expert, vlm
+    empty_device_cache(device)
+
+
+if __name__ == "__main__":
+    main()

scripts/verify_structured_vlm_handoff.py

@@ -0,0 +1,226 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+from dataclasses import asdict
+from functools import partial
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import load_file
+from transformers import LogitsProcessorList, StoppingCriteriaList
+
+from alpamayo_lite.action_space import RelativeDeltaActionSpace
+from alpamayo_lite.config import load_config
+from alpamayo_lite.data import load_manifest
+from alpamayo_lite.diffusion import FlowMatching
+from alpamayo_lite.evaluation import (
+    save_prediction_arrays,
+    save_trajectory_plot,
+    trajectory_metrics,
+)
+from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert, PrefixKV
+from alpamayo_lite.prompts import build_vlm_prefix
+from alpamayo_lite.runtime import (
+    empty_device_cache,
+    parameter_counts,
+    resolve_dtype,
+    seed_everything,
+    select_device,
+    write_json,
+)
+from alpamayo_lite.trajectory_tokens import (
+    DiscreteTrajectoryMask,
+    StopOneTokenAfterBoundary,
+    mask_after_boundary,
+)
+from alpamayo_lite.vlm import load_vlm, process_image_text, vlm_text_config
+
+
+def _resolve(root: Path, value: str) -> Path:
+    path = Path(value).expanduser()
+    return path.resolve() if path.is_absolute() else (root / path).resolve()
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="configs/local_overfit.yaml")
+    parser.add_argument("--vlm-dir", default="outputs/stage1_vlm_trajectory")
+    parser.add_argument("--stage2-dir", default="outputs/stage2_expert")
+    parser.add_argument("--output-dir", default="outputs/structured_vlm_handoff")
+    parser.add_argument("--max-new-tokens", type=int, default=96)
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    seed_everything(cfg.seed)
+    device = select_device()
+    dtype = resolve_dtype(cfg.vlm.dtype, device)
+    vlm_dir = _resolve(cfg.root, args.vlm_dir)
+    stage2_dir = _resolve(cfg.root, args.stage2_dir)
+    output_dir = _resolve(cfg.root, args.output_dir)
+    sample = load_manifest(cfg.resolve(cfg.paths.manifest))[0]
+    history, future = sample.load_trajectory()
+
+    vocabulary = _read_json(vlm_dir / "trajectory_vocabulary.json")
+    trajectory_start = int(vocabulary["start_id"])
+    trajectory_size = int(vocabulary["size"])
+    boundary_id = int(vocabulary["special_token_ids"]["<|traj_future_start|>"])
+
+    processor, vlm = load_vlm(vlm_dir, cfg.vlm, device, dtype)
+    vlm.requires_grad_(False)
+    vlm.eval()
+    inputs = process_image_text(
+        processor,
+        sample.load_image(),
+        build_vlm_prefix(history),
+        device,
+        dtype,
+    )
+    prompt_length = inputs["input_ids"].shape[1]
+    generated = vlm.generate(
+        **inputs,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        use_cache=True,
+        return_dict_in_generate=True,
+        output_logits=True,
+        logits_processor=LogitsProcessorList(
+            [DiscreteTrajectoryMask(trajectory_start, trajectory_size)]
+        ),
+        stopping_criteria=StoppingCriteriaList([StopOneTokenAfterBoundary(boundary_id)]),
+        pad_token_id=processor.tokenizer.pad_token_id,
+        eos_token_id=processor.tokenizer.eos_token_id,
+    )
+    new_ids = generated.sequences[:, prompt_length:]
+    visible_sequences = mask_after_boundary(
+        generated.sequences,
+        boundary_id,
+        processor.tokenizer.pad_token_id,
+    )
+    visible_new_ids = visible_sequences[:, prompt_length:]
+    prefix = PrefixKV.from_transformers_cache(generated.past_key_values)
+    boundary_positions = torch.nonzero(new_ids[0] == boundary_id, as_tuple=False).flatten()
+    if boundary_positions.numel() == 0:
+        raise RuntimeError("Structured VLM did not generate the diffusion handoff boundary")
+    boundary_index = int(boundary_positions[0])
+    expected_prefix_length = prompt_length + boundary_index + 1
+    if prefix.sequence_length != expected_prefix_length:
+        raise RuntimeError(
+            f"Generation cache does not end at the handoff boundary: "
+            f"cache={prefix.sequence_length}, expected={expected_prefix_length}"
+        )
+
+    expert_config = _read_json(stage2_dir / "expert_config.json")
+    architecture = ExpertArchitecture(**expert_config["architecture"])
+    expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
+        device=device, dtype=dtype
+    )
+    expert.load_state_dict(load_file(stage2_dir / "expert.safetensors"), strict=True)
+    expert.requires_grad_(False)
+    expert.eval()
+    flow = FlowMatching(
+        action_shape=tuple(expert_config["action_shape"]),
+        inference_steps=int(expert_config["inference_steps"]),
+    )
+    action_space = RelativeDeltaActionSpace(
+        torch.tensor(expert_config["action_mean"], device=device),
+        torch.tensor(expert_config["action_scale"], device=device),
+    ).to(device)
+    seed_everything(cfg.seed + 1)
+    predicted_action = flow.sample(
+        batch_size=1,
+        step_fn=partial(expert, prefix=prefix),
+        device=device,
+        dtype=dtype,
+    ).float()
+    history_xy = history[None, :, :2].to(device)
+    target_xy = future[None, :, :2].to(device)
+    target_action = action_space.trajectory_to_action(history_xy, target_xy)
+    predicted_xy = action_space.action_to_trajectory(predicted_action, history_xy)
+
+    post_boundary_index = boundary_index + 1
+    returned_masked_count = None
+    if 0 <= post_boundary_index < len(generated.logits or ()):
+        post_boundary_logits = generated.logits[post_boundary_index][0]
+        returned_masked_count = int(
+            torch.isneginf(
+                post_boundary_logits[trajectory_start : trajectory_start + trajectory_size]
+            ).sum().cpu()
+        )
+
+    result = {
+        "stage": "structured_vlm_to_diffusion_handoff",
+        "device": str(device),
+        "dtype": str(dtype),
+        "vlm_dir": str(vlm_dir),
+        "stage2_dir": str(stage2_dir),
+        "prompt_length": prompt_length,
+        "generated_token_count_including_one_cache_update_token": int(new_ids.shape[1]),
+        "boundary_index_in_generated_tokens": boundary_index,
+        "expected_cache_length_through_boundary": expected_prefix_length,
+        "visible_vlm_output": processor.tokenizer.decode(
+            visible_new_ids[0], skip_special_tokens=False
+        ),
+        "boundary_generated": bool((new_ids == boundary_id).any()),
+        "visible_discrete_trajectory_tokens": int(
+            (
+                (visible_new_ids >= trajectory_start)
+                & (visible_new_ids < trajectory_start + trajectory_size)
+            ).sum().cpu()
+        ),
+        "returned_post_boundary_masked_logits": returned_masked_count,
+        "handoff_prefix": {
+            "layers": prefix.num_layers,
+            "sequence_length": prefix.sequence_length,
+            "first_key_shape": list(prefix.keys[0].shape),
+            "all_finite": bool(
+                all(
+                    torch.isfinite(key).all() and torch.isfinite(value).all()
+                    for key, value in zip(prefix.keys, prefix.values)
+                )
+            ),
+        },
+        "expert_architecture": asdict(architecture),
+        "expert_parameters": parameter_counts(expert),
+        "diffusion_output_shape": list(predicted_action.shape),
+        "diffusion_output_finite": bool(torch.isfinite(predicted_action).all()),
+        "action_mse": float(
+            torch.nn.functional.mse_loss(predicted_action, target_action).cpu()
+        ),
+        **trajectory_metrics(predicted_xy, target_xy),
+        "quality_note": (
+            "This verifies exact structured-generation KV compatibility. The expert was trained "
+            "with the earlier plain-text Stage-1 prefix, so trajectory quality is diagnostic only."
+        ),
+    }
+    output_dir.mkdir(parents=True, exist_ok=True)
+    write_json(output_dir / "handoff_metrics.json", result)
+    save_prediction_arrays(
+        output_dir / "handoff_prediction.npz",
+        history[:, :2],
+        target_action[0],
+        predicted_action[0],
+        target_xy[0],
+        predicted_xy[0],
+    )
+    save_trajectory_plot(
+        output_dir / "handoff_trajectory.png",
+        history[:, :2],
+        target_xy[0],
+        predicted_xy[0],
+        prediction_label="diffusion from structured VLM KV",
+        title="Structured VLM handoff to diffusion",
+    )
+    print(json.dumps(result, indent=2, ensure_ascii=False))
+    del expert, vlm
+    empty_device_cache(device)
+
+
+if __name__ == "__main__":
+    main()

src/alpamayo_lite/__init__.py

@@ -0,0 +1,5 @@
+"""Small, local two-stage Alpamayo architecture validation."""
+
+from .config import ExperimentConfig, load_config
+
+__all__ = ["ExperimentConfig", "load_config"]

src/alpamayo_lite/action_space.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+
+@dataclass(frozen=True)
+class ActionStatistics:
+    mean: torch.Tensor
+    scale: torch.Tensor
+
+
+class RelativeDeltaActionSpace(nn.Module):
+    """Invertible local validation action space.
+
+    Alpamayo's release model predicts normalized acceleration and curvature. For the
+    architecture smoke test we use normalized per-waypoint XY deltas. It preserves
+    the important contract: diffusion predicts continuous actions, and a deterministic
+    decoder integrates those actions back to relative positions.
+    """
+
+    def __init__(self, mean: torch.Tensor | None = None, scale: torch.Tensor | None = None):
+        super().__init__()
+        if mean is None:
+            mean = torch.zeros(2, dtype=torch.float32)
+        if scale is None:
+            scale = torch.ones(2, dtype=torch.float32)
+        self.register_buffer("mean", mean.float().reshape(2))
+        self.register_buffer("scale", scale.float().reshape(2).clamp_min(1e-3))
+
+    @staticmethod
+    def fit(history_xy: torch.Tensor, future_xy: torch.Tensor) -> ActionStatistics:
+        if history_xy.shape[-1] != 2 or future_xy.shape[-1] != 2:
+            raise ValueError("history_xy and future_xy must end in XY coordinates")
+        start = history_xy[..., -1:, :]
+        previous = torch.cat([start, future_xy[..., :-1, :]], dim=-2)
+        deltas = future_xy - previous
+        flat = deltas.reshape(-1, 2).float()
+        mean = flat.mean(dim=0)
+        # A floor keeps a one-clip overfit target numerically well conditioned.
+        scale = flat.std(dim=0, unbiased=False).clamp_min(0.05)
+        return ActionStatistics(mean=mean, scale=scale)
+
+    def trajectory_to_action(
+        self, history_xy: torch.Tensor, future_xy: torch.Tensor
+    ) -> torch.Tensor:
+        start = history_xy[..., -1:, :]
+        previous = torch.cat([start, future_xy[..., :-1, :]], dim=-2)
+        deltas = future_xy - previous
+        return (deltas - self.mean) / self.scale
+
+    def action_to_trajectory(self, action: torch.Tensor, history_xy: torch.Tensor) -> torch.Tensor:
+        deltas = action * self.scale + self.mean
+        return history_xy[..., -1:, :] + torch.cumsum(deltas, dim=-2)

src/alpamayo_lite/config.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+from pathlib import Path
+from typing import Any, TypeVar
+
+import yaml
+
+
+@dataclass(frozen=True)
+class PathsConfig:
+    model_dir: str
+    pai_raw_dir: str
+    manifest: str
+    stage1_dir: str
+    stage2_dir: str
+
+
+@dataclass(frozen=True)
+class VLMConfig:
+    repo_id: str
+    dtype: str = "float16"
+    attn_implementation: str = "eager"
+    max_reasoning_tokens: int = 48
+    num_additional_image_tokens: int = 1
+
+
+@dataclass(frozen=True)
+class DataConfig:
+    repo_id: str
+    num_clips: int = 1
+    camera: str = "camera_front_wide_120fov"
+    num_frames: int = 1
+    history_steps: int = 8
+    source_future_steps: int = 64
+    expert_waypoints: int = 16
+    time_step: float = 0.1
+
+
+@dataclass(frozen=True)
+class Stage1Config:
+    steps: int = 80
+    learning_rate: float = 1e-4
+    weight_decay: float = 0.0
+    adam_epsilon: float = 1e-4
+    freeze_vision_tower: bool = True
+    freeze_multimodal_projector: bool = True
+    train_last_n_language_layers: int = 8
+    train_final_norm: bool = True
+    train_lm_head: bool = False
+    gradient_checkpointing: bool = True
+    log_every: int = 5
+
+
+@dataclass(frozen=True)
+class Stage2Config:
+    steps: int = 250
+    learning_rate: float = 3e-4
+    weight_decay: float = 0.0
+    adam_epsilon: float = 1e-4
+    expert_hidden_size: int = 448
+    expert_intermediate_size: int = 2432
+    expert_num_layers: int = 24
+    expert_num_attention_heads: int = 7
+    expert_num_key_value_heads: int = 2
+    expert_head_dim: int = 64
+    expert_non_causal_attention: bool = True
+    action_projection_hidden_size: int = 256
+    action_fourier_features: int = 20
+    inference_steps: int = 10
+    log_every: int = 10
+    checkpoint_every: int = 50
+
+
+@dataclass(frozen=True)
+class ExperimentConfig:
+    seed: int
+    paths: PathsConfig
+    vlm: VLMConfig
+    data: DataConfig
+    stage1: Stage1Config
+    stage2: Stage2Config
+    root: Path
+
+    def resolve(self, value: str) -> Path:
+        path = Path(value)
+        return path if path.is_absolute() else self.root / path
+
+
+T = TypeVar("T")
+
+
+def _construct(cls: type[T], values: dict[str, Any]) -> T:
+    known = {item.name for item in fields(cls)}
+    unknown = set(values) - known
+    if unknown:
+        raise ValueError(f"Unknown {cls.__name__} keys: {sorted(unknown)}")
+    return cls(**values)
+
+
+def load_config(path: str | Path) -> ExperimentConfig:
+    config_path = Path(path).expanduser().resolve()
+    raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
+    return ExperimentConfig(
+        seed=int(raw["seed"]),
+        paths=_construct(PathsConfig, raw["paths"]),
+        vlm=_construct(VLMConfig, raw["vlm"]),
+        data=_construct(DataConfig, raw["data"]),
+        stage1=_construct(Stage1Config, raw["stage1"]),
+        stage2=_construct(Stage2Config, raw["stage2"]),
+        root=config_path.parent.parent,
+    )

src/alpamayo_lite/data.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+@dataclass(frozen=True)
+class OverfitSample:
+    clip_id: str
+    t0_us: int
+    reasoning: str
+    image_path: Path
+    trajectory_path: Path
+
+    def load_image(self) -> Image.Image:
+        with Image.open(self.image_path) as image:
+            return image.convert("RGB")
+
+    def load_trajectory(self) -> tuple[torch.Tensor, torch.Tensor]:
+        with np.load(self.trajectory_path) as arrays:
+            history = torch.from_numpy(arrays["history_xyz"]).float()
+            future = torch.from_numpy(arrays["future_xyz"]).float()
+        return history, future
+
+
+def load_manifest(path: str | Path) -> list[OverfitSample]:
+    manifest_path = Path(path).expanduser().resolve()
+    root = manifest_path.parent.parent.parent
+    samples: list[OverfitSample] = []
+    for line in manifest_path.read_text(encoding="utf-8").splitlines():
+        if not line.strip():
+            continue
+        row = json.loads(line)
+        image_path = Path(row["image_path"])
+        trajectory_path = Path(row["trajectory_path"])
+        samples.append(
+            OverfitSample(
+                clip_id=str(row["clip_id"]),
+                t0_us=int(row["t0_us"]),
+                reasoning=str(row["reasoning"]),
+                image_path=image_path if image_path.is_absolute() else root / image_path,
+                trajectory_path=(
+                    trajectory_path if trajectory_path.is_absolute() else root / trajectory_path
+                ),
+            )
+        )
+    if not samples:
+        raise ValueError(f"No samples found in {manifest_path}")
+    return samples

src/alpamayo_lite/diffusion.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Small flow-matching implementation derived from NVlabs/alpamayo."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+from torch import nn
+
+
+class FlowMatching(nn.Module):
+    def __init__(self, action_shape: tuple[int, int], inference_steps: int = 10):
+        super().__init__()
+        self.action_shape = action_shape
+        self.inference_steps = inference_steps
+
+    def construct_training_data(self, target_action: torch.Tensor) -> dict[str, torch.Tensor]:
+        batch_size = target_action.shape[0]
+        t = torch.rand(batch_size, 1, 1, device=target_action.device)
+        noise = torch.randn_like(target_action)
+        noisy_action = t * target_action + (1.0 - t) * noise
+        return {
+            "target_action": target_action,
+            "noisy_action": noisy_action,
+            "timestep": t,
+            "noise": noise,
+        }
+
+    @staticmethod
+    def loss(training_data: dict[str, torch.Tensor], prediction: torch.Tensor) -> torch.Tensor:
+        target_velocity = training_data["target_action"] - training_data["noise"]
+        return torch.nn.functional.mse_loss(prediction.float(), target_velocity.float())
+
+    @torch.no_grad()
+    def sample(
+        self,
+        batch_size: int,
+        step_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        device: torch.device,
+        dtype: torch.dtype,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        x = torch.randn(
+            batch_size,
+            *self.action_shape,
+            device=device,
+            dtype=dtype,
+            generator=generator,
+        )
+        time_steps = torch.linspace(0.0, 1.0, self.inference_steps + 1, device=device)
+        for index in range(self.inference_steps):
+            t = time_steps[index].expand(batch_size).reshape(batch_size, 1, 1)
+            dt = time_steps[index + 1] - time_steps[index]
+            x = x + dt * step_fn(x, t)
+        return x

src/alpamayo_lite/evaluation.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+
+def trajectory_metrics(predicted_xy: torch.Tensor, target_xy: torch.Tensor) -> dict[str, float]:
+    distances = torch.linalg.vector_norm(predicted_xy.float() - target_xy.float(), dim=-1)
+    return {
+        "ade_m": float(distances.mean().cpu()),
+        "fde_m": float(distances[..., -1].mean().cpu()),
+        "max_error_m": float(distances.max().cpu()),
+    }
+
+
+def save_trajectory_plot(
+    path: str | Path,
+    history_xy: torch.Tensor,
+    target_xy: torch.Tensor,
+    predicted_xy: torch.Tensor,
+    *,
+    prediction_label: str = "diffusion",
+    title: str = "PAI overfit trajectory",
+) -> None:
+    output = Path(path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    history = history_xy.detach().cpu().numpy()
+    target = target_xy.detach().cpu().numpy()
+    predicted = predicted_xy.detach().cpu().numpy()
+    figure, axis = plt.subplots(figsize=(6, 5))
+    axis.plot(history[:, 0], history[:, 1], "o-", label="history", color="#555555")
+    axis.plot(target[:, 0], target[:, 1], "o-", label="target", color="#d62728")
+    axis.plot(
+        predicted[:, 0],
+        predicted[:, 1],
+        "o-",
+        label=prediction_label,
+        color="#1f77b4",
+    )
+    axis.scatter([0], [0], marker="x", color="black", label="current ego")
+    axis.set_xlabel("ego-relative x (m)")
+    axis.set_ylabel("ego-relative y (m)")
+    axis.set_title(title)
+    axis.axis("equal")
+    axis.grid(alpha=0.25)
+    axis.legend()
+    figure.tight_layout()
+    figure.savefig(output, dpi=160)
+    plt.close(figure)
+
+
+def save_prediction_arrays(
+    path: str | Path,
+    history_xy: torch.Tensor,
+    target_action: torch.Tensor,
+    predicted_action: torch.Tensor,
+    target_xy: torch.Tensor,
+    predicted_xy: torch.Tensor,
+) -> None:
+    output = Path(path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(
+        output,
+        history_xy=history_xy.detach().cpu().numpy(),
+        target_action=target_action.detach().cpu().numpy(),
+        predicted_action=predicted_action.detach().cpu().numpy(),
+        target_xy=target_xy.detach().cpu().numpy(),
+        predicted_xy=predicted_xy.detach().cpu().numpy(),
+    )

src/alpamayo_lite/expert.py

@@ -0,0 +1,237 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Past-KV-conditioned Qwen expert, following the Alpamayo connection pattern."""
+
+from __future__ import annotations
+
+import copy
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from transformers import Qwen2Config
+from transformers.cache_utils import DynamicCache
+from transformers.models.qwen2.modeling_qwen2 import Qwen2Model
+
+
+class FourierEncoder(nn.Module):
+    def __init__(self, dim: int, max_frequency: float = 100.0):
+        super().__init__()
+        if dim % 2:
+            raise ValueError("Fourier feature dimension must be even")
+        frequencies = torch.logspace(0, math.log10(max_frequency), steps=dim // 2)
+        self.register_buffer("frequencies", frequencies)
+
+    @property
+    def output_dim(self) -> int:
+        return self.frequencies.numel() * 2
+
+    def forward(self, values: torch.Tensor) -> torch.Tensor:
+        angles = values[..., None].float() * self.frequencies * (2.0 * torch.pi)
+        return torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1) * math.sqrt(2.0)
+
+
+class ActionTimeProjection(nn.Module):
+    def __init__(
+        self,
+        action_dim: int,
+        output_dim: int,
+        hidden_size: int,
+        fourier_features: int,
+    ):
+        super().__init__()
+        self.action_encoders = nn.ModuleList(
+            [FourierEncoder(fourier_features) for _ in range(action_dim)]
+        )
+        self.time_encoder = FourierEncoder(fourier_features)
+        input_dim = (action_dim + 1) * fourier_features
+        self.network = nn.Sequential(
+            nn.Linear(input_dim, hidden_size),
+            nn.SiLU(),
+            nn.LayerNorm(hidden_size),
+            nn.Linear(hidden_size, output_dim),
+            nn.LayerNorm(output_dim),
+        )
+
+    def forward(self, action: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+        batch_size, num_waypoints, _ = action.shape
+        action_features = torch.cat(
+            [encoder(action[..., index]) for index, encoder in enumerate(self.action_encoders)],
+            dim=-1,
+        )
+        scalar_timestep = timestep.reshape(batch_size, -1)[:, 0]
+        time_features = self.time_encoder(scalar_timestep)[:, None, :]
+        time_features = time_features.expand(batch_size, num_waypoints, -1)
+        features = torch.cat([action_features, time_features], dim=-1)
+        return self.network(features.to(dtype=self.network[0].weight.dtype))
+
+
+@dataclass(frozen=True)
+class PrefixKV:
+    keys: tuple[torch.Tensor, ...]
+    values: tuple[torch.Tensor, ...]
+
+    @property
+    def num_layers(self) -> int:
+        return len(self.keys)
+
+    @property
+    def sequence_length(self) -> int:
+        return int(self.keys[0].shape[-2])
+
+    @classmethod
+    def from_transformers_cache(cls, cache: object) -> "PrefixKV":
+        if hasattr(cache, "layers"):
+            keys = tuple(layer.keys.detach() for layer in cache.layers)
+            values = tuple(layer.values.detach() for layer in cache.layers)
+        else:
+            legacy = tuple(cache)
+            keys = tuple(item[0].detach() for item in legacy)
+            values = tuple(item[1].detach() for item in legacy)
+        return cls(keys=keys, values=values)
+
+    def dynamic_copy(self) -> DynamicCache:
+        entries = [(key.clone(), value.clone()) for key, value in zip(self.keys, self.values)]
+        try:
+            return DynamicCache(ddp_cache_data=entries)
+        except TypeError:
+            cache = DynamicCache()
+            for layer_index, (key, value) in enumerate(entries):
+                cache.update(key, value, layer_index)
+            return cache
+
+
+@dataclass(frozen=True)
+class ExpertArchitecture:
+    hidden_size: int
+    intermediate_size: int
+    num_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    action_projection_hidden_size: int
+    action_fourier_features: int
+    non_causal_attention: bool = True
+
+
+class KVCachedDiffusionExpert(nn.Module):
+    def __init__(
+        self,
+        vlm_text_config: Qwen2Config,
+        architecture: ExpertArchitecture,
+        action_dim: int = 2,
+    ):
+        super().__init__()
+        if architecture.hidden_size != architecture.num_attention_heads * architecture.head_dim:
+            raise ValueError("expert hidden_size must equal num_attention_heads * head_dim")
+        if architecture.num_layers != vlm_text_config.num_hidden_layers:
+            raise ValueError("direct past-KV reuse requires the same number of layers")
+        if architecture.num_key_value_heads != vlm_text_config.num_key_value_heads:
+            raise ValueError("direct past-KV reuse requires the same number of KV heads")
+        vlm_head_dim = getattr(
+            vlm_text_config,
+            "head_dim",
+            vlm_text_config.hidden_size // vlm_text_config.num_attention_heads,
+        )
+        if architecture.head_dim != vlm_head_dim:
+            raise ValueError("direct past-KV reuse requires the same attention head dimension")
+
+        expert_config = copy.deepcopy(vlm_text_config)
+        expert_config.hidden_size = architecture.hidden_size
+        expert_config.intermediate_size = architecture.intermediate_size
+        expert_config.num_hidden_layers = architecture.num_layers
+        expert_config.num_attention_heads = architecture.num_attention_heads
+        expert_config.num_key_value_heads = architecture.num_key_value_heads
+        expert_config.head_dim = architecture.head_dim
+        expert_config.vocab_size = 1
+        expert_config.tie_word_embeddings = False
+        expert_config.use_cache = True
+        expert_config._attn_implementation = "eager"
+
+        self.architecture = architecture
+        self.expert = Qwen2Model(expert_config)
+        # Inputs are always continuous action embeddings, exactly as in Alpamayo.
+        self.expert.embed_tokens = None
+        self.action_projection = ActionTimeProjection(
+            action_dim=action_dim,
+            output_dim=architecture.hidden_size,
+            hidden_size=architecture.action_projection_hidden_size,
+            fourier_features=architecture.action_fourier_features,
+        )
+        self.action_output = nn.Linear(architecture.hidden_size, action_dim)
+
+    def validate_prefix(self, prefix: PrefixKV) -> None:
+        if prefix.num_layers != self.architecture.num_layers:
+            raise ValueError(
+                f"KV layer mismatch: prefix={prefix.num_layers}, expert={self.architecture.num_layers}"
+            )
+        first_key = prefix.keys[0]
+        if first_key.shape[-3] != self.architecture.num_key_value_heads:
+            raise ValueError(
+                f"KV-head mismatch: prefix={first_key.shape[-3]}, "
+                f"expert={self.architecture.num_key_value_heads}"
+            )
+        if first_key.shape[-1] != self.architecture.head_dim:
+            raise ValueError(
+                f"head-dim mismatch: prefix={first_key.shape[-1]}, "
+                f"expert={self.architecture.head_dim}"
+            )
+
+    def forward(
+        self,
+        noisy_action: torch.Tensor,
+        timestep: torch.Tensor,
+        prefix: PrefixKV,
+    ) -> torch.Tensor:
+        self.validate_prefix(prefix)
+        batch_size, num_waypoints, _ = noisy_action.shape
+        if prefix.keys[0].shape[0] != batch_size:
+            raise ValueError("Prefix KV batch size does not match action batch size")
+
+        action_embeddings = self.action_projection(noisy_action, timestep)
+        prefix_length = prefix.sequence_length
+        position_ids = torch.arange(
+            prefix_length,
+            prefix_length + num_waypoints,
+            device=noisy_action.device,
+        )[None, :].expand(batch_size, -1)
+        if self.architecture.non_causal_attention:
+            # A prepared 4-D additive mask is returned unchanged by Transformers'
+            # mask builder.  This lets all diffusion tokens attend to one another,
+            # while still giving every query access to the frozen VLM prefix KV.
+            attention_mask = torch.zeros(
+                batch_size,
+                1,
+                num_waypoints,
+                prefix_length + num_waypoints,
+                dtype=action_embeddings.dtype,
+                device=noisy_action.device,
+            )
+        else:
+            attention_mask = torch.ones(
+                batch_size,
+                prefix_length + num_waypoints,
+                dtype=torch.long,
+                device=noisy_action.device,
+            )
+        cache_position = torch.arange(
+            prefix_length,
+            prefix_length + num_waypoints,
+            device=noisy_action.device,
+        )
+
+        kwargs: dict[str, object] = {}
+        if self.architecture.non_causal_attention:
+            kwargs["is_causal"] = False
+        outputs = self.expert(
+            inputs_embeds=action_embeddings,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=prefix.dynamic_copy(),
+            cache_position=cache_position,
+            use_cache=False,
+            **kwargs,
+        )
+        hidden = outputs.last_hidden_state[:, -num_waypoints:, :]
+        return self.action_output(hidden)

src/alpamayo_lite/prompts.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import re
+
+import torch
+
+
+def format_history(history_xyz: torch.Tensor) -> str:
+    history = history_xyz.detach().cpu().float()
+    return "; ".join(
+        f"t-{history.shape[-2] - 1 - index}:({point[0]:.2f},{point[1]:.2f})"
+        for index, point in enumerate(history)
+    )
+
+
+def clean_reasoning_text(text: str, max_characters: int = 320) -> str:
+    cleaned = re.sub(r"\s+", " ", text).strip()
+    if not cleaned:
+        return "Continue safely while following the visible lane."
+    return cleaned[:max_characters].rstrip()
+
+
+def build_vlm_prefix(history_xyz: torch.Tensor) -> str:
+    history_text = format_history(history_xyz)
+    return (
+        "<|im_start|>system\n"
+        "You are a driving vision-language model. Describe the immediate driving decision "
+        "briefly and concretely.<|im_end|>\n"
+        "<|im_start|>user\n"
+        "<image>\n"
+        f"Recent ego positions in the current ego-relative XY frame: {history_text}.\n"
+        "What should the vehicle do next based on the scene?<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+
+def build_stage1_text(history_xyz: torch.Tensor, reasoning: str) -> tuple[str, str]:
+    prefix = build_vlm_prefix(history_xyz)
+    target = clean_reasoning_text(reasoning) + "<|im_end|>"
+    return prefix, prefix + target
+
+
+def build_expert_prefix(history_xyz: torch.Tensor, reasoning: str) -> str:
+    return build_vlm_prefix(history_xyz) + clean_reasoning_text(reasoning) + "\nContinuous action:"

src/alpamayo_lite/runtime.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.backends.mps.is_available():
+        torch.mps.manual_seed(seed)
+
+
+def select_device() -> torch.device:
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def resolve_dtype(name: str, device: torch.device) -> torch.dtype:
+    mapping = {
+        "float32": torch.float32,
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+    }
+    if name not in mapping:
+        raise ValueError(f"Unsupported dtype: {name}")
+    dtype = mapping[name]
+    if device.type == "cpu" and dtype == torch.float16:
+        return torch.float32
+    return dtype
+
+
+def parameter_counts(module: torch.nn.Module) -> dict[str, int]:
+    total = sum(parameter.numel() for parameter in module.parameters())
+    trainable = sum(
+        parameter.numel() for parameter in module.parameters() if parameter.requires_grad
+    )
+    return {"total": total, "trainable": trainable, "frozen": total - trainable}
+
+
+def write_json(path: str | Path, value: Any) -> None:
+    output = Path(path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
+
+
+def empty_device_cache(device: torch.device) -> None:
+    if device.type == "mps":
+        torch.mps.empty_cache()
+    elif device.type == "cuda":
+        torch.cuda.empty_cache()

src/alpamayo_lite/trajectory_tokens.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from transformers import LogitsProcessor, StoppingCriteria
+
+
+STRUCTURED_SPECIAL_TOKENS = (
+    "<|cot_start|>",
+    "<|cot_end|>",
+    "<|meta_action_start|>",
+    "<|meta_action_end|>",
+    "<|traj_future_start|>",
+    "<|traj_future_end|>",
+)
+
+
+@dataclass(frozen=True)
+class TrajectoryVocabulary:
+    size: int
+    start_id: int
+    end_id: int
+    special_token_ids: dict[str, int]
+
+    def token_id(self, value: int) -> int:
+        if not 0 <= value < self.size:
+            raise ValueError(f"Trajectory token value is outside 0..{self.size - 1}: {value}")
+        return self.start_id + value
+
+
+def trajectory_token(value: int) -> str:
+    return f"<i{value}>"
+
+
+def add_trajectory_vocabulary(tokenizer: Any, size: int) -> TrajectoryVocabulary:
+    if size < 2:
+        raise ValueError("Trajectory vocabulary must contain at least two bins")
+    discrete_tokens = [trajectory_token(value) for value in range(size)]
+    existing = set(tokenizer.get_vocab())
+    collisions = [token for token in (*discrete_tokens, *STRUCTURED_SPECIAL_TOKENS) if token in existing]
+    if collisions:
+        raise ValueError(f"Trajectory tokens already exist in tokenizer: {collisions[:4]}")
+
+    added_discrete = tokenizer.add_tokens(discrete_tokens)
+    if added_discrete != size:
+        raise RuntimeError(f"Expected to add {size} trajectory tokens, added {added_discrete}")
+    tokenizer.add_tokens(list(STRUCTURED_SPECIAL_TOKENS), special_tokens=True)
+
+    discrete_ids = tokenizer.convert_tokens_to_ids(discrete_tokens)
+    start_id = int(discrete_ids[0])
+    expected_ids = list(range(start_id, start_id + size))
+    if list(discrete_ids) != expected_ids:
+        raise RuntimeError("Trajectory token IDs are not contiguous")
+    special_token_ids = {
+        token: int(tokenizer.convert_tokens_to_ids(token)) for token in STRUCTURED_SPECIAL_TOKENS
+    }
+    return TrajectoryVocabulary(
+        size=size,
+        start_id=start_id,
+        end_id=start_id + size - 1,
+        special_token_ids=special_token_ids,
+    )
+
+
+def quantize_action(
+    action: torch.Tensor,
+    size: int,
+    minimum: float,
+    maximum: float,
+) -> torch.LongTensor:
+    if maximum <= minimum:
+        raise ValueError("maximum must be greater than minimum")
+    normalized = (action.float().clamp(minimum, maximum) - minimum) / (maximum - minimum)
+    return (normalized * (size - 1)).round().long()
+
+
+def dequantize_action(
+    tokens: torch.Tensor,
+    size: int,
+    minimum: float,
+    maximum: float,
+) -> torch.Tensor:
+    if maximum <= minimum:
+        raise ValueError("maximum must be greater than minimum")
+    if torch.any(tokens < 0) or torch.any(tokens >= size):
+        raise ValueError("Trajectory token values are outside the configured vocabulary")
+    return tokens.float() / (size - 1) * (maximum - minimum) + minimum
+
+
+def build_structured_target(
+    reasoning: str,
+    meta_action: str,
+    token_values: torch.Tensor,
+) -> str:
+    flattened = token_values.detach().cpu().reshape(-1).tolist()
+    trajectory = "".join(trajectory_token(int(value)) for value in flattened)
+    return (
+        f"<|cot_start|>{reasoning}<|cot_end|>\n"
+        f"<|meta_action_start|>{meta_action}<|meta_action_end|>\n"
+        f"<|traj_future_start|>{trajectory}<|traj_future_end|><|im_end|>"
+    )
+
+
+def extract_trajectory_values(
+    token_ids: torch.Tensor,
+    vocabulary: TrajectoryVocabulary,
+) -> torch.LongTensor:
+    sequence = token_ids.detach().reshape(-1)
+    start_token_id = vocabulary.special_token_ids["<|traj_future_start|>"]
+    end_token_id = vocabulary.special_token_ids["<|traj_future_end|>"]
+    start_positions = torch.nonzero(sequence == start_token_id, as_tuple=False).flatten()
+    if start_positions.numel() == 0:
+        raise ValueError("Generated sequence does not contain <|traj_future_start|>")
+    start = int(start_positions[-1]) + 1
+    end_positions = torch.nonzero(sequence[start:] == end_token_id, as_tuple=False).flatten()
+    if end_positions.numel() == 0:
+        raise ValueError("Generated sequence does not contain <|traj_future_end|>")
+    end = start + int(end_positions[0])
+    values = sequence[start:end] - vocabulary.start_id
+    if torch.any(values < 0) or torch.any(values >= vocabulary.size):
+        raise ValueError("Generated trajectory block contains non-trajectory tokens")
+    return values.long()
+
+
+class DiscreteTrajectoryMask(LogitsProcessor):
+    """Source-style mask that prevents the VLM from sampling discrete trajectory tokens."""
+
+    def __init__(self, start_id: int, size: int):
+        self.start_id = start_id
+        self.size = size
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        del input_ids
+        scores[:, self.start_id : self.start_id + self.size] = float("-inf")
+        return scores
+
+
+class StopOneTokenAfterBoundary(StoppingCriteria):
+    """Match Alpamayo's handoff behavior so generation updates the cache once after the marker."""
+
+    def __init__(self, boundary_token_id: int):
+        self.boundary_token_id = boundary_token_id
+        self.boundary_found: torch.BoolTensor | None = None
+
+    def __call__(
+        self,
+        input_ids: torch.LongTensor,
+        scores: torch.FloatTensor,
+        **kwargs: Any,
+    ) -> torch.BoolTensor:
+        del scores, kwargs
+        if self.boundary_found is None:
+            self.boundary_found = torch.zeros(
+                input_ids.shape[0], dtype=torch.bool, device=input_ids.device
+            )
+        if self.boundary_found.all():
+            return torch.ones_like(self.boundary_found)
+        self.boundary_found |= input_ids[:, -1] == self.boundary_token_id
+        return torch.zeros_like(self.boundary_found)
+
+
+def mask_after_boundary(
+    token_ids: torch.Tensor,
+    boundary_token_id: int,
+    pad_token_id: int,
+) -> torch.Tensor:
+    result = token_ids.clone()
+    boundary_mask = result == boundary_token_id
+    sequence_length = result.shape[-1]
+    positions = torch.arange(sequence_length, device=result.device).expand_as(result)
+    first_boundary = torch.where(
+        boundary_mask,
+        positions,
+        torch.full_like(positions, sequence_length),
+    ).min(dim=-1, keepdim=True).values
+    result[positions > first_boundary] = pad_token_id
+    return result

src/alpamayo_lite/vlm.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+from PIL import Image
+from transformers import AutoProcessor, LlavaForConditionalGeneration
+
+from .config import Stage1Config, VLMConfig
+from .expert import PrefixKV
+from .runtime import parameter_counts
+
+
+def load_vlm(
+    model_path: str | Path,
+    config: VLMConfig,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[Any, LlavaForConditionalGeneration]:
+    processor = AutoProcessor.from_pretrained(model_path, use_fast=False)
+    # The selected third-party checkpoint was saved with this value set to 0,
+    # although its CLIP tower has one CLS token.  LLaVA's "default" feature
+    # selection removes CLS after counting it, so the processor must start at 1
+    # to expand <image> into the same 576 tokens produced by the vision tower.
+    processor.num_additional_image_tokens = config.num_additional_image_tokens
+    model = LlavaForConditionalGeneration.from_pretrained(
+        model_path,
+        dtype=dtype,
+        attn_implementation=config.attn_implementation,
+        low_cpu_mem_usage=True,
+    )
+    model.to(device)
+    return processor, model
+
+
+def _core_modules(model: LlavaForConditionalGeneration) -> tuple[Any, Any, Any, Any]:
+    core = model.model
+    vision_tower = core.vision_tower
+    projector = core.multi_modal_projector
+    language_model = core.language_model
+    language_backbone = language_model
+    return vision_tower, projector, language_model, language_backbone
+
+
+def configure_stage1_trainability(
+    model: LlavaForConditionalGeneration, config: Stage1Config
+) -> dict[str, Any]:
+    model.requires_grad_(False)
+    vision_tower, projector, language_model, backbone = _core_modules(model)
+    layers = backbone.layers
+    train_last = config.train_last_n_language_layers
+    if not (1 <= train_last <= len(layers)):
+        raise ValueError(f"train_last_n_language_layers must be within 1..{len(layers)}")
+    first_trainable = len(layers) - train_last
+    for layer in layers[first_trainable:]:
+        layer.requires_grad_(True)
+    if config.train_final_norm:
+        backbone.norm.requires_grad_(True)
+    if config.train_lm_head:
+        model.lm_head.requires_grad_(True)
+
+    if config.freeze_vision_tower:
+        vision_tower.requires_grad_(False)
+    if config.freeze_multimodal_projector:
+        projector.requires_grad_(False)
+
+    if config.gradient_checkpointing:
+        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    model.config.use_cache = False
+    language_model.config.use_cache = False
+
+    counts = parameter_counts(model)
+    counts.update(
+        {
+            "language_layers_total": len(layers),
+            "language_layers_trainable": train_last,
+            "first_trainable_language_layer": first_trainable,
+            "vision_trainable": any(p.requires_grad for p in vision_tower.parameters()),
+            "projector_trainable": any(p.requires_grad for p in projector.parameters()),
+            "lm_head_trainable": any(p.requires_grad for p in model.lm_head.parameters()),
+        }
+    )
+    return counts
+
+
+def enforce_frozen_module_eval(model: LlavaForConditionalGeneration) -> None:
+    vision_tower, projector, _, _ = _core_modules(model)
+    vision_tower.eval()
+    projector.eval()
+
+
+def _move_inputs(
+    inputs: dict[str, torch.Tensor], device: torch.device, dtype: torch.dtype
+) -> dict[str, torch.Tensor]:
+    moved: dict[str, torch.Tensor] = {}
+    for key, value in inputs.items():
+        if not isinstance(value, torch.Tensor):
+            moved[key] = value
+        elif value.is_floating_point():
+            moved[key] = value.to(device=device, dtype=dtype)
+        else:
+            moved[key] = value.to(device=device)
+    return moved
+
+
+def process_image_text(
+    processor: Any,
+    image: Image.Image,
+    text: str,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> dict[str, torch.Tensor]:
+    encoded = processor(text=text, images=image, return_tensors="pt")
+    return _move_inputs(dict(encoded), device=device, dtype=dtype)
+
+
+def prepare_stage1_batch(
+    processor: Any,
+    image: Image.Image,
+    prefix: str,
+    full_text: str,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[dict[str, torch.Tensor], torch.Tensor, int]:
+    prefix_inputs = process_image_text(processor, image, prefix, device, dtype)
+    full_inputs = process_image_text(processor, image, full_text, device, dtype)
+    prefix_length = prefix_inputs["input_ids"].shape[1]
+    labels = full_inputs["input_ids"].clone()
+    labels[:, :prefix_length] = -100
+    if "attention_mask" in full_inputs:
+        labels[full_inputs["attention_mask"] == 0] = -100
+    return full_inputs, labels, prefix_length
+
+
+@torch.no_grad()
+def generate_reasoning(
+    model: LlavaForConditionalGeneration,
+    processor: Any,
+    image: Image.Image,
+    prefix: str,
+    device: torch.device,
+    dtype: torch.dtype,
+    max_new_tokens: int,
+) -> str:
+    model.eval()
+    inputs = process_image_text(processor, image, prefix, device, dtype)
+    prompt_length = inputs["input_ids"].shape[1]
+    generated = model.generate(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        do_sample=False,
+        use_cache=True,
+        pad_token_id=processor.tokenizer.pad_token_id,
+        eos_token_id=processor.tokenizer.eos_token_id,
+    )
+    new_tokens = generated[:, prompt_length:]
+    return processor.batch_decode(new_tokens, skip_special_tokens=True)[0].strip()
+
+
+@torch.no_grad()
+def extract_prefix_kv(
+    model: LlavaForConditionalGeneration,
+    processor: Any,
+    image: Image.Image,
+    text: str,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> PrefixKV:
+    model.eval()
+    model.config.use_cache = True
+    model.model.language_model.config.use_cache = True
+    inputs = process_image_text(processor, image, text, device, dtype)
+    outputs = model(**inputs, use_cache=True, return_dict=True)
+    if outputs.past_key_values is None:
+        raise RuntimeError("VLM did not return past_key_values")
+    return PrefixKV.from_transformers_cache(outputs.past_key_values)
+
+
+def vlm_text_config(model: LlavaForConditionalGeneration) -> Any:
+    return model.config.text_config

tests/test_action_space.py

@@ -0,0 +1,13 @@
+import torch
+
+from alpamayo_lite.action_space import RelativeDeltaActionSpace
+
+
+def test_action_round_trip() -> None:
+    history = torch.tensor([[[-1.0, 0.0], [0.0, 0.0]]])
+    future = torch.tensor([[[1.0, 0.2], [2.1, 0.5], [3.2, 0.9]]])
+    statistics = RelativeDeltaActionSpace.fit(history, future)
+    action_space = RelativeDeltaActionSpace(statistics.mean, statistics.scale)
+    action = action_space.trajectory_to_action(history, future)
+    reconstructed = action_space.action_to_trajectory(action, history)
+    torch.testing.assert_close(reconstructed, future)

tests/test_complete_comparison.py

@@ -0,0 +1,46 @@
+import runpy
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+SCRIPT = runpy.run_path(
+    str(Path(__file__).resolve().parents[1] / "scripts" / "build_complete_comparison.py")
+)
+_decode_action = SCRIPT["_decode_action"]
+_extract_component = SCRIPT["_extract_component"]
+_prediction_metrics = SCRIPT["_prediction_metrics"]
+
+
+def test_decode_action_uses_statistics_and_nonzero_history_anchor() -> None:
+    history_xy = np.array([[8.0, -4.0], [10.0, -3.0]])
+    action = np.array([[1.0, -1.0], [0.0, 2.0]])
+    decoded = _decode_action(
+        action,
+        history_xy,
+        mean=np.array([0.5, 0.25]),
+        scale=np.array([2.0, 0.5]),
+    )
+    np.testing.assert_allclose(decoded, np.array([[12.5, -3.25], [13.0, -2.0]]))
+
+
+def test_extract_component_requires_one_ordered_nonempty_pair() -> None:
+    text = "<|cot_start|>keep left<|cot_end|>"
+    assert _extract_component(text, "cot") == "keep left"
+    with pytest.raises(ValueError):
+        _extract_component(text + text, "cot")
+    with pytest.raises(ValueError):
+        _extract_component("<|cot_start|><|cot_end|>", "cot")
+
+
+def test_prediction_metrics_use_normalized_action_and_euclidean_xy() -> None:
+    target_action = np.zeros((2, 2))
+    predicted_action = np.ones((2, 2))
+    target_xy = np.zeros((2, 2))
+    predicted_xy = np.array([[3.0, 4.0], [0.0, 0.0]])
+    assert _prediction_metrics(predicted_action, predicted_xy, target_action, target_xy) == {
+        "action_mse": 1.0,
+        "ade_m": 2.5,
+        "fde_m": 0.0,
+        "max_error_m": 5.0,
+    }

tests/test_diffusion.py

@@ -0,0 +1,11 @@
+import torch
+
+from alpamayo_lite.diffusion import FlowMatching
+
+
+def test_flow_matching_target_velocity() -> None:
+    flow = FlowMatching((4, 2), inference_steps=2)
+    target = torch.ones(1, 4, 2)
+    data = flow.construct_training_data(target)
+    exact_velocity = data["target_action"] - data["noise"]
+    assert float(flow.loss(data, exact_velocity)) == 0.0

tests/test_expert_contract.py

@@ -0,0 +1,69 @@
+import torch
+from transformers import Qwen2Config
+
+from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert, PrefixKV
+
+
+def tiny_expert(non_causal_attention: bool = False) -> KVCachedDiffusionExpert:
+    vlm_config = Qwen2Config(
+        vocab_size=32,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        max_position_embeddings=128,
+    )
+    architecture = ExpertArchitecture(
+        hidden_size=32,
+        intermediate_size=64,
+        num_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        head_dim=16,
+        action_projection_hidden_size=16,
+        action_fourier_features=4,
+        non_causal_attention=non_causal_attention,
+    )
+    return KVCachedDiffusionExpert(vlm_config, architecture)
+
+
+def test_expert_consumes_compatible_past_kv() -> None:
+    expert = tiny_expert()
+    prefix = PrefixKV(
+        keys=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
+        values=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
+    )
+    output = expert(torch.randn(1, 3, 2), torch.rand(1, 1, 1), prefix)
+    assert output.shape == (1, 3, 2)
+
+
+def test_expert_rejects_incompatible_kv_heads() -> None:
+    expert = tiny_expert()
+    prefix = PrefixKV(
+        keys=tuple(torch.randn(1, 1, 5, 16) for _ in range(2)),
+        values=tuple(torch.randn(1, 1, 5, 16) for _ in range(2)),
+    )
+    try:
+        expert.validate_prefix(prefix)
+    except ValueError as error:
+        assert "KV-head mismatch" in str(error)
+    else:
+        raise AssertionError("Expected incompatible KV heads to be rejected")
+
+
+def test_non_causal_expert_allows_future_action_attention() -> None:
+    torch.manual_seed(3)
+    expert = tiny_expert(non_causal_attention=True).eval()
+    prefix = PrefixKV(
+        keys=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
+        values=tuple(torch.randn(1, 2, 5, 16) for _ in range(2)),
+    )
+    first = torch.randn(1, 3, 2)
+    changed_future = first.clone()
+    changed_future[:, -1] += 0.137
+    timestep = torch.rand(1, 1, 1)
+    with torch.no_grad():
+        first_output = expert(first, timestep, prefix)
+        second_output = expert(changed_future, timestep, prefix)
+    assert not torch.allclose(first_output[:, 0], second_output[:, 0])

tests/test_trajectory_tokens.py

@@ -0,0 +1,39 @@
+import torch
+
+from alpamayo_lite.trajectory_tokens import (
+    DiscreteTrajectoryMask,
+    TrajectoryVocabulary,
+    dequantize_action,
+    extract_trajectory_values,
+    quantize_action,
+)
+
+
+def test_quantized_action_round_trip_is_bounded() -> None:
+    action = torch.tensor([[-4.0, -1.25], [0.0, 1.5], [3.9, 4.0]])
+    tokens = quantize_action(action, size=128, minimum=-4.0, maximum=4.0)
+    reconstructed = dequantize_action(tokens, size=128, minimum=-4.0, maximum=4.0)
+    assert float((reconstructed - action).abs().max()) <= 8.0 / 127.0
+
+
+def test_extract_trajectory_values_between_markers() -> None:
+    vocabulary = TrajectoryVocabulary(
+        size=8,
+        start_id=100,
+        end_id=107,
+        special_token_ids={"<|traj_future_start|>": 200, "<|traj_future_end|>": 201},
+    )
+    sequence = torch.tensor([1, 200, 103, 107, 100, 201, 2])
+    torch.testing.assert_close(
+        extract_trajectory_values(sequence, vocabulary), torch.tensor([3, 7, 0])
+    )
+
+
+def test_discrete_trajectory_mask_changes_only_trajectory_range() -> None:
+    input_ids = torch.tensor([[1, 2]])
+    scores = torch.randn(1, 16)
+    original = scores.clone()
+    masked = DiscreteTrajectoryMask(start_id=5, size=4)(input_ids, scores)
+    assert torch.isneginf(masked[:, 5:9]).all()
+    torch.testing.assert_close(masked[:, :5], original[:, :5])
+    torch.testing.assert_close(masked[:, 9:], original[:, 9:])