Alpamayo 2 Super
A 34-billion parameter foundation model for autonomous vehicles. It pairs a 32B vision-language backbone with a 2B diffusion action expert to produce reasoning text and sampled future trajectories from multi-camera driving video.
The two-stage idea in one sentence
The VLM reasons in text (Chain-of-Causation) and emits a future-trajectory anchor token; the diffusion expert then denoises a continuous action conditioned on the VLM's KV cache to produce the final 64-waypoint trajectory. Text and trajectory share one backbone, but the continuous action comes from a separate flow-matching head.
Architecture at a glance
A pure VLM can reason about a scene but quantizing trajectories into discrete tokens loses geometric precision. A pure diffusion head can output smooth trajectories but can't explain why. Alpamayo 2 Super gets both: interpretable Chain-of-Causation text and a high-precision continuous trajectory, sharing one set of visual features.
01The Data Contract
Every inference call starts by loading one clip from nvidia/PhysicalAI-Autonomous-Vehicles and reshaping it into a fixed ego-centric tensor contract. Understanding this contract is the key to everything downstream.
What one sample contains
Everything is transformed into the ego frame at t₀: xyz_local = R_t₀⁻¹ · (xyz_world − xyz_t₀). So history ends at the origin and the future is expressed relative to where the car is now. This is done in load_physical_aiavdataset() before the model ever sees the data.
Temporal layout — the 4 context frames
Images are sampled at t₀ − 0.3 s, t₀ − 0.2 s, t₀ − 0.1 s, t₀. The final frame index is always t₀ (ego_t0_frame_idx = num_frames − 1). Crucially, the auto-labeling task never sees future camera frames — only context through t₀, even though the offline teacher that made its labels could look ahead.
02Trajectory Tokenization
Continuous xyz waypoints can't go straight into a VLM. Alpamayo 2 Super uses a delta tokenizer that quantizes per-step deltas into discrete bins, plus a set of special tokens that frame each modality in the conversation.
Delta encoding (history)
History is encoded as deltas ending at the origin. With pad_origin_at_beginning=True, a zero origin is prepended so 16 waypoints → 16 deltas. Each delta is normalized to [0, num_bins−1] and flattened: 16 × 3 axes = 48 tokens.
Two vocabularies
<i0>…<i999><i1000>…<i3999>The future span is not filled with real tokens — it's left as the <|traj_future_start|> anchor. The discrete future tokens are only a training target; the real inference future comes from the diffusion expert.
The conversation template
A system message sets the persona; the user message interleaves camera-labeled images and the history-trajectory span; the assistant message is where the model writes CoC and the future anchor. The order is configurable via components_order.
03The VLM Backbone
A 32B Qwen-VL model does the heavy lifting: it prefills the prompt once, then autoregressively generates the Chain-of-Causation text and stops at the <|traj_future_start|> anchor — which hands control to the expert.
Shared-prefill generation
For trajectory sampling you may want n trajectories per prompt. Rather than re-prefilling each time, the prompt is prefilled once, the KV cache is batch_repeat_interleave'd, and generation runs across all copies in parallel. This is the _generate_with_shared_prefill path.
Masking the trajectory tokens
While the VLM generates CoC text, it must never emit a discrete trajectory token (those belong to training only). A logits processor sets the entire trajectory-token span to −inf every step.
Stopping at the anchor
The real EOS for this generation is not the text EOS — it's <|traj_future_start|>. StopAfterEOS tracks per-sequence completion, then replace_padding_after_eos cleans up trailing pad so the expert sees a clean boundary.
The expert then starts reading the KV cache right after this anchor offset.
04The Action Space — a Unicycle
The expert doesn't output waypoints directly. It outputs normalized acceleration and curvature per waypoint — 64 × 2 values — which are integrated through a unicycle kinematic model to recover xyz. This keeps trajectories physically realizable.
The control inputs
a·accel_std + accel_mean)κ·curv_std + curv_mean)From actions to trajectory — the forward map
Given an initial velocity v₀ (estimated from history), the model integrates: velocity from acceleration, heading from curvature × velocity, and position from velocity × heading. This is the action_to_traj path used at inference.
During training the ground-truth future trajectory is converted into actions via traj_to_action (solving constrained least-squares for smooth accel/curvature), then the expert learns to denoise those actions. So the expert lives in action space, not xyz space.
05Flow Matching & the Expert
The expert is a 2B transformer trained with flow matching (a continuous-time cousin of diffusion). At inference it integrates a learned vector field from pure noise to a clean action using fixed-step Euler updates.
Training target
Sample a timestep t ~ Beta(1.5, 1), build noisy action x_t = t·x + (1−t)·noise, and regress the network to the constant-velocity target v = x − noise.
Inference sampling
Start from x ~ N(0, I) at t=0 and take num_inference_steps (default 10) Euler steps to t=1: x ← x + dt · v_θ(x, t). That's it — no DDPM noise injection, just straight-line flow integration.
What's inside one step_fn call?
Each Euler step runs the full expert transformer with the VLM's KV cache as a prefix. The noisy action + timestep are projected to tokens, attend to the cached prompt, and a linear head reads out the predicted velocity.
The expert reuses the VLM's prefilled KV cache but crops it back to prefill_seq_len after each step so the cache doesn't accumulate expert tokens across diffusion iterations. Every step sees the same prompt prefix.
06The Full Inference Pipeline
Click through each stage of sample_trajectories_from_data() — the single entry point that turns one PhysicalAI-AV clip into a predicted trajectory and Chain-of-Causation text.
load_physical_aiavdataset() fetches the 7-camera ring + egomotion, transforms to ego-t₀ frame. select_task_input(data, "trajectory") drops to 6 cameras × 4 frames.helper.prepare_model_inputs() builds the chat messages (system + images + history span + prompt), applies the chat template, and runs the Qwen processor to get input_ids + pixel_values.fuse_traj_tokens() replaces the history placeholder span with real encoded history tokens (48). The future span stays as placeholders (training-only).batch_repeat_interleave it for n = num_traj_samples × num_traj_sets parallel generations.top_p=0.98, temp=0.6. The trajectory-token logits are masked to −inf; text EOS is masked; generation stops at <|traj_future_start|>.find_eos_offset() locates where each sequence emitted the anchor; build_expert_pos_ids_and_attn_mask() builds MRoPE position IDs and a 4D attention mask so expert tokens attend only to the right prefix.action_in_proj (Fourier + MLP) → expert transformer with VLM KV prefix → action_out_proj → velocity. Cache cropped back each step.action_space.action_to_traj() integrates the sampled (accel, curvature) through the unicycle model from the last history pose → pred_xyz, pred_rot shaped [B, n_sets, n_samples, 64, 3].extract_text_tokens() decodes the CoC text per trajectory. plot_inference_result() renders the 6-camera grid + BEV with predicted vs. ground-truth trajectory and computes minADE/minFDE.Sample output (real model run)
From the checked-in outputs/blog.json — clip 030c760c… at t0_us=5100000:
07How It's Trained
Training is split across two objectives that share the VLM backbone: a discrete next-token loss on text + future-trajectory tokens, and a flow-matching MSE loss on continuous actions in the expert.
VLM forward — two masked losses
The training forward() fuses real history and future tokens into the prompt, runs the VLM, then splits the loss: a future_traj cross-entropy over the future-token span, and an others cross-entropy over everything else (CoC text, special tokens). Each is weighted via config.loss_weights.
Expert forward — flow-matching loss
The expert converts the GT future into actions (traj_to_action), adds noise via construct_training_data, runs the denoiser, and regresses to the velocity target x − noise.
Are VLM and expert co-trained?
Controlled by config.cotrain_expert_vlm. When False (the release default), the VLM is frozen (requires_grad_(False)) and only the expert trains against the cached prompt features. When True, gradients flow through both.
The discrete future-trajectory tokens act as an auxiliary teaching signal for the VLM — they force the backbone to understand future motion — while the actual deployed trajectory comes from the expert's continuous actions. So the VLM "knows" trajectories in two representations, but only the continuous one ships.
08Beyond Trajectories — Text Tasks
The same VLM backbone drives three pure-text AV tasks (no expert involved). Each uses a different prompt template and camera profile, all routed through generate_text().
Meta-action
Generates CoC text followed by a structured block of Longitudinal, Lateral, Lane actions. Stops at <|traj_future_start|>; split_cot_and_meta_action() separates the two halves.
Auto-labeling
Outputs a 4-field JSON: critical_components_analysis, ego_vehicle_motion_analysis, trajectory_analysis, chain_of_causation. Conditions on 4 context frames through t₀ + a future ego trajectory.
VQA & grounding
Free-form scene Q&A plus bounding-box grounding, using the no-special-token prompt format. Uses the only profile that includes camera_rear_tele (id 4).
Shared machinery
All three tasks reuse the same logits processor (mask trajectory tokens), the same chat-template pipeline, and the same extract_text_tokens() decoder. The only differences are the components_prompt list and whether a future trajectory is fed in as context.
This is what makes Alpamayo 2 Super a foundation model: a single set of weights serves trajectory prediction, scene reasoning, auto-labeling, and visual grounding — all from the same 32B backbone, with the 2B expert bolted on only when continuous motion is needed.
09Code Map — Where Things Live
A reference for navigating the repository. Each entry is the file (and key symbol) you'd open to understand or modify that part of the system.
Model definition
models/alpamayo2_super.py — Alpamayo2Supermodels/expert.py — ExpertModelmodels/action_in_proj.py — Fourier + MLPconfig.py — Alpamayo2SuperConfigDenoising & actions
diffusion/flow_matching.pyaction_space/unicycle_accel_curvature.pymodels/delta_tokenizer.pymodels/token_utils.pyInference entry points
inference_smoke.pyload_physical_aiavdataset.pyinput_profiles.pytext_tasks.pyPrompting & viz
chat_template/conversation.pyhelper.py — prepare_model_inputsvisualization.py + viz_utils.pymodels/expert_utils.py