# alpamayo-0.5B

`alpamayo-0.5B` 是 [NVIDIA Alpamayo](https://github.com/NVlabs/alpamayo) 的最小复现，
面向单机、单样本的架构与接口验证。
项目使用带 Qwen2.5-0.5B 文本骨干的 VLM 和缩小版 diffusion expert，覆盖以下完整链路：

1. 图像、自然语言和 ego-relative 历史坐标进入 VLM；
2. VLM 生成驾驶指令，也可以在未截断验证分支中生成结构化指令、meta-action 和
   离散轨迹 token；
3. VLM 的 24 层 past-KV 作为 expert 的前缀；
4. flow-matching diffusion 预测连续 action，再积分回未来 XY 位置。

项目目标是用可运行、可训练的最小实现验证连接方式和输出契约，而不是复刻
NVIDIA Alpamayo-R1 的完整规模、训练数据或模型能力。该实现不能用于真实车辆控制。

## 复现范围

| 范围 | 本项目 | 与 NVIDIA Alpamayo 的差异 |
| --- | --- | --- |
| VLM | Qwen2.5-0.5B 文本骨干，完整模型约 0.80B 参数 | 不使用 Alpamayo-R1 原始权重 |
| VLM–expert 连接 | 24 层 past-KV，保持 KV heads 与 head dimension 一致 | 只验证接口和缓存传递 |
| Diffusion expert | 105.6M Qwen2 expert，flow matching，10 步采样 | 使用缩小配置和二维 XY-delta action |
| 数据 | 一个 PAI `train` 事件，单样本过拟合 | 不覆盖多场景训练或泛化评估 |
| 输出 | 驾驶指令、离散轨迹 token、连续未来轨迹 | 不等同于上游完整 action space 与 tokenizer |
| 运行环境 | 单机 MPS、CUDA 或 CPU | 不包含分布式训练配置 |

最小复现包含模型下载、PAI 样本物化、两阶段训练、checkpoint 恢复、reload evaluation、
结构化 VLM cache handoff 和结果比较脚本。仓库不分发模型权重或 gated 数据。

## 架构

```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 |

### Expert 参数量

参数量由 `configs/local_overfit.yaml` 中的 Stage 2 配置实例化
`KVCachedDiffusionExpert` 后直接统计：

| 模块 | 参数量 |
| --- | ---: |
| 24 层 Qwen2 expert | 105,425,408 |
| Action/time projection | 148,736 |
| Action output | 1,026 |
| **合计** | **105,575,170（约 105.6M）** |

全部 `105,575,170` 个参数参与 Stage 2 训练。仅模型权重约占 `201.37 MiB`（FP16）或
`402.74 MiB`（FP32），不包含梯度、优化器状态、VLM 权重和运行时 KV cache。

主配置位于 [`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/ood_reasoning.parquet` 中的 `coc` 字段 |
| 驾驶指令 | “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 输出与数据集驾驶指令逐字一致。

Stage 1 输出到本地 `outputs/stage1_vlm/`，供后续两个分支使用：

- plain-text KV 分支训练现有 diffusion expert；
- structured VLM 分支继续训练轨迹 logits。

### VLM structured 未截断轨迹-token 分支

[`scripts/train_vlm_trajectory_logits.py`](scripts/train_vlm_trajectory_logits.py) 从 Stage 1
checkpoint 继续训练，并添加：

- 128 个连续 token：`<i0>` … `<i127>`；
- `<|cot_start|>`、`<|cot_end|>`；
- `<|meta_action_start|>`、`<|meta_action_end|>`；
- `<|traj_future_start|>`、`<|traj_future_end|>`。

当前 Qwen LM head 已预留 151,936 行，而原 tokenizer 长度为 151,666；扩展后的 151,800
个 token 不需要增大权重矩阵。每个 16 x 2 normalized XY-delta action 被均匀量化到
`[-4, 4]` 的 128 bins，形成 32 个监督 token。

当前结构化 checkpoint 训练 160 steps，loss 从 `14.17349` 降至 `0.0000347825`，最终
训练 accuracy 为 1.0。未截断 greedy generation 生成 32/32 个正确轨迹 token。完整
输出格式如下；`cot` 是上游兼容标记，承载的是单句驾驶指令：

```text
<|cot_start|>Steer left following the temporary lane delineated by traffic cones through the construction zone.<|cot_end|>
<|meta_action_start|>STEER_LEFT<|meta_action_end|>
<|traj_future_start|><i9><i37><i71><i38><i70><i41><i68><i46><i66><i49><i84><i63><i63><i63><i65><i67><i66><i69><i65><i67><i82><i80><i61><i77><i60><i79><i59><i78><i56><i76><i72><i86><|traj_future_end|><|im_end|>
```

另有 source-style handoff 验证：generation 到 `<|traj_future_start|>` 后屏蔽全部 128 个
trajectory logits，并把 exact generation cache 交给 diffusion。这个模式验证了 KV 接口，
但不会把 VLM 离散轨迹 token 作为 diffusion 分支的可见输出。

### Stage 2：KV-conditioned diffusion expert

[`scripts/train_stage2_expert.py`](scripts/train_stage2_expert.py) 使用 Stage 1 plain-text prefix
的 past-KV，训练 105.6M expert 预测 flow-matching vector field。目标 action 是按单样本统计量
标准化的相邻 waypoint XY delta，预测 action 经反标准化和累积求和恢复未来轨迹。

当前产物记录为 1,000 total steps：早期 250-step 权重缺少 optimizer state，之后以 legacy
warm start 运行了剩余 750 steps。因此它是 1,000-step 参数训练记录，但不是严格连续、
optimizer 完整恢复的 1,000-step run。新 checkpoint 已保存 expert、optimizer、step 和 RNG，
后续中断可用 `--resume` 精确恢复。

普通 reload evaluation 仍使用训练 expert 时的旧 plain-text prefix（sequence length 770）。
structured exact-cache handoff 的 prefix 长度为 777；接口验证成功，但旧 expert 没有在这个
新 prefix 分布上联合训练，所以该分支的轨迹质量只作诊断。

## 目录

```text
alpamayo-0.5B/
├── configs/                     # 单样本实验配置
├── scripts/                     # 下载、数据准备、训练、评估和比较入口
├── src/alpamayo_lite/           # VLM、action space、expert、diffusion 实现
├── tests/                       # 单元测试和产物一致性测试
├── results/complete_comparison/ # 仓库中保留的轻量比较 JSON、CSV 和三方图
├── references/alpamayo-upstream/# 可选的本地只读上游 checkout；gitignored
├── models/                      # 本地模型；gitignored
├── data/                        # gated/raw/processed 数据；gitignored
├── outputs/                     # checkpoint 和评估产物；gitignored
└── .venv/                       # Python 3.12 虚拟环境；gitignored
```

## 安装

需要 Python 3.12。以下命令把环境和所有运行产物留在仓库目录内：

```bash
git clone <repository-url> alpamayo-0.5B
cd alpamayo-0.5B
python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e '.[dev]'
```

登录 Hugging Face 时使用交互式认证，不要把 access token 写进配置、脚本、README、日志或
shell history：

```bash
hf auth login
```

## Notebook 复现

[`notebooks/alpamayo_0_5b_minimal_reproduction.ipynb`](notebooks/alpamayo_0_5b_minimal_reproduction.ipynb)
提供从环境安装到实验结果的顺序执行流程，包含论文链接、架构说明、权重与数据准备、连接
验证、两阶段训练、结构化轨迹分支、checkpoint 重载和三方结果比较。

安装并启动：

```bash
python -m pip install -e '.[notebook]'
python -m jupyter lab notebooks/alpamayo_0_5b_minimal_reproduction.ipynb
```

Notebook 自动选择运行模式：本地存在完整产物时使用 `reuse`，全新 clone 使用 `full`。
设置 `ALPAMAYO_PROFILE=smoke` 可执行低步数连通性检查。外部模型、manifest 和 checkpoint
可通过 Notebook 中列出的环境变量直接导入。

## 运行

所有命令均从仓库根目录执行。

下载 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 |
| --- | --- | --- | ---: |
| 驾驶指令 | 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`（由驾驶指令规则派生） | `STEER_LEFT` | 是 |
| Diffusion conditioning text | 同上方驾驶指令 | 同上方驾驶指令 | 是 |

## 本地产物

仓库保留以下不含权重或 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) 从驾驶指令
  规则派生的标签，不是 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/` 进行接口研究，但该目录不会随本仓库分发。