Flow matching: how π₀ generates actions

The anatomy lesson left one box deliberately unopened: the action expert. You now know where it sits — a smaller transformer bolted onto the VLM backbone, reading its features through attention — but not how it turns those features into a $50 \times 7$ matrix of joint targets for your WidowX AI. That matrix is continuous and multimodal: when two identical mugs sit on the table, the demonstrations contain both grasps, and a policy that regresses their mean reaches between them. Phase 03 taught you the standard fix — diffusion — and its price: on the order of a hundred denoising passes per action chunk. Tolerable for a 90M-parameter UNet; π₀'s denoiser conditions on a 3-billion-parameter vision-language backbone. Multiply 3B by 100 passes and the arm stands still while the GPU melts. Flow matching is Physical Intelligence's answer: a training objective that makes few-step generation a first-class design goal rather than an inference-time hack.

Conceptual Foundation: Probability Paths and Velocity Fields

Before deriving the loss, we must rigorously define the geometric object being learned: the probability path. In diffusion models, the timestep $t$ typically runs from $0$ (data) to $T$ (noise), describing a forward corruption process. Flow matching inverts this intuition. We define a continuous flow time $\tau \in [0, 1]$ where $\tau = 0$ corresponds to the source distribution (standard Gaussian noise $\epsilon$) and $\tau = 1$ corresponds to the target data distribution (action chunk $A$). The probability path $p_\tau(x)$ is the marginal distribution of the interpolated point $x$ at time $\tau$. Unlike diffusion, which relies on a stochastic differential equation (SDE) with a noise schedule, flow matching uses a deterministic ordinary differential equation (ODE) driven by a velocity field $v(x, \tau)$. The path is defined by the straight-line interpolation $x_\tau = (1-\tau)\epsilon + \tau A$. This choice is not arbitrary; it minimizes the expected path length, making the ODE trajectory as straight as possible in expectation, which is the key to few-step integration.

The central theoretical result is that the marginal velocity field $v^*(x, \tau)$, which transports the full noise distribution to the full data distribution, is the conditional expectation of the conditional velocities. For a fixed point $x$ at time $\tau$, multiple training pairs $(A, \epsilon)$ may pass through $x$. The network learns the average of their velocities: $v^*(x, \tau) = \mathbb{E}[A - \epsilon \mid x_\tau = x]$. This is not an approximation; it is the exact solution to the optimal transport problem under the linear interpolation assumption. The 'straight-line' assumption ensures that the conditional velocity $A - \epsilon$ is constant along the path, making the expectation well-defined and the ODE integrable with simple Euler steps.

Why sample $\tau$ from a beta distribution skewed toward the noisy end ($\tau \approx 0$)? At $\tau = 0$, the point $x$ is pure noise, and the posterior over $A$ is broad. The velocity field must be accurate here to correctly 'steer' the noise toward the correct mode. If $\tau$ were sampled uniformly, the network would spend too much capacity on the easy, high-signal end ($\tau \approx 1$) where the data is already clear. By oversampling the noisy end, we ensure the field is well-conditioned where the transport is most ambiguous, preserving multimodality. This is a form of curriculum learning: the network learns to resolve ambiguity early in the trajectory.

Worked Example: Latency Arithmetic via the Roofline Model

To justify the claim that flow matching is 'cheap,' we must move from qualitative assertions to quantitative bounds. The roofline model provides a hardware-agnostic upper bound on performance based on two constraints: compute throughput (FLOPs/s) and memory bandwidth (Bytes/s). For a forward pass of a neural network, the operation is often memory-bound if the arithmetic intensity (FLOPs per byte) is low. In the action expert, the batch size is 1, and the sequence length is small (51 tokens). The dominant cost is reading the model weights from memory, not performing matrix multiplications. We will calculate the minimum latency for one forward pass of the 300M-parameter expert on an RTX 4090.

First, calculate the total weight size in bytes. $P \times \text{bytes/param} = 300 \times 10^6 \times 2 = 600 \times 10^6$ bytes $= 0.6$ GB. The minimum time to read these weights is $t_{\text{min}} = \frac{\text{Weight Size}}{\text{Bandwidth}} = \frac{0.6 \text{ GB}}{1008 \text{ GB/s}} \approx 0.000595$ s $= 0.595$ ms. This is the theoretical lower bound. In practice, kernel launch overhead, attention computation, and non-ideal memory access patterns add latency. Empirically, a 300M model at batch 1 on an RTX 4090 takes approximately 1.5–2.0 ms per forward pass. We will use 1.6 ms as a conservative estimate for the expert step.

Recall the Phase 03 arithmetic. A diffusion policy's UNet at ~90M parameters costs 1–2 ms per pass on your RTX workstation; DDPM sampling at 50–100 steps lands the chunk in 50–200 ms. Annoying but workable — it is why diffusion policies visibly pause-and-go, and why they deploy with action chunking to amortize the pause over many control ticks.

Now scale the denoiser to a VLA. A 3B-parameter forward pass at batch 1 is memory-bandwidth-bound: 6 GB of bf16 weights over the RTX 4090's ~1 TB/s gives a hard floor of ~6 ms, and realistic kernels with a ~500-token multimodal prefix land at 20–60 ms per pass. Run the DDPM recipe naively — full network, every step — and 50 steps costs 1–3 s of compute to produce one second of motion (a 50-action chunk at 50 Hz). The robot cannot even keep up open-loop: it finishes executing a chunk before the next one exists. Even a 10-step DDIM sampler leaves you at 200–600 ms with degraded sample quality, because DDPM's training objective never promised anything about coarse step counts.

There are exactly two levers: make the passes fewer and make each pass cheaper. Flow matching is the first lever — an objective whose sampler is a well-behaved ODE you can integrate in ~10 Euler steps. The action-expert architecture is the second — only a ~300M model runs inside the loop. This lesson takes the levers in that order.

Flow matching: learn the velocity, not the noise

Setup. Given an observation $o$ (camera images, language instruction, proprioception), we want to sample an action chunk $A \sim p(A \mid o)$. Flow matching's move: pick a trivially samplable source distribution — standard Gaussian noise — and learn to transport noise samples to data samples along paths so simple that a handful of integration steps traverses them. The simplest possible path is a straight line. Pair each training chunk $A$ with a fresh noise draw $\epsilon$ and define the interpolation:

Train a network $v_\theta$ to predict that velocity from the interpolated point, the flow time, and the observation. This is the conditional flow matching objective — the entire training recipe fits in one line:

• $A \in \mathbb{R}^{H \times d}$ — one action chunk from the dataset: $H = 50$ future actions; in OpenPI, $d$ is padded to a fixed 32 so one expert serves many embodiments (your WidowX AI fills 7 slots: 6 joints + gripper), and actions are quantile-normalized to roughly unit scale.
• $\epsilon$ — a fresh standard-normal tensor of the same shape, drawn per training example, per step. It is the source sample, not corruption noise to be predicted.
• $\tau \in [0,1]$ — flow time, the analogue of the diffusion timestep. π₀ does not sample it uniformly: it uses a beta-shaped distribution skewed toward the noisy end, truncated just below 1 (at 0.999), spending more gradient where the transport problem is hardest.
• $A^{\tau}$ — the point on the straight segment between $\epsilon$ and $A$; the network's main input.
• $v_\theta(\cdot)$ — the action expert. Its regression target $A - \epsilon$ is the segment's constant velocity.
• $o$ — the conditioning: in π₀, the backbone's KV cache over image and language tokens, plus the proprioceptive state.

One question decides whether this objective is legitimate: many different $(A, \epsilon)$ pairs produce paths passing through the same point $x$ at the same time $\tau$, each demanding a different velocity. What does least squares do with conflicting targets? Exactly what it always does — converge to their conditional mean. Fix a point and minimize pointwise:

That conditional expectation is precisely the marginal velocity field whose ODE transports the full noise distribution to the full data distribution — the central theorem of Flow Matching for Generative Modeling (opens in a new tab): the conditional and marginal objectives have identical gradients, so you train on trivial straight-line targets and still inherit a valid generative model. Note what did not get lost: multimodality. The field is deterministic, but randomness enters through the starting noise $\epsilon$ — different starts flow to different modes. Both mug grasps survive; nothing was averaged away at the distribution level, only at the field level.

Sampling is ODE integration: ten Euler steps

Inference is now numerical integration of a learned ODE. Start from pure noise, follow the field, take $K$ equal Euler steps:

How low can $K$ go? The derivation above answers this exactly, and the answer is one of the most clarifying results in this whole phase. At $\tau = 0$ the input $A^{0} = \epsilon$ carries no information about which chunk it was paired with — noise is drawn independently of data. So the conditional expectation collapses to $\mathbb{E}[A - \epsilon \mid A^0 = x, o] = \mathbb{E}[A \mid o] - x$, and a single full-length Euler step gives:

So $K$ interpolates between two policies you already know: at $K = 1$, MSE behavioral cloning (the gripper reaches between the two mugs); as $K \to \infty$, exact samples from the demonstrated distribution. π₀'s empirical claim is that real manipulation data needs only ~10 steps to get within task-relevant accuracy. Here is the entire algorithm, both halves, in numpy:

import numpy as np

def fm_training_step(v_theta, actions, obs_features, rng):
    """actions: (B, H, D) normalized chunks; obs_features: backbone KV."""
    B, H, D = actions.shape
    eps = rng.standard_normal((B, H, D))
    tau = rng.uniform(0.0, 0.999, size=(B, 1, 1))  # uniform here; pi0 uses a beta skewed toward the noisy end
    x_tau = (1.0 - tau) * eps + tau * actions      # point on the straight path
    target = actions - eps                         # constant conditional velocity
    pred = v_theta(x_tau, tau, obs_features)
    return np.mean((pred - target) ** 2)           # backprop this

def sample_chunk(v_theta, obs_features, H, D, rng, K=10):
    """Euler-integrate the learned field from noise (tau=0) toward data (tau=1)."""
    x = rng.standard_normal((H, D))
    dt = 1.0 / K
    for k in range(K):
        x = x + dt * v_theta(x, k * dt, obs_features)
    return x   # one chunk: H=50 rows of joint targets for the arm

Now the latency arithmetic, Fermi-style, for an RTX 4090-class card — you will validate every term in the serving lesson. Prefill (once per chunk): two RealSense views at $224^2$ through SigLIP give $2 \times 256 = 512$ vision tokens plus ~20 language tokens: call it 530 tokens through the 3B backbone, $2 \times 3\mathrm{B} \times 530 \approx 3.2$ TFLOPs — 40–70 ms at realistic utilization including the vision encoder. Per integration step: only the ~300M expert runs, over 51 tokens (50 noisy actions + 1 state), attending to the cached backbone KV. Weight traffic is 0.6 GB of bf16 — a ~0.6 ms bandwidth floor — landing at 2–4 ms per step with attention and kernel overheads; ten steps, 20–40 ms. Total: roughly 80–120 ms for one second of motion, better than 8× real time, with headroom to replan several times per second. Compare:

The π₀ action expert: small where it counts

Put the architecture numbers from the anatomy lesson back on the table, now with the flow-matching lens. The π₀ backbone (opens in a new tab) is PaliGemma — a SigLIP vision encoder feeding a Gemma language model, ~3B parameters — and it processes images and language once per observation. The action expert is a separate transformer of roughly 300M parameters with a narrower hidden width, trained from scratch. Its inputs are the noisy chunk $A^{\tau}$ (50 tokens), the robot's proprioceptive state, and the flow time $\tau$; through blockwise attention it reads the backbone's cached keys and values but maintains its own weights. During sampling, the ten Euler steps loop over the expert alone — the backbone never runs twice for the same observation.

Why does 300M suffice next to 3B? Because the hard problem was already solved upstream. The backbone does open-vocabulary perception and instruction grounding — the part that needs web-scale knowledge. What remains is a conditional regression in a low-dimensional, physically smooth space: given “the mug's handle is here and we are grasping it,” predict a velocity field over $50 \times 32$ normalized action values. Joint trajectories at 50 Hz are bounded and strongly autocorrelated — orders of magnitude less entropy than images or text. Spending 3B parameters there would buy little and cost 10× on every one of the ten steps in the loop. It is the asymmetry your serving stack exploits with draft models: match capacity to the conditional entropy of the subproblem.

Training recipe: mixtures and knowledge insulation

π₀'s training follows the LLM playbook you already run: pretrain broad, post-train narrow. Pretraining mixes Physical Intelligence's own dexterous teleoperation corpus with open cross-embodiment data (OXE-style) — on the order of 10,000 hours across many robot types, with action vectors padded to the shared 32-wide layout so one expert spans embodiments. This stage deliberately includes messy, imperfect data: like web-scale pretraining, its job is coverage — recovery behaviors, embodiment variety, task breadth. Post-training then fine-tunes on a curated, consistent demonstration set for the target task — often just 1–20 hours — to sharpen execution. The division of labor is exactly pretraining versus SFT: the base mixture determines what the policy can recover from; the post-training set determines what fluent execution looks like. When you fine-tune openpi (opens in a new tab) on your own WidowX AI demos in two lessons, you are doing the second stage only.

There is a published subtlety here that will shape your fine-tuning configuration. Train the flow-matching expert jointly with the backbone from day one and the expert's gradients flow back through the shared attention into the VLM — measurably eroding its web-scale semantics, visible as degraded language following and weaker generalization. The π₀.₅ line of work, formalized in the paper “Knowledge Insulating Vision-Language-Action Models,” names the fix knowledge insulation: stop the gradient at the expert-to-backbone boundary, so the expert learns continuous control from backbone features without writing into them. The backbone still adapts to robotics — but through a representation-compatible channel: predicting discretized action tokens (the FAST tokenization of the next lesson) alongside web-data co-training, the same next-token objective it was born with. Reported effects: substantially faster training, preserved instruction following, better generalization. The mental model: the backbone learns robotics as language; the expert learns control as flow; a gradient stop keeps the second from corrupting the first.

Week 13: the three knobs you will ablate

Everything above reduces to three runtime knobs, and each one maps to a behavior you can measure on the arm with the instrumentation you built in Phase 01. Design the grid now, before the fine-tuning lesson, so data collection serves it:

• Integration steps $K \in \{1, 2, 5, 10, 20\}$. Offline: distributional distance of sampled chunks against a K = 50 reference on held-out observations. On the robot: success rate on a deliberately bimodal scene (two identical objects) and gripper-path jitter. Prediction from the math: K = 1 reaches between the objects every time; quality plateaus by K ≈ 5–10; latency cost is only ~3 ms per added step.
• Action horizon (execute 25 vs 50 of the 50-step chunk). The staleness of the last executed action is 0.5 s vs 1.0 s plus generation time. Measure tracking error against a target you displace mid-episode, and total episode time (shorter horizons replan more, paying prefill more often — 40–70 ms each time).
• Replanning rate (execute $m$ actions, then regenerate). m = 10 gives 5 Hz replanning and fast reaction to a mid-episode push; m = 50 is fully open-loop. Measure joint-velocity discontinuity at chunk boundaries in the encoder stream — naive replanning splices trajectories that disagree, and the jerk spikes are visible in your latency-probe logs. This axis is the capstone's home turf.

Where this goes next: the previous lesson, Anatomy of a vision-language-action model, gave you the map of π₀; this one gave you its engine — a velocity field you can derive, sample, and budget to the millisecond. But flow matching is only one of Physical Intelligence's two answers to action generation. The next lesson, π₀-FAST: actions as tokens done right, takes the opposite bet: discretize actions and let the backbone decode them autoregressively — the same discrete channel knowledge insulation exploits during training. That sets up a question your capstone lives on: tokens decode slower but play natively with the LLM serving stack you know; continuous flow is faster but needs its own expert. Measure both before you choose.