commit fb1bd51046456097958fb12bd416f31c28f03a43 Author: seeker Date: Sat Jul 11 13:08:42 2026 +0800 Add 0.5B Alpamayo validation diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..57e9b87 --- /dev/null +++ b/.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 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..91c82ac --- /dev/null +++ b/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：`` … ``； +- `<|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|><|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 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/` 进行接口研究，但该目录不会随本仓库分发。 diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..fc9807e --- /dev/null +++ b/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. diff --git a/configs/local_overfit.yaml b/configs/local_overfit.yaml new file mode 100644 index 0000000..2ae5af6 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..9e25413 --- /dev/null +++ b/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"] diff --git a/results/complete_comparison/comparison.json b/results/complete_comparison/comparison.json new file mode 100644 index 0000000..95a0c7c --- /dev/null +++ b/results/complete_comparison/comparison.json @@ -0,0 +1,1381 @@ +{ + "schema_version": 1, + "sample": { + "manifest_index": 0, + "clip_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20", + "t0_us": 1503899, + "camera": "camera_front_wide_120fov", + "history_steps": 8, + "expert_waypoints": 16, + "source_time_step_s": 0.1, + "image_path": "data/processed/images/sample_00_314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899.jpg", + "trajectory_path": "data/processed/trajectories/sample_00_314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899.npz" + }, + "prediction_protocol": { + "vlm": { + "decoding": "greedy untruncated structured generation", + "trajectory_representation": "32 discrete normalized-XY-delta tokens", + "trajectory_vocabulary_bins": 128 + }, + "diffusion": { + "decoding": "flow-matching sample", + "seed": 12, + "inference_steps": 10 + }, + "units": { + "action_mse": "normalized_action_squared", + "trajectory": "metres" + } + }, + "text": { + "ground_truth_reasoning": "Steer left following the temporary lane delineated by traffic cones through the construction zone.", + "vlm_reasoning": "Steer left following the temporary lane delineated by traffic cones through the construction zone.", + "vlm_reasoning_exact_match": true, + "vlm_meta_action": "STEER_LEFT", + "target_meta_action": "STEER_LEFT", + "vlm_meta_action_exact_match": true, + "vlm_structured_output": "<|cot_start|>Steer left following the temporary lane delineated by traffic cones through the construction zone.<|cot_end|>\n<|meta_action_start|>STEER_LEFT<|meta_action_end|>\n<|traj_future_start|><|traj_future_end|><|im_end|>", + "target_structured_output": "<|cot_start|>Steer left following the temporary lane delineated by traffic cones through the construction zone.<|cot_end|>\n<|meta_action_start|>STEER_LEFT<|meta_action_end|>\n<|traj_future_start|><|traj_future_end|><|im_end|>", + "vlm_structured_exact_match": true, + "diffusion_conditioning_reasoning": "Steer left following the temporary lane delineated by traffic cones through the construction zone.", + "diffusion_reasoning_exact_match": true + }, + "arrays": { + "ground_truth": { + "action": [ + [ + -3.4397106170654297, + -1.6662895679473877 + ], + [ + 0.46647655963897705, + -1.6230756044387817 + ], + [ + 0.3931064307689667, + -1.3917449712753296 + ], + [ + 0.26794740557670593, + -1.1304054260253906 + ], + [ + 0.1750280112028122, + -0.931678295135498 + ], + [ + 1.3138797283172607, + -0.017029158771038055 + ], + [ + -0.0019493416184559464, + -0.05992039293050766 + ], + [ + 0.07472366094589233, + 0.23706550896167755 + ], + [ + 0.13519389927387238, + 0.3722502589225769 + ], + [ + 0.06780000776052475, + 0.21779732406139374 + ], + [ + 1.1512386798858643, + 1.035483717918396 + ], + [ + -0.15684714913368225, + 0.8499115109443665 + ], + [ + -0.2395608276128769, + 0.9665660858154297 + ], + [ + -0.2766667306423187, + 0.9057282209396362 + ], + [ + -0.44114238023757935, + 0.7915734648704529 + ], + [ + 0.5104826092720032, + 1.4437676668167114 + ] + ], + "xy_m": [ + [ + 0.9489696025848389, + -0.00039592807297594845 + ], + [ + 4.71621561050415, + 0.002558943582698703 + ], + [ + 8.430525779724121, + 0.023451128974556923 + ], + [ + 12.054534912109375, + 0.06460751593112946 + ], + [ + 15.611503601074219, + 0.12117315828800201 + ], + [ + 19.990142822265625, + 0.24866046011447906 + ], + [ + 23.419424057006836, + 0.37282198667526245 + ], + [ + 26.904024124145508, + 0.5200117230415344 + ], + [ + 30.432252883911133, + 0.6776836514472961 + ], + [ + 33.91185760498047, + 0.8233793377876282 + ], + [ + 38.173152923583984, + 1.0324782133102417 + ], + [ + 41.49067687988281, + 1.2271878719329834 + ], + [ + 44.7485237121582, + 1.4309428930282593 + ], + [ + 47.97959899902344, + 1.6299805641174316 + ], + [ + 51.09200668334961, + 1.8201667070388794 + ], + [ + 54.8910026550293, + 2.0609238147735596 + ] + ], + "discrete_token_values": [ + 9, + 37, + 71, + 38, + 70, + 41, + 68, + 46, + 66, + 49, + 84, + 63, + 63, + 63, + 65, + 67, + 66, + 69, + 65, + 67, + 82, + 80, + 61, + 77, + 60, + 79, + 59, + 78, + 56, + 76, + 72, + 86 + ] + }, + "vlm": { + "action": [ + [ + -3.4330708980560303, + -1.6692912578582764 + ], + [ + 0.4724407196044922, + -1.6062991619110107 + ], + [ + 0.40944862365722656, + -1.4173228740692139 + ], + [ + 0.2834644317626953, + -1.1023621559143066 + ], + [ + 0.15748023986816406, + -0.9133858680725098 + ], + [ + 1.2913384437561035, + -0.03149604797363281 + ], + [ + -0.03149604797363281, + -0.03149604797363281 + ], + [ + 0.09448814392089844, + 0.2204723358154297 + ], + [ + 0.15748023986816406, + 0.34645652770996094 + ], + [ + 0.09448814392089844, + 0.2204723358154297 + ], + [ + 1.1653542518615723, + 1.039370059967041 + ], + [ + -0.15748023986816406, + 0.8503937721252441 + ], + [ + -0.2204723358154297, + 0.9763779640197754 + ], + [ + -0.2834646701812744, + 0.9133858680725098 + ], + [ + -0.4724409580230713, + 0.7874016761779785 + ], + [ + 0.535433292388916, + 1.4173226356506348 + ] + ], + "xy_m": [ + [ + 0.9537601470947266, + -0.0006286799907684326 + ], + [ + 4.725309371948242, + 0.0036270394921302795 + ], + [ + 8.451410293579102, + 0.02253592759370804 + ], + [ + 12.086614608764648, + 0.06586679071187973 + ], + [ + 15.630922317504883, + 0.12385082244873047 + ], + [ + 19.993297576904297, + 0.25021636486053467 + ], + [ + 23.401260375976562, + 0.37658190727233887 + ], + [ + 26.90011978149414, + 0.5224850177764893 + ], + [ + 30.444427490234375, + 0.6781569123268127 + ], + [ + 33.94328689575195, + 0.8240600228309631 + ], + [ + 38.21476745605469, + 1.033460259437561 + ], + [ + 41.53183364868164, + 1.2282073497772217 + ], + [ + 44.80345153808594, + 1.4327231645584106 + ], + [ + 48.029624938964844, + 1.6323546171188354 + ], + [ + 51.11945343017578, + 1.8222172260284424 + ], + [ + 54.93644714355469, + 2.0609238147735596 + ] + ], + "discrete_token_values": [ + 9, + 37, + 71, + 38, + 70, + 41, + 68, + 46, + 66, + 49, + 84, + 63, + 63, + 63, + 65, + 67, + 66, + 69, + 65, + 67, + 82, + 80, + 61, + 77, + 60, + 79, + 59, + 78, + 56, + 76, + 72, + 86 + ] + }, + "diffusion": { + "action": [ + [ + -3.84375, + -0.443359375 + ], + [ + -0.68017578125, + -1.453125 + ], + [ + -0.204345703125, + -2.953125 + ], + [ + 1.0224609375, + -0.82275390625 + ], + [ + -0.126953125, + -0.8857421875 + ], + [ + 2.5078125, + -0.9423828125 + ], + [ + 0.2435302734375, + -0.045867919921875 + ], + [ + 1.005859375, + 0.17138671875 + ], + [ + 0.12469482421875, + 0.52001953125 + ], + [ + 0.432373046875, + -0.6748046875 + ], + [ + 1.95703125, + 0.78125 + ], + [ + 2.060546875, + 1.314453125 + ], + [ + -0.9697265625, + 2.46875 + ], + [ + -0.52783203125, + 1.48046875 + ], + [ + -1.5810546875, + 0.258056640625 + ], + [ + -1.0625, + 1.236328125 + ] + ], + "xy_m": [ + [ + 0.657459020614624, + 0.09442976862192154 + ], + [ + 3.5974063873291016, + 0.11056257039308548 + ], + [ + 6.880660533905029, + 0.010385759174823761 + ], + [ + 11.049043655395508, + 0.07539737224578857 + ], + [ + 14.38813591003418, + 0.1355248987674713 + ], + [ + 19.628185272216797, + 0.19126051664352417 + ], + [ + 23.234577178955078, + 0.31651169061660767 + ], + [ + 27.39098358154297, + 0.4586087167263031 + ], + [ + 30.911636352539062, + 0.6277386546134949 + ], + [ + 34.65427780151367, + 0.7042222023010254 + ], + [ + 39.496944427490234, + 0.8936078548431396 + ], + [ + 44.41429901123047, + 1.124337911605835 + ], + [ + 47.14533615112305, + 1.4445719718933105 + ], + [ + 50.19519805908203, + 1.688174843788147 + ], + [ + 52.48516845703125, + 1.8369922637939453 + ], + [ + 55.14927673339844, + 2.061664581298828 + ] + ], + "seed": 12 + }, + "history_xy_m": [ + [ + -6.5544657707214355, + 0.035949498414993286 + ], + [ + -5.636332988739014, + 0.025562627241015434 + ], + [ + -4.717358112335205, + 0.019754931330680847 + ], + [ + -3.7914953231811523, + 0.012032918632030487 + ], + [ + -2.853987693786621, + 0.008905102498829365 + ], + [ + -1.9065805673599243, + 0.006025215145200491 + ], + [ + -0.9533587694168091, + 0.0025723553262650967 + ], + [ + 0.0, + 0.0 + ] + ], + "source_future_indices": [ + 0, + 4, + 8, + 12, + 16, + 21, + 25, + 29, + 33, + 37, + 42, + 46, + 50, + 54, + 58, + 63 + ], + "waypoint_times_s": [ + 0.1, + 0.5, + 0.9, + 1.3, + 1.7000000000000002, + 2.2, + 2.6, + 3.0, + 3.4000000000000004, + 3.8000000000000003, + 4.3, + 4.7, + 5.1000000000000005, + 5.5, + 5.9, + 6.4 + ] + }, + "waypoints": [ + { + "index": 1, + "source_future_index": 0, + "time_s": 0.1, + "action": { + "ground_truth": [ + -3.4397106170654297, + -1.6662895679473877 + ], + "vlm": [ + -3.4330708980560303, + -1.6692912578582764 + ], + "diffusion": [ + -3.84375, + -0.443359375 + ] + }, + "action_error_l2": { + "vlm": 0.00728670095756034, + "diffusion": 1.2879464584308167 + }, + "xy_m": { + "ground_truth": [ + 0.9489696025848389, + -0.00039592807297594845 + ], + "vlm": [ + 0.9537601470947266, + -0.0006286799907684326 + ], + "diffusion": [ + 0.657459020614624, + 0.09442976862192154 + ] + }, + "position_error_m": { + "vlm": 0.004796195383473365, + "diffusion": 0.30654580759535116 + } + }, + { + "index": 2, + "source_future_index": 4, + "time_s": 0.5, + "action": { + "ground_truth": [ + 0.46647655963897705, + -1.6230756044387817 + ], + "vlm": [ + 0.4724407196044922, + -1.6062991619110107 + ], + "diffusion": [ + -0.68017578125, + -1.453125 + ] + }, + "action_error_l2": { + "vlm": 0.01780506186402779, + "diffusion": 1.1591785017050988 + }, + "xy_m": { + "ground_truth": [ + 4.71621561050415, + 0.002558943582698703 + ], + "vlm": [ + 4.725309371948242, + 0.0036270394921302795 + ], + "diffusion": [ + 3.5974063873291016, + 0.11056257039308548 + ] + }, + "position_error_m": { + "vlm": 0.009156272498882664, + "diffusion": 1.1240101695561981 + } + }, + { + "index": 3, + "source_future_index": 8, + "time_s": 0.9, + "action": { + "ground_truth": [ + 0.3931064307689667, + -1.3917449712753296 + ], + "vlm": [ + 0.40944862365722656, + -1.4173228740692139 + ], + "diffusion": [ + -0.204345703125, + -2.953125 + ] + }, + "action_error_l2": { + "vlm": 0.03035286444028776, + "diffusion": 1.6717824758008162 + }, + "xy_m": { + "ground_truth": [ + 8.430525779724121, + 0.023451128974556923 + ], + "vlm": [ + 8.451410293579102, + 0.02253592759370804 + ], + "diffusion": [ + 6.880660533905029, + 0.010385759174823761 + ] + }, + "position_error_m": { + "vlm": 0.02090455721909409, + "diffusion": 1.5499203153987877 + } + }, + { + "index": 4, + "source_future_index": 12, + "time_s": 1.3, + "action": { + "ground_truth": [ + 0.26794740557670593, + -1.1304054260253906 + ], + "vlm": [ + 0.2834644317626953, + -1.1023621559143066 + ], + "diffusion": [ + 1.0224609375, + -0.82275390625 + ] + }, + "action_error_l2": { + "vlm": 0.032050009363179546, + "diffusion": 0.8148252128373736 + }, + "xy_m": { + "ground_truth": [ + 12.054534912109375, + 0.06460751593112946 + ], + "vlm": [ + 12.086614608764648, + 0.06586679071187973 + ], + "diffusion": [ + 11.049043655395508, + 0.07539737224578857 + ] + }, + "position_error_m": { + "vlm": 0.03210440328783258, + "diffusion": 1.0055491476438747 + } + }, + { + "index": 5, + "source_future_index": 16, + "time_s": 1.7000000000000002, + "action": { + "ground_truth": [ + 0.1750280112028122, + -0.931678295135498 + ], + "vlm": [ + 0.15748023986816406, + -0.9133858680725098 + ], + "diffusion": [ + -0.126953125, + -0.8857421875 + ] + }, + "action_error_l2": { + "vlm": 0.025348316840923476, + "diffusion": 0.3054549600301843 + }, + "xy_m": { + "ground_truth": [ + 15.611503601074219, + 0.12117315828800201 + ], + "vlm": [ + 15.630922317504883, + 0.12385082244873047 + ], + "diffusion": [ + 14.38813591003418, + 0.1355248987674713 + ] + }, + "position_error_m": { + "vlm": 0.019602459875540927, + "diffusion": 1.2234518707065785 + } + }, + { + "index": 6, + "source_future_index": 21, + "time_s": 2.2, + "action": { + "ground_truth": [ + 1.3138797283172607, + -0.017029158771038055 + ], + "vlm": [ + 1.2913384437561035, + -0.03149604797363281 + ], + "diffusion": [ + 2.5078125, + -0.9423828125 + ] + }, + "action_error_l2": { + "vlm": 0.02678433110733247, + "diffusion": 1.5105478634480827 + }, + "xy_m": { + "ground_truth": [ + 19.990142822265625, + 0.24866046011447906 + ], + "vlm": [ + 19.993297576904297, + 0.25021636486053467 + ], + "diffusion": [ + 19.628185272216797, + 0.19126051664352417 + ] + }, + "position_error_m": { + "vlm": 0.003517572516526143, + "diffusion": 0.36648058822783336 + } + }, + { + "index": 7, + "source_future_index": 25, + "time_s": 2.6, + "action": { + "ground_truth": [ + -0.0019493416184559464, + -0.05992039293050766 + ], + "vlm": [ + -0.03149604797363281, + -0.03149604797363281 + ], + "diffusion": [ + 0.2435302734375, + -0.045867919921875 + ] + }, + "action_error_l2": { + "vlm": 0.04099940539405987, + "diffusion": 0.24588150277253198 + }, + "xy_m": { + "ground_truth": [ + 23.419424057006836, + 0.37282198667526245 + ], + "vlm": [ + 23.401260375976562, + 0.37658190727233887 + ], + "diffusion": [ + 23.234577178955078, + 0.31651169061660767 + ] + }, + "position_error_m": { + "vlm": 0.018548754984252572, + "diffusion": 0.19323358343645852 + } + }, + { + "index": 8, + "source_future_index": 29, + "time_s": 3.0, + "action": { + "ground_truth": [ + 0.07472366094589233, + 0.23706550896167755 + ], + "vlm": [ + 0.09448814392089844, + 0.2204723358154297 + ], + "diffusion": [ + 1.005859375, + 0.17138671875 + ] + }, + "action_error_l2": { + "vlm": 0.025806359338943326, + "diffusion": 0.9334492066902851 + }, + "xy_m": { + "ground_truth": [ + 26.904024124145508, + 0.5200117230415344 + ], + "vlm": [ + 26.90011978149414, + 0.5224850177764893 + ], + "diffusion": [ + 27.39098358154297, + 0.4586087167263031 + ] + }, + "position_error_m": { + "vlm": 0.004621804667577402, + "diffusion": 0.49081548705534744 + } + }, + { + "index": 9, + "source_future_index": 33, + "time_s": 3.4000000000000004, + "action": { + "ground_truth": [ + 0.13519389927387238, + 0.3722502589225769 + ], + "vlm": [ + 0.15748023986816406, + 0.34645652770996094 + ], + "diffusion": [ + 0.12469482421875, + 0.52001953125 + ] + }, + "action_error_l2": { + "vlm": 0.034088085117141044, + "diffusion": 0.14814178485892907 + }, + "xy_m": { + "ground_truth": [ + 30.432252883911133, + 0.6776836514472961 + ], + "vlm": [ + 30.444427490234375, + 0.6781569123268127 + ], + "diffusion": [ + 30.911636352539062, + 0.6277386546134949 + ] + }, + "position_error_m": { + "vlm": 0.01218380133562631, + "diffusion": 0.4819782284527732 + } + }, + { + "index": 10, + "source_future_index": 37, + "time_s": 3.8000000000000003, + "action": { + "ground_truth": [ + 0.06780000776052475, + 0.21779732406139374 + ], + "vlm": [ + 0.09448814392089844, + 0.2204723358154297 + ], + "diffusion": [ + 0.432373046875, + -0.6748046875 + ] + }, + "action_error_l2": { + "vlm": 0.026821862343970004, + "diffusion": 0.9641845528178778 + }, + "xy_m": { + "ground_truth": [ + 33.91185760498047, + 0.8233793377876282 + ], + "vlm": [ + 33.94328689575195, + 0.8240600228309631 + ], + "diffusion": [ + 34.65427780151367, + 0.7042222023010254 + ] + }, + "position_error_m": { + "vlm": 0.03143666093157371, + "diffusion": 0.7519216522735415 + } + }, + { + "index": 11, + "source_future_index": 42, + "time_s": 4.3, + "action": { + "ground_truth": [ + 1.1512386798858643, + 1.035483717918396 + ], + "vlm": [ + 1.1653542518615723, + 1.039370059967041 + ], + "diffusion": [ + 1.95703125, + 0.78125 + ] + }, + "action_error_l2": { + "vlm": 0.014640800071050068, + "diffusion": 0.84494760155749 + }, + "xy_m": { + "ground_truth": [ + 38.173152923583984, + 1.0324782133102417 + ], + "vlm": [ + 38.21476745605469, + 1.033460259437561 + ], + "diffusion": [ + 39.496944427490234, + 0.8936078548431396 + ] + }, + "position_error_m": { + "vlm": 0.04162611833154021, + "diffusion": 1.331055566937441 + } + }, + { + "index": 12, + "source_future_index": 46, + "time_s": 4.7, + "action": { + "ground_truth": [ + -0.15684714913368225, + 0.8499115109443665 + ], + "vlm": [ + -0.15748023986816406, + 0.8503937721252441 + ], + "diffusion": [ + 2.060546875, + 1.314453125 + ] + }, + "action_error_l2": { + "vlm": 0.0007958515720084112, + "diffusion": 2.2655319837630143 + }, + "xy_m": { + "ground_truth": [ + 41.49067687988281, + 1.2271878719329834 + ], + "vlm": [ + 41.53183364868164, + 1.2282073497772217 + ], + "diffusion": [ + 44.41429901123047, + 1.124337911605835 + ] + }, + "position_error_m": { + "vlm": 0.04116939340134956, + "diffusion": 2.925430648852423 + } + }, + { + "index": 13, + "source_future_index": 50, + "time_s": 5.1000000000000005, + "action": { + "ground_truth": [ + -0.2395608276128769, + 0.9665660858154297 + ], + "vlm": [ + -0.2204723358154297, + 0.9763779640197754 + ], + "diffusion": [ + -0.9697265625, + 2.46875 + ] + }, + "action_error_l2": { + "vlm": 0.021462606388743272, + "diffusion": 1.6702390584698135 + }, + "xy_m": { + "ground_truth": [ + 44.7485237121582, + 1.4309428930282593 + ], + "vlm": [ + 44.80345153808594, + 1.4327231645584106 + ], + "diffusion": [ + 47.14533615112305, + 1.4445719718933105 + ] + }, + "position_error_m": { + "vlm": 0.054956668638742615, + "diffusion": 2.3968511884068464 + } + }, + { + "index": 14, + "source_future_index": 54, + "time_s": 5.5, + "action": { + "ground_truth": [ + -0.2766667306423187, + 0.9057282209396362 + ], + "vlm": [ + -0.2834646701812744, + 0.9133858680725098 + ], + "diffusion": [ + -0.52783203125, + 1.48046875 + ] + }, + "action_error_l2": { + "vlm": 0.010239704174774942, + "diffusion": 0.6272245881452143 + }, + "xy_m": { + "ground_truth": [ + 47.97959899902344, + 1.6299805641174316 + ], + "vlm": [ + 48.029624938964844, + 1.6323546171188354 + ], + "diffusion": [ + 50.19519805908203, + 1.688174843788147 + ] + }, + "position_error_m": { + "vlm": 0.050082240312057325, + "diffusion": 2.2163631852922747 + } + }, + { + "index": 15, + "source_future_index": 58, + "time_s": 5.9, + "action": { + "ground_truth": [ + -0.44114238023757935, + 0.7915734648704529 + ], + "vlm": [ + -0.4724409580230713, + 0.7874016761779785 + ], + "diffusion": [ + -1.5810546875, + 0.258056640625 + ] + }, + "action_error_l2": { + "vlm": 0.03157538269426274, + "diffusion": 1.2585866160107093 + }, + "xy_m": { + "ground_truth": [ + 51.09200668334961, + 1.8201667070388794 + ], + "vlm": [ + 51.11945343017578, + 1.8222172260284424 + ], + "diffusion": [ + 52.48516845703125, + 1.8369922637939453 + ] + }, + "position_error_m": { + "vlm": 0.027523236355242352, + "diffusion": 1.393263373166715 + } + }, + { + "index": 16, + "source_future_index": 63, + "time_s": 6.4, + "action": { + "ground_truth": [ + 0.5104826092720032, + 1.4437676668167114 + ], + "vlm": [ + 0.535433292388916, + 1.4173226356506348 + ], + "diffusion": [ + -1.0625, + 1.236328125 + ] + }, + "action_error_l2": { + "vlm": 0.03635761627740968, + "diffusion": 1.5866018569828055 + }, + "xy_m": { + "ground_truth": [ + 54.8910026550293, + 2.0609238147735596 + ], + "vlm": [ + 54.93644714355469, + 2.0609238147735596 + ], + "diffusion": [ + 55.14927673339844, + 2.061664581298828 + ] + }, + "position_error_m": { + "vlm": 0.045444488525390625, + "diffusion": 0.2582751406784516 + } + } + ], + "metrics": { + "vlm": { + "computed_from_npz": { + "action_mse": 0.0003439451570161356, + "ade_m": 0.026104664266543905, + "fde_m": 0.045444488525390625, + "max_error_m": 0.054956668638742615 + }, + "reported_in_metrics_json": { + "action_mse": 0.0003439451684243977, + "ade_m": 0.02610466256737709, + "fde_m": 0.045444488525390625, + "max_error_m": 0.05495666712522507 + }, + "trajectory_token_matches": 32, + "trajectory_token_count": 32, + "trajectory_token_accuracy": 1.0 + }, + "diffusion": { + "computed_from_npz": { + "action_mse": 0.745336398308135, + "ade_m": 1.125946622105056, + "fde_m": 0.2582751406784516, + "max_error_m": 2.925430648852423 + }, + "reported_in_metrics_json": { + "action_mse": 0.7453364133834839, + "ade_m": 1.1259467601776123, + "fde_m": 0.2582751512527466, + "max_error_m": 2.9254307746887207 + }, + "seed": 12 + }, + "diffusion_minus_vlm": { + "action_mse": 0.7449924531511188, + "ade_m": 1.0998419578385121, + "fde_m": 0.21283065215306096, + "max_error_m": 2.87047398021368 + } + }, + "validation": { + "ground_truth_action_exact_match": true, + "ground_truth_action_max_abs_difference": 0.0, + "ground_truth_xy_exact_match": true, + "ground_truth_xy_max_abs_difference": 0.0, + "dataset_target_xy_exact_match": true, + "dataset_target_xy_max_abs_difference": 0.0, + "dataset_history_xy_exact_match": true, + "dataset_history_xy_max_abs_difference": 0.0, + "ground_truth_text_matches_manifest": true, + "vlm_discrete_tokens_exact_match": true, + "action_decode_max_abs_differences": { + "ground_truth": 6.767095328541473e-08, + "vlm": 4.464556582206569e-06, + "diffusion": 2.836197381839156e-06 + }, + "vlm_reported_metrics_match_npz": { + "consistent": true, + "absolute_tolerance": 1e-05, + "absolute_differences": { + "action_mse": 1.1408262084577209e-11, + "ade_m": 1.699166814195241e-09, + "fde_m": 0.0, + "max_error_m": 1.5135175474023654e-09 + } + }, + "diffusion_reported_metrics_match_npz": { + "consistent": true, + "absolute_tolerance": 1e-05, + "absolute_differences": { + "action_mse": 1.5075348858140103e-08, + "ade_m": 1.3807255627718007e-07, + "fde_m": 1.0574294995180367e-08, + "max_error_m": 1.2583629782625394e-07 + } + } + }, + "data_sources": { + "manifest": { + "path": "data/processed/overfit_manifest.jsonl", + "sha256": "3af0fe4e1b1315b83ae7214dd923950c5aa2b5a6a14918bc3f940013c431da97" + }, + "vlm_npz": { + "path": "outputs/stage1_vlm_trajectory/trajectory_logits.npz", + "sha256": "1570bb17b39821b5e73900f6b1e9332dd0ea3f765770fbe6a282799945adf2ab", + "keys": [ + "boundary_masked_trajectory_logits", + "boundary_raw_trajectory_logits", + "generated_action", + "generated_values", + "generated_xy", + "quantized_target_values", + "target_action", + "target_xy", + "teacher_target_values", + "teacher_trajectory_logits" + ] + }, + "vlm_metrics": { + "path": "outputs/stage1_vlm_trajectory/trajectory_logits_metrics.json", + "sha256": "5bd5033ce99b2ba20fb1933d3dc997af7b44183d6f22735de549a4f5c71f88f6" + }, + "diffusion_npz": { + "path": "outputs/stage2_expert/reload_prediction.npz", + "sha256": "21005dbece8064e3a01b5c9e314f576d90b11934e789ed094ad46eee1d160030", + "keys": [ + "history_xy", + "predicted_action", + "predicted_xy", + "target_action", + "target_xy" + ] + }, + "diffusion_metrics": { + "path": "outputs/stage2_expert/stage2_metrics.json", + "sha256": "fc0ee1da1d0bc512f42e4530caabd3f4c738294e905e401a96807abb7bcc7698" + }, + "expert_config": { + "path": "outputs/stage2_expert/expert_config.json", + "sha256": "60b83a95c5883cc65aeda9ff1dd9faa652dcf50fa9968640306ed6305b46a38d" + }, + "trajectory_npz": { + "path": "data/processed/trajectories/sample_00_314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899.npz", + "sha256": "145e480d94c06fd968acbf066cf512595ab2b6d02170ed2f570b4a93ada88336", + "keys": [ + "future_xyz", + "history_xyz", + "source_future_xyz", + "waypoint_indices" + ] + } + }, + "artifacts": { + "comparison_json": "results/complete_comparison/comparison.json", + "waypoints_csv": "results/complete_comparison/waypoints.csv", + "trajectory_plot": "results/complete_comparison/trajectory_comparison.png" + } +} diff --git a/results/complete_comparison/trajectory_comparison.png b/results/complete_comparison/trajectory_comparison.png new file mode 100644 index 0000000..03f7f43 Binary files /dev/null and b/results/complete_comparison/trajectory_comparison.png differ diff --git a/results/complete_comparison/waypoints.csv b/results/complete_comparison/waypoints.csv new file mode 100644 index 0000000..f046c19 --- /dev/null +++ b/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 diff --git a/scripts/build_complete_comparison.py b/scripts/build_complete_comparison.py new file mode 100644 index 0000000..b9294dc --- /dev/null +++ b/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() diff --git a/scripts/download_model.py b/scripts/download_model.py new file mode 100644 index 0000000..415bde4 --- /dev/null +++ b/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() diff --git a/scripts/evaluate_pipeline.py b/scripts/evaluate_pipeline.py new file mode 100644 index 0000000..75c5350 --- /dev/null +++ b/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() diff --git a/scripts/materialize_pai_subset.py b/scripts/materialize_pai_subset.py new file mode 100644 index 0000000..23be009 --- /dev/null +++ b/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() diff --git a/scripts/train_stage1_vlm.py b/scripts/train_stage1_vlm.py new file mode 100644 index 0000000..4483f20 --- /dev/null +++ b/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() diff --git a/scripts/train_stage2_expert.py b/scripts/train_stage2_expert.py new file mode 100644 index 0000000..7bee022 --- /dev/null +++ b/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() diff --git a/scripts/train_vlm_trajectory_logits.py b/scripts/train_vlm_trajectory_logits.py new file mode 100644 index 0000000..eb61160 --- /dev/null +++ b/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("")) + end_id = int(tokenizer.convert_tokens_to_ids(f"")) + 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() diff --git a/scripts/verify_connection.py b/scripts/verify_connection.py new file mode 100644 index 0000000..49b96d7 --- /dev/null +++ b/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() diff --git a/scripts/verify_structured_vlm_handoff.py b/scripts/verify_structured_vlm_handoff.py new file mode 100644 index 0000000..d827f24 --- /dev/null +++ b/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() diff --git a/src/alpamayo_lite/__init__.py b/src/alpamayo_lite/__init__.py new file mode 100644 index 0000000..ee681dc --- /dev/null +++ b/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"] diff --git a/src/alpamayo_lite/action_space.py b/src/alpamayo_lite/action_space.py new file mode 100644 index 0000000..60c0c5c --- /dev/null +++ b/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) diff --git a/src/alpamayo_lite/config.py b/src/alpamayo_lite/config.py new file mode 100644 index 0000000..64b96e9 --- /dev/null +++ b/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, + ) diff --git a/src/alpamayo_lite/data.py b/src/alpamayo_lite/data.py new file mode 100644 index 0000000..effb872 --- /dev/null +++ b/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 diff --git a/src/alpamayo_lite/diffusion.py b/src/alpamayo_lite/diffusion.py new file mode 100644 index 0000000..ae7c001 --- /dev/null +++ b/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 diff --git a/src/alpamayo_lite/evaluation.py b/src/alpamayo_lite/evaluation.py new file mode 100644 index 0000000..a8078c9 --- /dev/null +++ b/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(), + ) diff --git a/src/alpamayo_lite/expert.py b/src/alpamayo_lite/expert.py new file mode 100644 index 0000000..66a3952 --- /dev/null +++ b/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) diff --git a/src/alpamayo_lite/prompts.py b/src/alpamayo_lite/prompts.py new file mode 100644 index 0000000..53608df --- /dev/null +++ b/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" + "\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:" diff --git a/src/alpamayo_lite/runtime.py b/src/alpamayo_lite/runtime.py new file mode 100644 index 0000000..7b6a768 --- /dev/null +++ b/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() diff --git a/src/alpamayo_lite/trajectory_tokens.py b/src/alpamayo_lite/trajectory_tokens.py new file mode 100644 index 0000000..103ded0 --- /dev/null +++ b/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"" + + +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 diff --git a/src/alpamayo_lite/vlm.py b/src/alpamayo_lite/vlm.py new file mode 100644 index 0000000..2a8a16b --- /dev/null +++ b/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 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 diff --git a/tests/test_action_space.py b/tests/test_action_space.py new file mode 100644 index 0000000..3b9241b --- /dev/null +++ b/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) diff --git a/tests/test_complete_comparison.py b/tests/test_complete_comparison.py new file mode 100644 index 0000000..772b575 --- /dev/null +++ b/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, + } diff --git a/tests/test_diffusion.py b/tests/test_diffusion.py new file mode 100644 index 0000000..00a3b6a --- /dev/null +++ b/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 diff --git a/tests/test_expert_contract.py b/tests/test_expert_contract.py new file mode 100644 index 0000000..84e1470 --- /dev/null +++ b/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]) diff --git a/tests/test_trajectory_tokens.py b/tests/test_trajectory_tokens.py new file mode 100644 index 0000000..05d4791 --- /dev/null +++ b/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:])