§13.3 Flow matching as a control objective

A velocity field, not a denoiser

The mental picture is a flow of particles. At time $\tau = 0$ every particle is a fresh draw of noise. At time $\tau = 1$ every particle has landed on a real action chunk. In between, each particle follows a smooth path, and the whole population morphs from the Gaussian blob into the shape of the demonstration data. If you knew the velocity of a particle at every position and every time, generating a sample would be trivial: draw noise, then integrate that velocity forward from $\tau=0$ to $\tau=1$. The learned object in flow matching is exactly that velocity field, a function $v_\theta(a_\tau, \tau, \text{scene})$ that reads the current noisy chunk, the current time, and the backbone’s scene representation, and returns which way to push.

This is the point of contact with §13.2. The action expert is $v_\theta$. Every one of the ten integration steps at inference is one evaluation of the velocity field, one forward pass through the expert, conditioned through the shared attention on the scene the backbone already encoded. So training the expert means learning a good velocity field, and the elegance of flow matching is that you can learn it with an ordinary regression loss, no simulation of the flow during training, no adversarial game, no variational bound.

The loss is a regression, and that should worry you until it doesn’t

Here is the move that makes flow matching cheap to train, and it is worth slowing down for because it looks like it should reintroduce the averaging trap that killed naive regression in §13.1.

Pick a straight path between noise and data. For a clean action chunk $a_1$ drawn from the demonstrations and a noise chunk $a_0$ drawn from the Gaussian, define the interpolated point at time $\tau$ as

$a_\tau = (1-\tau)\, a_0 + \tau\, a_1.$

Differentiate with respect to $\tau$ and the velocity along this particular path is just $a_1 - a_0$, a constant. That is the target. The training step is: sample a demonstration chunk, sample a noise chunk, sample a random time $\tau \in [0,1]$, form $a_\tau$, and regress the expert’s output onto the constant velocity $a_1 - a_0$:

$\mathcal{L}(\theta) = \mathbb{E}_{a_1,\, a_0,\, \tau}\Big[\, \big\| v_\theta(a_\tau, \tau, \text{scene}) - (a_1 - a_0) \big\|^2 \,\Big].$

A squared-error loss. The same loss that averages left and right turns into a crash. Why does it not collapse here?

Because the target is conditioned on both the noise draw and the time, and that conditioning is what breaks the tie. When the scene admits two good action chunks, a left fold and a right fold, the demonstrations contain both. During training, a given noisy point $a_\tau$ gets pulled toward the left chunk in one sampled pair and toward the right chunk in another. The expert cannot satisfy both with a single vector, so at that shared point it learns the average velocity, which does point roughly straight ahead. That sounds like the trap. It isn’t, and the reason is the whole trick: the region of noise space that flows toward the left mode and the region that flows toward the right mode are mostly different regions, separated by the noise draw $a_0$. A particle that starts in the left basin has an unambiguous target throughout almost all of its path and only sees the averaged, ambiguous velocity in the thin seam between basins. Integrate from a specific noise sample and you commit to one basin early and ride a clean path into one mode. Draw a different noise sample and you land in the other. The multimodality lives in the noise you draw at inference, not in a distribution the network has to emit in one shot, and that is how a plain regression loss ends up sampling honestly from a multimodal action distribution.

Why π0 uses this instead of diffusion

Diffusion policies (Chapter 10) attack the same multimodality requirement and get there. π0’s designers chose flow matching over diffusion, and the reason is the one §13.1 kept circling: steps cost latency, and control has no latency to spare.

A diffusion sampler learns to reverse a noising process, and the standard training recipe produces a curved, stochastic reverse trajectory that wants many small denoising steps to follow accurately. Cut the step count too aggressively and sample quality falls off. Flow matching with the straight-line path above learns paths that are much closer to straight, because the target velocity is literally constant along each training path. Straighter paths integrate accurately with a coarse solver, so you can use a handful of Euler steps instead of dozens of denoising steps and still land on a clean chunk. This is the rectified-flow idea from Liu et al. (arXiv 2209.03003), and it is the property §13.2 was cashing in when it said the expert iterates “ten or so” times. Ten forward passes through a 300M expert is a latency budget a 50 Hz two-arm robot can afford; fifty passes through a diffusion head might not be.

One honest qualification. The straight-path construction does not by itself guarantee a straight flow, since the learned field averages over many crossing paths and the actual sampling trajectory can bend where basins meet. Rectified flow can be iterated to straighten the field further, and in practice π0 uses a small fixed step count that works well enough without that extra machinery. The engineering claim is not “provably one step,” it is “few enough steps to close the loop,” and that is what matters on hardware.

One training step, in code

The objective is small enough to write out. This is the core of the π0-style training step for the action expert, stripped to the shape of the computation, with the backbone’s scene encoding passed in as a conditioning tensor.

def flow_matching_loss(expert, scene_repr, action_chunk):
    # action_chunk: (B, H, action_dim), a clean expert demonstration
    B = action_chunk.shape[0]

    a1 = action_chunk                       # target: clean chunk
    a0 = torch.randn_like(a1)               # source: Gaussian noise
    tau = torch.rand(B, 1, 1)               # random time in [0, 1]

    a_tau = (1 - tau) * a0 + tau * a1       # point on the straight path
    target_v = a1 - a0                      # constant velocity along it

    pred_v = expert(a_tau, tau, scene_repr) # the action expert IS v_theta
    return ((pred_v - target_v) ** 2).mean()

Notice what is absent. No forward simulation of the flow during training, no reverse process to unroll, no discretization of the action into bins. The whole objective is one interpolation and one squared error, and the network learns a field it never had to roll out to be scored on. That is why flow matching trains stably at the scale π0 needs.

Inference inverts the picture with an ODE solver, and the plainest one, forward Euler, is enough to see the shape:

def sample_chunk(expert, scene_repr, shape, n_steps=10):
    a = torch.randn(shape)                  # start from noise
    dt = 1.0 / n_steps
    for k in range(n_steps):
        tau = torch.full((shape[0], 1, 1), k * dt)
        a = a + expert(a, tau, scene_repr) * dt   # step along the field
    return a                                # a clean action chunk

Ten additions, each gated by one expert evaluation, and the scene representation is computed once outside the loop. Hold this next to OpenVLA emitting dozens of action tokens by serial decode through a 7B backbone and the latency asymmetry from §13.2 stops being an assertion and becomes arithmetic.

The toy version you can hold in your head

The chapter’s hands-on exercise builds this on a 2-D toy dataset, and it is worth previewing because two dimensions is where the whole objective becomes visible. Take target points arranged in two clusters, a stand-in for two demonstrated action modes. Train a tiny $v_\theta$ with exactly the loss above, no scene conditioning at all. Then plot the velocity field over the plane. You will see arrows that sweep the central noise cloud outward and split it, funneling the left half of the noise toward the left cluster and the right half toward the right, with a visible seam down the middle where the field goes slack because that is the ambiguous region no single velocity can resolve. Drop a hundred noise points and integrate them and they fan into the two clusters in the right proportion. That picture is the entire section: a regression loss that produces a splitting flow, and multimodality that comes out of where each particle started rather than out of anything the network emits at once. Scale the toy up, condition the field on a PaliGemma scene encoding, make the points 50-step action chunks for a two-arm robot, and you have π0’s action expert.

Flow matching, then, is not a detail of π0’s plumbing; it is the reason the continuous head clears all three bars from §13.1 at once. It emits a full chunk from one noise draw, it keeps genuine multimodality by routing it through the noise, and it samples in few enough steps to run at control rate. What none of this has shown yet is what that combination actually buys on a real robot against the token-head baselines, which is where §13.4 takes the argument next.