Reading the Alpamayo-R1 code

Why read the code

The concepts page describes Alpamayo-R1 from its paper. A paper says what the authors meant to build; the code says what runs. The released repository, github.com/NVlabs/alpamayo, is small (about 4,000 lines of Python for inference) and readable in an afternoon, and none of what follows needs the gated 22 GB weights or a GPU. It is Apache-2.0 licensed; the weights have their own licence.

This page follows one call of the test script through the code. We read commit 11a0e01 (the repository is in limited maintenance and points to newer releases, so details may move). Where we quote the code we name the file; the three runnable pieces at the end are our own small re-implementations, written to match what the code does and tested on a CPU.

A map of the repository

1. What goes in

load_physical_aiavdataset returns, for one moment in a clip:

• Images from four cameras (front wide, front telephoto, cross left, cross right), four frames each: 16 images.
• Ego history: 16 past positions and rotations at 10 Hz, 1.6 s.
• Ego future: 64 positions at 10 Hz, 6.4 s, used only to score the prediction.

These are the code's defaults. The paper describes a 2 s history; the released loader's default is 1.6 s. Small differences like this are why it is worth reading both.

2. The prompt

helper.create_message builds a chat in the backbone's format. The system message is "You are a driving assistant that generates safe and accurate actions." The user message holds the 16 images, then 48 placeholder tokens between <|traj_history_start|> and <|traj_history_end|>, then the instruction "output the chain-of-thought reasoning of the driving process, then output the future trajectory." The assistant's turn is opened with <|cot_start|>, so generation begins inside the reasoning.

The 48 placeholders are 16 history steps times three coordinates. They are filled in the next step.

3. The ego history as tokens

The backbone reads tokens, so the history has to become tokens. DeltaTrajectoryTokenizer takes the step-to-step motion (each position minus the one before), maps each coordinate's range (by default −4 to 4 m for x and y) onto a fixed number of bins (1,000 by default), and rounds. The token ids are offsets into a block of trajectory tokens added to the vocabulary (<i0>, <i1>, …), and fuse_traj_tokens writes them over the placeholders.

Our version of the same idea:

def delta_tokenize(
    positions: torch.Tensor, bins: int = 1000, low: float = -4.0, high: float = 4.0
) -> torch.Tensor:
    """Turns a path into tokens: each step's (dx, dy) quantised into ``bins`` levels.

    The first step is measured from the origin. Steps outside [low, high] metres are
    clipped to the end bins.
    """
    steps = torch.diff(positions, dim=0, prepend=torch.zeros(1, positions.shape[1]))
    scaled = (steps - low) / (high - low) * (bins - 1)
    return scaled.round().long().clamp(0, bins - 1)


def delta_detokenize(
    tokens: torch.Tensor, bins: int = 1000, low: float = -4.0, high: float = 4.0
) -> torch.Tensor:
    """Inverts ``delta_tokenize``: bin centres back to steps, then a running sum."""
    steps = tokens.float() / (bins - 1) * (high - low) + low
    return torch.cumsum(steps, dim=0)

Encoding steps rather than positions keeps every token in a small range wherever the car is. The cost is that decoding is a running sum, so rounding errors add up along the path. For a 6.4 s curve at 10 m/s, 1,000 bins reconstruct the path to within 2.6 cm; 100 bins to within 24 cm. For an input that is plenty, which is why the history can afford it.

The same tokenizer family produced the discrete future trajectory the model is trained to write after its reasoning (the 128 tokens the paper describes). At inference, as the next section shows, those tokens are switched off.

4. Writing the reasoning

sample_trajectories_from_data_with_vlm_rollout in models/alpamayo_r1.py is the function this page follows. Its first half is ordinary language-model generation, with two details that matter.

• Trajectory tokens are masked. An ExpertLogitsProcessor sets the logits of every trajectory token to minus infinity before sampling, so the backbone can only write text. Its docstring gives the reason: the discrete trajectory tokens are not used by the expert, and masking them gives better Chain-of-Causation text.
• Generation stops one token late. A custom StopAfterEOS stops after the token that follows <|traj_future_start|>. The comment explains why: the key-value cache is only updated for a token when the next one is generated, and the expert needs the cache to include the end of the reasoning.

The test script samples with temperature=0.6, top_p=0.98 and up to 256 new tokens, and asks for one trajectory sample (num_traj_samples=1) so that it fits in GPU memory. More samples give more trajectories and reasoning traces for the same scene.

What is left after generation is the key-value cache of the prompt plus the reasoning. That cache is the whole interface between the two halves of the model.

5. The action expert reads the cache

The expert is built in the model's constructor with AutoModel.from_config from a copy of the backbone's text configuration (with optional overrides), and its token-embedding table is deleted. It is a transformer with the same layer shapes as the language model, which is what lets its layers attend to keys and values the language model computed. It never sees text tokens; its inputs are action tokens.

Each flow step calls an inner step_fn:

1. action_in_proj turns the noisy chunk into one token per waypoint (64 tokens). Each coordinate (acceleration, curvature) and the flow time get Fourier features (sines and cosines at logarithmically spaced frequencies), an MLP mixes them, and a layer norm finishes.
2. The expert runs on those 64 tokens with past_key_values set to the reasoning model's cache. Position ids continue from where the reasoning ended, and an attention mask blocks the padding that shorter samples in a batch have after their reasoning. Attention among the action tokens is not causal (expert_non_causal_attention defaults to true): every waypoint sees every other.
3. The cache is cropped back to its original length, so the next step starts clean.
4. action_out_proj maps each output token to a velocity for that waypoint's two controls.

The mask, in our version:

def expert_attention_mask(
    prompt_length: int, valid_until: int, actions: int
) -> torch.Tensor:
    """Which positions each action token may attend to, as an additive mask.

    The reasoning model's cache holds ``prompt_length`` positions: the prompt, the
    reasoning text and, for shorter samples in a batch, padding after the end of the
    text at ``valid_until``. The action tokens come after the cache. Each action token
    may see the cache up to ``valid_until`` and every action token (in both directions,
    since the expert is not causal); the padding is blocked.

    Returns:
        Shape (actions, prompt_length + actions): 0 where allowed, -inf where blocked.
    """
    mask = torch.zeros(actions, prompt_length + actions)
    mask[:, valid_until:prompt_length] = float("-inf")
    return mask

This design is why the paper's latency table can show 5 flow steps costing about 9 ms after 70 ms of reasoning: the expensive part, reading the images and writing the reasoning, happens once, and each flow step is one pass of the expert over 64 tokens against a fixed cache.

6. Flow matching, as in the course

diffusion/flow_matching.py uses the same convention as the Flow matching module: noise at $t = 0$, data at $t = 1$, the noisy chunk $x_t = t\,x + (1-t)\,\varepsilon$, and the network regresses the velocity $x - \varepsilon$ with a squared error. Sampling starts from Gaussian noise and takes Euler steps from 0 to 1.

Two settings are worth knowing:

• Number of steps. The class default is 10. The paper's latency table reports 5. The checkpoint's own configuration decides which runs; we have not read it, because it sits behind the gated access.
• Where training spends its effort. Training times are drawn from a Beta(1.5, 1) distribution and flipped, so about two thirds of samples have $t$ below 0.5: training spends more of its effort near the noisy end than a uniform draw would.

7. From controls to a trajectory

The chunk the expert produces is not positions. It is 64 pairs of acceleration and curvature, normalised by a stored mean and standard deviation. action_to_traj in action_space/unicycle_accel_curvature.py rolls them out from the current state: speed is the running sum of acceleration, heading turns by curvature times distance travelled, and positions integrate the velocity with the trapezoid rule. The starting speed comes from the history.

def unicycle_rollout(
    accel: torch.Tensor, curvature: torch.Tensor, v0: float, dt: float = DT
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Integrates (acceleration, curvature) controls into a path from the origin.

    Speed is the running sum of acceleration; heading turns by curvature times the
    distance travelled in each step; position uses the average of the start and end
    velocity of each step (the trapezoid rule). The car starts at the origin facing +x.

    Args:
        accel: Accelerations, shape (n,), in m/s^2.
        curvature: Curvatures, shape (n,), in 1/m.
        v0: Speed at the start, in m/s.
        dt: Time step in seconds.

    Returns:
        Positions (n, 2), headings (n,) and speeds (n,) after each step.
    """
    speed = torch.cat([torch.tensor([v0]), v0 + torch.cumsum(accel * dt, dim=0)])
    turn = curvature * speed[:-1] * dt + curvature * accel * dt**2 / 2
    heading = torch.cat([torch.zeros(1), torch.cumsum(turn, dim=0)])
    vx, vy = speed * torch.cos(heading), speed * torch.sin(heading)
    x = torch.cumsum((vx[:-1] + vx[1:]) * dt / 2, dim=0)
    y = torch.cumsum((vy[:-1] + vy[1:]) * dt / 2, dim=0)
    return torch.stack([x, y], dim=-1), heading[1:], speed[1:]

Our tests check the geometry: constant curvature 0.02 m⁻¹ at 10 m/s traces a circle of radius 50 m and turns 1.28 rad in 6.4 s. They also show something a downstream controller must handle: braking at 2 m/s² from 6 m/s reaches zero speed after 3 s and then keeps going, into negative speed. The released rollout, like ours, has no stop: speed is a running sum with no clamp.

The reverse direction, used to make training targets from logged paths, is a regularised least-squares fit (traj_to_action): find accelerations and curvatures whose rollout matches the log, with a penalty that keeps them smooth. Curvature is set to zero below 0.6 m/s, where heading is too noisy to define it. Bounds of ±9.8 m/s² and ±0.33 m⁻¹ are stored for checking actions.

8. What comes out

The function returns positions and rotations for each sample, shaped [batch, sets, samples, 64, …], and, with return_extra=True, the decoded reasoning text. The test script prints the reasoning and the minADE over the samples against the logged future: with one sample, that is simply its average displacement error.

What the code adds to the paper

Reading both, a few things stand out. Each is a fact about the released code's defaults, not a claim about what NVIDIA trained.

None of these is a contradiction; checkpoints carry their own configuration. They are reminders to check the code and the model card before quoting a number from the paper for the model you downloaded.

Reference implementation (PyTorch)

libs/tutorial/alpamayo_pieces.py; tests are in libs/tutorial/tests/test_alpamayo_pieces.py. They check that zero controls drive straight at constant speed, that constant curvature traces a circle of radius one over the curvature, that braking runs through zero speed, that delta tokens round-trip with an error that grows along the path and with fewer bins, that large steps are clipped, and that the expert's mask blocks exactly the padding. Run uv run --group tutorial python -m libs.tutorial.alpamayo_pieces.

Check yourself

1. Where does the reasoning end and the trajectory begin, in tokens, and what tells generation to stop?
2. Why does the expert need the same layer shapes as the backbone's language model?
3. Why are the discrete trajectory tokens masked during reasoning generation?
4. Why is the history encoded as step-to-step motion, and what does that cost when decoding?
5. A flow step costs one expert pass. Why is that cheap compared with the reasoning?
6. What does the released rollout do when a braking command would take the speed below zero?

Quiz

Answer each question; feedback appears as soon as you choose.

Sources

• NVlabs, Alpamayo repository, commit 11a0e01, Apache-2.0; files named in the text are under src/alpamayo_r1/.
• Alpamayo-R1 paper, arXiv 2511.00088.
• Lipman et al., Flow Matching for Generative Modeling, 2022, cited in the repository's flow-matching class.