alpamayo-0.5B

alpamayo-0.5B 是 NVIDIA 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 数据。

架构

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。

单样本 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 将世界坐标平移、旋转到 t0 ego-relative frame,并把单样本 manifest、图像和轨迹写入本地 data/processed/。 原始/处理后的 PAI 数据都被 .gitignore 排除,不应提交到仓库。

训练阶段

Stage 1:自然语言驾驶输出

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 从 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 是上游兼容标记,承载的是单句驾驶指令:

<|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 使用 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 分布上联合训练,所以该分支的轨迹质量只作诊断。

目录

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。以下命令把环境和所有运行产物留在仓库目录内:

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:

hf auth login

运行

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

下载 0.5B VLM,并记录 checkpoint provenance:

python scripts/download_model.py --config configs/local_overfit.yaml

下载并物化一个 PAI 样本,然后验证 VLM KV 与 expert 的形状契约:

python scripts/materialize_pai_subset.py --config configs/local_overfit.yaml
python scripts/verify_connection.py --config configs/local_overfit.yaml

训练 Stage 1 自然语言 VLM:

python scripts/train_stage1_vlm.py \
  --config configs/local_overfit.yaml \
  --steps 80

训练并验证 VLM 的 structured/untruncated trajectory logits:

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:

python scripts/train_stage2_expert.py \
  --config configs/local_overfit.yaml \
  --steps 1000

若 outputs/stage2_expert/checkpoints/latest/ 中已有完整 checkpoint,则精确续训到总计 1,000 steps:

python scripts/train_stage2_expert.py \
  --config configs/local_overfit.yaml \
  --steps 1000 \
  --resume

重新加载 VLM/expert,固定从 seed 12 开始做 16 个 diffusion rollouts:

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:

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:

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 三方比较

文字对比

字段 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 样本的轻量结果,可直接审计完整三方对比:

models/、data/ 和 outputs/ 默认不进入 Git。运行完成后,以下链接仅在本地工作区存在; 不要为了让它们在 GitHub 页面可用而提交 gated data 或 checkpoint 权重。

没有列出或链接任何 checkpoint 权重文件。

测试状态

发布或提交结果前运行:

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 从驾驶指令 规则派生的标签,不是 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 发布。上游归属、修改说明及不随仓库分发的模型/数据 声明见 THIRD_PARTY_NOTICES.md。本地工作区可另行检出 references/alpamayo-upstream/ 进行接口研究,但该目录不会随本仓库分发。

S
Description
NVIDIA Alpamayo 自动驾驶视觉语言动作模型的最小复现。
Readme 463 KiB
Languages
Python 86.8%
Jupyter Notebook 13.2%