Section 22.5: Flow matching for actions

Flow Matching For Continuous Action Paths

With the data interface in hand, the next step is the objective that turns those demonstration chunks into a learnable velocity field. Flow matching trains a vector field that moves samples from a simple base distribution toward the data distribution. For action chunks, imagine drawing an initial noisy chunk \(A_0\) and a demonstration chunk \(A_1\), then training a network \(v_\theta(A_t,t,o)\) to predict the velocity along a path between them. A simple linear path is:

$$A_t = (1-t)A_0 + tA_1, \quad u_t = A_1 - A_0, \quad \mathcal{L}_{FM}=\mathbb{E}\left[\|v_\theta(A_t,t,o)-u_t\|_2^2\right].$$

The diagram below traces this construction visually: it shows the straight path from a noise sample to a demonstration action and the constant velocity the network learns to predict at any point along it.

Why linearity matters for physical robots. A straight interpolation path keeps the target velocity \(u_t = A_1 - A_0\) constant across all timesteps \(t\). In practice this is more than a training convenience: it also tends to be a safety-relevant property, though the section does not claim it as a formal guarantee. Curved paths force the network to predict a different velocity at every \(t\), and any misprediction accumulates into joint-velocity spikes during integration. On a real arm, a sudden velocity spike can saturate actuators, trigger safety stops, or damage end-effector hardware. A constant target removes that accumulation. The training signal stays the same everywhere on the path, so the learned field is typically smoother and integration errors tend to stay bounded over the short chunk horizon.

How the regression works. At each training step the network sees a mid-path sample \(A_t\) and must output the single constant vector pointing from \(A_0\) to \(A_1\). Because the target does not change with \(t\), the network can treat time as a mild conditioning signal rather than the primary input. During inference, a simple Euler integrator advances \(A_t\) by \(\Delta t \cdot v_\theta(A_t, t, o)\) at each step. With a nearly constant field, even two to four Euler steps stay close to the true path, which is why inference is fast: the field is easy to follow, not just easy to compute.

The algorithm above is the specification; the network itself is small enough to write out in full. Code Fragment 22.5.3 gives a minimal velocity-field network that follows the algorithm's steps directly: an MLP conditioned on the noisy action, the timestep, and the observation embedding, trained with the flow-matching loss and sampled with fixed-step Euler integration.

import torch
import torch.nn as nn

class VelocityField(nn.Module):
    """Minimal velocity-field network for flow matching over action chunks.
    Predicts v_theta(A_t, t, o): the same object trained in the algorithm above."""
    def __init__(self, action_dim, obs_dim, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(action_dim + 1 + obs_dim, hidden), nn.Mish(),
            nn.Linear(hidden, hidden), nn.Mish(),
            nn.Linear(hidden, action_dim),
        )

    def forward(self, a_t, t, obs):
        t = t.view(-1, 1).expand(a_t.shape[0], 1)
        return self.net(torch.cat([a_t, t, obs], dim=-1))

def flow_matching_loss(model, a1, obs):
    """One training step: sample noise and t, interpolate, regress the constant velocity."""
    a0 = torch.randn_like(a1)
    t = torch.rand(a1.shape[0], device=a1.device)
    a_t = (1 - t.view(-1, 1)) * a0 + t.view(-1, 1) * a1
    u_t = a1 - a0
    pred = model(a_t, t, obs)
    return ((pred - u_t) ** 2).mean()

@torch.no_grad()
def sample_action(model, obs, action_dim, n_steps=4):
    """Fixed-step Euler inference: the same integration loop as algorithm step 8."""
    a_t = torch.randn(obs.shape[0], action_dim, device=obs.device)
    dt = 1.0 / n_steps
    for step in range(n_steps):
        t = torch.full((obs.shape[0],), step * dt, device=obs.device)
        a_t = a_t + dt * model(a_t, t, obs)
    return a_t

The practical appeal is straight-line sampling speed. A learned vector field generates actions with far fewer integration steps than DDPM, which matters for control loops with tight latency budgets. Chi et al. (2023) report that a DDPM-based Diffusion Policy requires 100 denoising steps at inference. On a typical desktop GPU, those steps take roughly 0.1 s per chunk, though exact timing depends on the network size and GPU model, limiting effective control to about 10 Hz. Put concretely, a peg-insertion task that requires 50 Hz closed-loop correction gets zero corrective updates during those 100 ms. The arm flies blind for five full control cycles while the policy thinks. A flow-matching variant using an ODE solver with 10 function evaluations can reach comparable quality in roughly 10 ms on similar hardware. That speed opens the door to 50 Hz or higher control rates, which contact-rich tasks such as peg insertion and cloth folding require.

A policy that finishes sampling after the robot needed to act is not a slow policy; it is a spectator.

Setting aside the multimodal failure case for a moment, it helps to see the linear path's velocity target computed on a single concrete sample, since that target is exactly what the network regresses toward during training. Code Fragment 22.5.2 computes the target velocity for a two-dimensional action chunk under the linear flow path.

# Compute the flow-matching target for a simple action path.
# The vector field should point from noise toward the demonstrated action.
import numpy as np

noise_action = np.array([-0.5, 0.2])
demo_action = np.array([0.3, 0.6])
t = 0.25
intermediate = (1 - t) * noise_action + t * demo_action
velocity_target = demo_action - noise_action
print("intermediate:", intermediate.round(2).tolist())
print("velocity target:", velocity_target.round(2).tolist())