nvlabs/alpamayo2

Purpose and Entry Points

Alpamayo 2 Super is a Python inference package for a 34B autonomous-vehicle foundation model: a 32B vision-language backbone generates Chain-of-Causation/text-task outputs, and a 2B diffusion action expert samples future ego trajectories. Primary executable entry points are the end-to-end smoke CLI at src/alpamayo2_super/inference_smoke.py, its compatibility alias src/alpamayo2_super/test_inference.py, notebooks under notebooks/, and the explicit two-GPU CFG demo at examples/two_gpu_nav_cfg_demo.py.

Model Composition and Loading

The core model assembly belongs in src/alpamayo2_super/models/alpamayo2_super.py, with the trajectory/action expert implemented by src/alpamayo2_super/models/expert.py and action-conditioning projection in src/alpamayo2_super/models/action_in_proj.py. Configuration and checkpoint/model-ID handling are centralized in src/alpamayo2_super/config.py and src/alpamayo2_super/common/config_utils.py. The runtime accepts the gated Hugging Face model nvidia/Alpamayo2-Super or a local release-native safetensors checkpoint.

AV Input Boundary and Profiles

PhysicalAI-AV dataset retrieval and sample decoding are isolated behind src/alpamayo2_super/load_physical_aiavdataset.py. Manifests such as examples/validation_samples.json identify clips and timestamps; the loader resolves corresponding camera, ego-history, calibration, and future-trajectory data from Hugging Face’s nvidia/PhysicalAI-Autonomous-Vehicles dataset. src/alpamayo2_super/input_profiles.py enforces task-specific camera/frame selections, creating a stable boundary between raw dataset assets and model-ready multimodal inputs.

Prompting and Text Tasks

Conversation/prompt construction is represented by src/alpamayo2_super/chat_template/conversation.py. Task routing and output contracts for reasoning, VQA, grounding, meta-actions, and auto-labeling are concentrated in src/alpamayo2_super/text_tasks.py. The backbone’s generated text is both a standalone result and conditioning/context for trajectory generation; downstream consumers should preserve task profile and decoding metadata alongside outputs.

Trajectory Sampling Pipeline

The action-space contract and trajectory parameterization live in src/alpamayo2_super/action_space/action_space.py and src/alpamayo2_super/action_space/unicycle_accel_curvature.py. Diffusion/flow-matching sampling is separated into src/alpamayo2_super/diffusion/base.py and src/alpamayo2_super/diffusion/flow_matching.py. This stage consumes model conditioning and emits one or more future ego trajectories, which the CLI evaluates against available ground truth (for example, minADE). Geometry transforms used across model/output processing reside in src/alpamayo2_super/geometry/coordinates.py and src/alpamayo2_super/geometry/rotation.py.

Execution, Devices, and Distributed Concerns

Normal CLI/notebook inference uses CUDA PyTorch on Linux and requires enough GPU VRAM for the 34B checkpoint; CUDA 12.x and FlashAttention are material installation/runtime dependencies. Device and distributed primitives are grouped in src/alpamayo2_super/common/distributed.py. The advanced CFG path deliberately splits VLM generation onto one GPU and expert denoising plus guided/unguided KV caches onto another, rather than relying on automatic Hugging Face device mapping (examples/two_gpu_nav_cfg_demo.py).

Outputs and Visualization

Inference produces decoded text, sampled trajectories, metrics, and task/profile metadata; the smoke path serializes JSON plus visualization artifacts. Public plotting APIs are exposed through src/alpamayo2_super/visualization.py, with camera projection, six-camera figures, compact layouts, grounding, and task-specific rendering implemented in src/alpamayo2_super/viz_utils.py. Calibration-dependent overlays are optional but explicitly report whether projection was available.