Toolkit: generating trajectories with flow matching

Why this is in a driving course

At a junction there are several good trajectories. A network trained with squared error predicts their average, which can be a bad trajectory (the imitation and prediction modules measure this). Generative models fix it by sampling one good answer instead of averaging them. Flow matching is the generative method behind pi0's action expert and Alpamayo-R1's trajectory decoder. The full treatment is the Flow matching module.

The idea: move noise to data

Start with a point of pure Gaussian noise. Move it, step by step, until it is a sample from the data, say a 64-step trajectory. If you knew the right velocity at every point and time, you could integrate it like an ODE and arrive at a realistic sample. Flow matching trains a network to predict that velocity.

Time runs from $t = 0$ (noise) to $t = 1$ (data). Training uses the simplest path there is: a straight line from a noise sample $x_0$ to a data sample $x_1$,

whose velocity is constant, $x_1 - x_0$. The network $v_\theta(x_t, t)$ is trained to predict it, with squared error:

def flow_matching_loss(velocity: VelocityField, x1: torch.Tensor) -> torch.Tensor:
    """Regresses the velocity of the straight line from noise to data.

    Args:
        velocity: A network taking (x_t, t) and returning a velocity of x_t's shape.
        x1: A batch of data points, shape (batch, dim).

    Returns:
        The mean squared error between the predicted and target velocity.
    """
    x0 = torch.randn_like(x1)
    t = torch.rand(len(x1), 1)
    x_t = (1 - t) * x0 + t * x1
    target = x1 - x0
    return ((velocity(x_t, t) - target) ** 2).mean()

That is the whole training loop: pick data, pick noise, pick a time, regress the velocity. No adversary, no likelihood, no noise schedule to tune. The released Alpamayo-R1 code uses this same convention.

Why it does not average

Each training pair pulls the network towards one straight line. Where lines from different data points pass through the same place, the network learns their average direction at that place and time. Early on ($t$ near 0) that average points at the middle of the data; later, as the point gets closer to one cluster, the average points at that cluster. Integrating the field therefore splits noise between the modes instead of landing between them.

Sampling with a few Euler steps

To generate, start from noise and take Euler steps along the predicted velocity:

@torch.no_grad()
def euler_sample(
    velocity: VelocityField, noise: torch.Tensor, steps: int
) -> torch.Tensor:
    """Integrates dx/dt = v(x, t) from t = 0 (noise) to t = 1 with Euler steps."""
    x = noise
    for i in range(steps):
        t = torch.full((len(x), 1), i / steps)
        x = x + velocity(x, t) / steps
    return x

The number of steps is a speed-quality dial. One step is the average again; a handful is usually enough because the learned paths are nearly straight. Alpamayo-R1's paper reports 5 steps costing under 9 ms on its GPU, against about 70 ms for its reasoning.

Conditioning: generating this car's trajectory

A driving model does not want any realistic trajectory; it wants one for this scene. The velocity network therefore also takes the scene as input. In an action expert, the noisy chunk becomes tokens that attend to the scene features (or, in Alpamayo-R1, to the reasoning model's key-value cache), so the predicted velocity depends on what the cameras saw and what the reasoning said. Training is unchanged: the data sample is the logged trajectory for that scene.

Where this shows up in the course

• Imitation: why averaging demonstrations is unsafe.
• VLA models: the action expert against discrete tokens, measured on an ambiguous instruction.
• Alpamayo-R1: its trajectory decoder, the number of flow steps and its latency.
• Deployment: fewer steps as a way to cut latency.

Check yourself

1. What does the network in flow matching predict, and what is its training target?
2. Why does a squared-error regression of the trajectory itself average the modes, while flow matching does not?
3. What happens if you sample with a single Euler step, and why?
4. How does a flow-matching action expert know which scene it is generating for?
5. Why is the number of steps a latency decision?

Quiz

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

Go deeper

The Flow matching module trains a flow on a toy, compares solvers, derives the exact velocity field and straightens the paths with reflow; Sampling and guidance covers diffusion policies and guidance.