{
"cells": [
{
"cell_type": "markdown",
"id": "title-and-toc",
"metadata": {},
"source": [
"# alpamayo-0.5B：NVIDIA Alpamayo 最小复现\n",
"\n",
"本 Notebook 从环境安装开始，依次完成模型与数据准备、VLM–expert 连接验证、两阶段训练、结构化轨迹分支、重载评估和结果可视化。\n",
"\n",
"## 目录\n",
"\n",
"1. [论文与复现范围](#论文与复现范围)\n",
"2. [架构介绍](#架构介绍)\n",
"3. [环境安装](#环境安装)\n",
"4. [运行配置](#运行配置)\n",
"5. [权重与数据准备](#权重与数据准备)\n",
"6. [样本检查](#样本检查)\n",
"7. [连接验证](#连接验证)\n",
"8. [Stage 1：驾驶指令 VLM](#stage-1驾驶指令-vlm)\n",
"9. [结构化轨迹 token 分支](#结构化轨迹-token-分支)\n",
"10. [Stage 2：Diffusion expert](#stage-2diffusion-expert)\n",
"11. [评估与实验结果](#评估与实验结果)\n",
"12. [复现检查清单](#复现检查清单)\n"
]
},
{
"cell_type": "markdown",
"id": "paper-and-scope",
"metadata": {},
"source": [
"\n",
"## 1. 论文与复现范围\n",
"\n",
"- 论文：[Alpamayo-R1: Bridging Reasoning and Action Prediction for Generalizable Autonomous Driving in the Long Tail](https://arxiv.org/abs/2511.00088)\n",
"- NVIDIA 上游仓库：[NVlabs/alpamayo](https://github.com/NVlabs/alpamayo)\n",
"- NVIDIA 模型页：[nvidia/Alpamayo-R1-10B](https://huggingface.co/nvidia/Alpamayo-R1-10B)\n",
"- 本项目数据源：[nvidia/PhysicalAI-Autonomous-Vehicles](https://huggingface.co/datasets/nvidia/PhysicalAI-Autonomous-Vehicles)\n",
"\n",
"本项目复现的是 VLM、past-KV、diffusion expert 与连续轨迹解码之间的核心连接方式。它使用 Qwen2.5-0.5B 文本骨干、约 0.80B 的完整 VLM 和 105.6M diffusion expert，在一个 PAI 样本上完成可训练、可保存、可恢复、可评估的最小闭环。它不使用 Alpamayo-R1 原始权重，也不代表完整模型规模、训练数据或驾驶能力。\n"
]
},
{
"cell_type": "markdown",
"id": "architecture",
"metadata": {},
"source": [
"\n",
"## 2. 架构介绍\n",
"\n",
"```text\n",
"前视图像 ─┐\n",
"驾驶问题 ─┼─> LLaVA VLM (CLIP + Qwen2.5-0.5B) ─┬─> 驾驶指令 / 结构化轨迹 token\n",
"历史 XY ─┘ └─> 24 层 past-KV\n",
" │\n",
"noisy XY-delta + diffusion timestep ─> 105.6M Qwen2 expert <──┘\n",
" │\n",
" vector field\n",
" │\n",
" 10-step denoising\n",
" │\n",
" 16 × 2 future XY\n",
"```\n",
"\n",
"Expert 使用 24 层、hidden size 512、8 个 query heads、2 个 KV heads 和 head dimension 64。精确参数量为 105,575,170，其中 Qwen2 主体 105,425,408、action/time projection 148,736、输出层 1,026。\n"
]
},
{
"cell_type": "markdown",
"id": "environment-guide",
"metadata": {},
"source": [
"\n",
"## 3. 环境安装\n",
"\n",
"推荐使用 Python 3.12。若尚未启动 Notebook，可在仓库根目录执行：\n",
"\n",
"```bash\n",
"python3.12 -m venv .venv\n",
"source .venv/bin/activate\n",
"python -m pip install --upgrade pip\n",
"python -m pip install -e '.[notebook]'\n",
"python -m jupyter lab notebooks/alpamayo_0_5b_minimal_reproduction.ipynb\n",
"```\n",
"\n",
"下面的单元格会定位仓库根目录，并在当前 kernel 尚未安装项目时执行 editable install。\n"
]
},
{
"cell_type": "code",
"id": "environment-setup",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from __future__ import annotations\n",
"\n",
"import importlib.util\n",
"import os\n",
"import platform\n",
"import subprocess\n",
"import sys\n",
"from pathlib import Path\n",
"\n",
"def find_repo_root(start: Path) -> Path:\n",
" for candidate in (start.resolve(), *start.resolve().parents):\n",
" if (candidate / 'pyproject.toml').exists() and (candidate / 'scripts').is_dir():\n",
" return candidate\n",
" raise FileNotFoundError('请从 alpamayo-0.5B 仓库内启动此 Notebook')\n",
"\n",
"ROOT = find_repo_root(Path.cwd())\n",
"os.chdir(ROOT)\n",
"if importlib.util.find_spec('alpamayo_lite') is None:\n",
" subprocess.run([sys.executable, '-m', 'pip', 'install', '-e', str(ROOT)], check=True)\n",
"\n",
"print('Repository:', ROOT)\n",
"print('Python:', sys.version.split()[0])\n",
"print('Platform:', platform.platform())\n"
]
},
{
"cell_type": "markdown",
"id": "runtime-guide",
"metadata": {},
"source": [
"\n",
"## 4. 运行配置\n",
"\n",
"支持三种执行方式：\n",
"\n",
"- `reuse`：复用本地已有权重、数据和 checkpoint；当前工作区具备完整产物时自动选择。\n",
"- `full`：全新工作区默认选择，执行 80-step Stage 1、160-step structured VLM 和 1,000-step expert。\n",
"- `smoke`：每个训练阶段只运行少量步数，用于检查代码、显存和接口；不用于复现发布指标。\n",
"\n",
"可在启动 Jupyter 前设置以下环境变量导入本地资源：\n",
"\n",
"```bash\n",
"export ALPAMAYO_PROFILE=reuse # reuse | full | smoke\n",
"export ALPAMAYO_MODEL_DIR=/path/to/qwen2.5_0.5B_vl\n",
"export ALPAMAYO_MANIFEST=/path/to/overfit_manifest.jsonl\n",
"export ALPAMAYO_PAI_RAW_DIR=/path/to/pai/raw\n",
"export ALPAMAYO_STAGE1_DIR=/path/to/stage1_vlm\n",
"export ALPAMAYO_STRUCTURED_DIR=/path/to/stage1_vlm_trajectory\n",
"export ALPAMAYO_STAGE2_DIR=/path/to/stage2_expert\n",
"```\n",
"\n",
"若自动下载 PAI 数据，Hugging Face 账号必须先接受 gated dataset 许可，并通过 `hf auth login` 或 `HF_TOKEN` 提供凭据。\n"
]
},
{
"cell_type": "code",
"id": "runtime-config",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"import torch\n",
"import yaml\n",
"from IPython.display import display\n",
"\n",
"BASE_CONFIG = ROOT / 'configs/local_overfit.yaml'\n",
"NOTEBOOK_OUTPUT = ROOT / 'outputs/notebook'\n",
"NOTEBOOK_OUTPUT.mkdir(parents=True, exist_ok=True)\n",
"\n",
"default_stage1 = ROOT / 'outputs/stage1_vlm'\n",
"default_stage2 = ROOT / 'outputs/stage2_expert'\n",
"has_reusable_outputs = (default_stage1 / 'stage1_metrics.json').exists() and (default_stage2 / 'expert_config.json').exists()\n",
"PROFILE = os.environ.get('ALPAMAYO_PROFILE', 'reuse' if has_reusable_outputs else 'full').lower()\n",
"if PROFILE not in {'reuse', 'full', 'smoke'}:\n",
" raise ValueError('ALPAMAYO_PROFILE must be reuse, full, or smoke')\n",
"\n",
"profile_steps = {\n",
" 'reuse': {'stage1': 80, 'structured': 160, 'stage2': 1000, 'rollouts': 16},\n",
" 'full': {'stage1': 80, 'structured': 160, 'stage2': 1000, 'rollouts': 16},\n",
" 'smoke': {'stage1': 1, 'structured': 1, 'stage2': 1, 'rollouts': 2},\n",
"}[PROFILE]\n",
"\n",
"raw = yaml.safe_load(BASE_CONFIG.read_text(encoding='utf-8'))\n",
"raw['paths']['model_dir'] = os.environ.get('ALPAMAYO_MODEL_DIR', str(ROOT / 'models/qwen2.5_0.5B_vl'))\n",
"raw['paths']['pai_raw_dir'] = os.environ.get('ALPAMAYO_PAI_RAW_DIR', str(ROOT / 'data/raw/pai'))\n",
"raw['paths']['manifest'] = os.environ.get('ALPAMAYO_MANIFEST', str(ROOT / 'data/processed/overfit_manifest.jsonl'))\n",
"raw['paths']['stage1_dir'] = os.environ.get(\n",
" 'ALPAMAYO_STAGE1_DIR',\n",
" str(default_stage1 if PROFILE == 'reuse' else NOTEBOOK_OUTPUT / 'stage1_vlm'),\n",
")\n",
"raw['paths']['stage2_dir'] = os.environ.get(\n",
" 'ALPAMAYO_STAGE2_DIR',\n",
" str(default_stage2 if PROFILE == 'reuse' else NOTEBOOK_OUTPUT / 'stage2_expert'),\n",
")\n",
"raw['stage1']['steps'] = profile_steps['stage1']\n",
"raw['stage2']['steps'] = profile_steps['stage2']\n",
"RUNTIME_CONFIG = NOTEBOOK_OUTPUT / 'runtime_config.yaml'\n",
"RUNTIME_CONFIG.write_text(yaml.safe_dump(raw, sort_keys=False), encoding='utf-8')\n",
"\n",
"MODEL_DIR = Path(raw['paths']['model_dir']).expanduser().resolve()\n",
"MANIFEST = Path(raw['paths']['manifest']).expanduser().resolve()\n",
"STAGE1_DIR = Path(raw['paths']['stage1_dir']).expanduser().resolve()\n",
"STAGE2_DIR = Path(raw['paths']['stage2_dir']).expanduser().resolve()\n",
"default_structured = ROOT / 'outputs/stage1_vlm_trajectory'\n",
"STRUCTURED_DIR = Path(os.environ.get(\n",
" 'ALPAMAYO_STRUCTURED_DIR',\n",
" str(default_structured if PROFILE == 'reuse' else NOTEBOOK_OUTPUT / 'stage1_vlm_trajectory'),\n",
")).expanduser().resolve()\n",
"HANDOFF_DIR = (ROOT / 'outputs/structured_vlm_handoff') if PROFILE == 'reuse' else (NOTEBOOK_OUTPUT / 'structured_vlm_handoff')\n",
"COMPARISON_DIR = (ROOT / 'outputs/complete_comparison') if PROFILE == 'reuse' else (NOTEBOOK_OUTPUT / 'complete_comparison')\n",
"\n",
"def run_script(name: str, *args: object) -> None:\n",
" command = [sys.executable, str(ROOT / 'scripts' / name), *map(str, args)]\n",
" print('$', ' '.join(command))\n",
" subprocess.run(command, cwd=ROOT, check=True)\n",
"\n",
"device = 'cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu')\n",
"print(json.dumps({\n",
" 'profile': PROFILE,\n",
" 'steps': profile_steps,\n",
" 'device': device,\n",
" 'model_dir': str(MODEL_DIR),\n",
" 'manifest': str(MANIFEST),\n",
" 'stage1_dir': str(STAGE1_DIR),\n",
" 'stage2_dir': str(STAGE2_DIR),\n",
"}, indent=2, ensure_ascii=False))\n"
]
},
{
"cell_type": "markdown",
"id": "resources-guide",
"metadata": {},
"source": [
"\n",
"## 5. 权重与数据准备\n",
"\n",
"模型目录存在时直接复用，否则从 `thkim0305/qwen2.5_0.5B_vl` 下载固定 revision。Manifest 存在时直接复用，否则从 gated PAI 数据集物化单个样本。\n"
]
},
{
"cell_type": "code",
"id": "prepare-resources",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"if not (MODEL_DIR / 'model.safetensors').exists():\n",
" run_script('download_model.py', '--config', RUNTIME_CONFIG)\n",
"else:\n",
" print('Reuse model:', MODEL_DIR)\n",
"\n",
"if not MANIFEST.exists():\n",
" if not (os.environ.get('HF_TOKEN') or (Path.home() / '.cache/huggingface/token').exists()):\n",
" raise RuntimeError(\n",
" 'PAI manifest 不存在。请先接受数据集许可并运行 hf auth login，'\n",
" '或设置 ALPAMAYO_MANIFEST 指向本地 manifest。'\n",
" )\n",
" run_script('materialize_pai_subset.py', '--config', RUNTIME_CONFIG)\n",
"else:\n",
" print('Reuse manifest:', MANIFEST)\n",
"\n",
"assert (MODEL_DIR / 'config.json').exists(), MODEL_DIR\n",
"assert MANIFEST.exists(), MANIFEST\n"
]
},
{
"cell_type": "markdown",
"id": "sample-guide",
"metadata": {},
"source": [
"\n",
"## 6. 样本检查\n",
"\n",
"显示单样本图像、历史轨迹和未来监督轨迹，确认 manifest 中的相对路径可以从仓库根目录解析。\n"
]
},
{
"cell_type": "code",
"id": "inspect-sample",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"from PIL import Image\n",
"\n",
"sample = json.loads(MANIFEST.read_text(encoding='utf-8').splitlines()[0])\n",
"image_path = Path(sample['image_path'])\n",
"trajectory_path = Path(sample['trajectory_path'])\n",
"image_path = image_path if image_path.is_absolute() else ROOT / image_path\n",
"trajectory_path = trajectory_path if trajectory_path.is_absolute() else ROOT / trajectory_path\n",
"trajectory = np.load(trajectory_path)\n",
"history_xy = trajectory['history_xyz'][:, :2]\n",
"future_xy = trajectory['future_xyz'][:, :2]\n",
"\n",
"fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))\n",
"axes[0].imshow(Image.open(image_path).convert('RGB'))\n",
"axes[0].set_title(f\"{sample['camera']} @ {sample['t0_us']} us\")\n",
"axes[0].axis('off')\n",
"axes[1].plot(history_xy[:, 0], history_xy[:, 1], 'o-', label='history')\n",
"axes[1].plot(future_xy[:, 0], future_xy[:, 1], 'o-', label='future target')\n",
"axes[1].axis('equal')\n",
"axes[1].grid(True, alpha=0.3)\n",
"axes[1].legend()\n",
"axes[1].set_title(sample['reasoning'])\n",
"plt.show()\n"
]
},
{
"cell_type": "markdown",
"id": "connection-guide",
"metadata": {},
"source": [
"\n",
"## 7. 连接验证\n",
"\n",
"该步骤加载 VLM，提取 24 层 past-KV，实例化 105.6M expert，并检查层数、KV heads、head dimension 和输出形状。\n"
]
},
{
"cell_type": "code",
"id": "verify-connection",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"run_script('verify_connection.py', '--config', RUNTIME_CONFIG)\n",
"connection = json.loads((ROOT / 'outputs/connection_test.json').read_text(encoding='utf-8'))\n",
"connection\n"
]
},
{
"cell_type": "markdown",
"id": "stage1-guide",
"metadata": {},
"source": [
"\n",
"## 8. Stage 1：驾驶指令 VLM\n",
"\n",
"Stage 1 冻结 vision tower、multimodal projector 和 LM head，只训练最后 8 个 Qwen2 language layers 与 final norm。`reuse` 模式在 checkpoint 存在时跳过训练。\n"
]
},
{
"cell_type": "code",
"id": "train-stage1",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"stage1_metrics_path = STAGE1_DIR / 'stage1_metrics.json'\n",
"if PROFILE == 'reuse' and stage1_metrics_path.exists():\n",
" print('Reuse Stage 1:', STAGE1_DIR)\n",
"else:\n",
" run_script('train_stage1_vlm.py', '--config', RUNTIME_CONFIG, '--steps', profile_steps['stage1'])\n",
"stage1_metrics = json.loads(stage1_metrics_path.read_text(encoding='utf-8'))\n",
"stage1_metrics\n"
]
},
{
"cell_type": "markdown",
"id": "structured-guide",
"metadata": {},
"source": [
"\n",
"## 9. 结构化轨迹 token 分支\n",
"\n",
"该分支从 Stage 1 checkpoint 继续训练，加入 128 个离散 action token、meta-action 和轨迹边界标记。完整配置训练 160 步并生成 32 个轨迹 token。`smoke` 模式只验证训练入口，不保证结构化输出收敛。\n"
]
},
{
"cell_type": "code",
"id": "train-structured",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"structured_metrics_path = STRUCTURED_DIR / 'trajectory_logits_metrics.json'\n",
"if PROFILE == 'smoke':\n",
" structured_metrics = None\n",
" print('Smoke profile: skip structured generation, which requires convergence before decoding.')\n",
"elif PROFILE == 'reuse' and structured_metrics_path.exists():\n",
" print('Reuse structured VLM:', STRUCTURED_DIR)\n",
"else:\n",
" run_script(\n",
" 'train_vlm_trajectory_logits.py',\n",
" '--config', RUNTIME_CONFIG,\n",
" '--output-dir', STRUCTURED_DIR,\n",
" '--steps', profile_steps['structured'],\n",
" '--trajectory-vocab-size', 128,\n",
" '--action-min', -4,\n",
" '--action-max', 4,\n",
" '--max-new-tokens', 96,\n",
" )\n",
"if PROFILE != 'smoke':\n",
" structured_metrics = json.loads(structured_metrics_path.read_text(encoding='utf-8'))\n",
"structured_metrics\n"
]
},
{
"cell_type": "markdown",
"id": "stage2-guide",
"metadata": {},
"source": [
"\n",
"## 10. Stage 2：Diffusion expert\n",
"\n",
"Stage 2 冻结 Stage 1 VLM，使用其 prefix cache 训练 105,575,170 参数的 expert。目标为标准化的相邻 waypoint XY delta，flow matching 输出经反标准化和累积求和还原未来轨迹。\n"
]
},
{
"cell_type": "code",
"id": "train-stage2",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"stage2_metrics_path = STAGE2_DIR / 'stage2_metrics.json'\n",
"if PROFILE == 'reuse' and stage2_metrics_path.exists():\n",
" print('Reuse Stage 2:', STAGE2_DIR)\n",
"else:\n",
" run_script('train_stage2_expert.py', '--config', RUNTIME_CONFIG, '--steps', profile_steps['stage2'])\n",
"stage2_metrics = json.loads(stage2_metrics_path.read_text(encoding='utf-8'))\n",
"expert_config = json.loads((STAGE2_DIR / 'expert_config.json').read_text(encoding='utf-8'))\n",
"print('Expert parameters:', stage2_metrics['expert_parameters'])\n",
"stage2_metrics\n"
]
},
{
"cell_type": "markdown",
"id": "evaluation-guide",
"metadata": {},
"source": [
"\n",
"## 11. 评估与实验结果\n",
"\n",
"完整模式执行 16 个 diffusion rollouts、structured exact-cache handoff 和 GT/VLM/diffusion 三方比较。`smoke` 模式只执行少量 rollout，并展示仓库中保存的发布结果。\n"
]
},
{
"cell_type": "code",
"id": "reload-evaluation",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"evaluation_path = NOTEBOOK_OUTPUT / 'reload_evaluation.json'\n",
"run_script(\n",
" 'evaluate_pipeline.py',\n",
" '--config', RUNTIME_CONFIG,\n",
" '--stage2-dir', STAGE2_DIR,\n",
" '--num-rollouts', profile_steps['rollouts'],\n",
" '--base-seed', 12,\n",
" '--output', evaluation_path,\n",
")\n",
"evaluation = json.loads(evaluation_path.read_text(encoding='utf-8'))\n",
"evaluation\n"
]
},
{
"cell_type": "code",
"id": "comparison-results",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"if PROFILE != 'smoke':\n",
" run_script(\n",
" 'verify_structured_vlm_handoff.py',\n",
" '--config', RUNTIME_CONFIG,\n",
" '--vlm-dir', STRUCTURED_DIR,\n",
" '--stage2-dir', STAGE2_DIR,\n",
" '--output-dir', HANDOFF_DIR,\n",
" )\n",
" run_script(\n",
" 'build_complete_comparison.py',\n",
" '--vlm-npz', STRUCTURED_DIR / 'trajectory_logits.npz',\n",
" '--vlm-metrics', STRUCTURED_DIR / 'trajectory_logits_metrics.json',\n",
" '--diffusion-npz', STAGE2_DIR / 'reload_prediction.npz',\n",
" '--diffusion-metrics', STAGE2_DIR / 'stage2_metrics.json',\n",
" '--expert-config', STAGE2_DIR / 'expert_config.json',\n",
" '--manifest', MANIFEST,\n",
" '--diffusion-seed', 12,\n",
" '--output-dir', COMPARISON_DIR,\n",
" )\n",
" result_dir = COMPARISON_DIR\n",
"else:\n",
" print('Smoke profile: display checked-in comparison artifacts.')\n",
" result_dir = ROOT / 'results/complete_comparison'\n",
"\n",
"comparison = json.loads((result_dir / 'comparison.json').read_text(encoding='utf-8'))\n",
"metrics = comparison['metrics']\n",
"print(json.dumps({\n",
" 'vlm': metrics['vlm']['computed_from_npz'],\n",
" 'diffusion': metrics['diffusion']['computed_from_npz'],\n",
" 'trajectory_token_accuracy': metrics['vlm']['trajectory_token_accuracy'],\n",
"}, indent=2))\n",
"display(Image.open(result_dir / 'trajectory_comparison.png'))\n"
]
},
{
"cell_type": "markdown",
"id": "published-results",
"metadata": {},
"source": [
"### 仓库当前发布结果\n",
"\n",
"| 输出 | Action MSE | ADE | FDE | Max error |\n",
"| --- | ---: | ---: | ---: | ---: |\n",
"| VLM trajectory tokens | 0.0003439452 | 0.026104664 m | 0.045444489 m | 0.054956669 m |\n",
"| Diffusion（seed 12） | 0.7453363983 | 1.125946622 m | 0.258275141 m | 2.925430649 m |\n",
"\n",
"VLM 的 32 个生成 token 与量化目标全部匹配；其非零误差来自 128-bin 反量化。Diffusion 结果依赖采样 seed，仓库记录的 16-seed 平均 ADE 为 1.530302633 m。\n"
]
},
{
"cell_type": "markdown",
"id": "reproduction-checklist",
"metadata": {},
"source": [
"\n",
"## 12. 复现检查清单\n",
"\n",
"顺序执行全部单元格后，应满足：\n",
"\n",
"- 模型目录包含 `model.safetensors` 与 `download_provenance.json`。\n",
"- Manifest 可加载一张前视图像和一份包含 history/future 的 NPZ。\n",
"- `verify_connection.py` 返回 24 层 KV，并报告 105,575,170 个 expert 参数。\n",
"- Stage 1、structured VLM 与 Stage 2 均生成指标 JSON 和权重文件。\n",
"- Reload evaluation 可重复加载 checkpoint 并输出 rollout 指标。\n",
"- 完整模式生成 handoff 指标、比较 JSON、waypoint CSV 和三方轨迹图。\n",
"\n",
"模型输出和本项目不能用于真实车辆控制或安全决策。不要提交 Hugging Face token、gated PAI 数据、下载权重或训练 checkpoint。\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (alpamayo-0.5B)",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}