π0 — Visual Experiments¶

Goal: Understand flow matching — how π0 generates smooth, multimodal continuous action chunks by integrating a learned velocity field.

What we'll explore¶

#	Experiment	Question answered
1	Flow matching	How does a learned velocity field carry noise to a target action distribution?
2	Multimodal actions	Does the generative flow keep all valid action modes a regressor would average?
3	Integration steps	How many Euler steps does sampling need, and why are straight paths cheap?
4	Action chunking	How does predicting a chunk of actions reduce per-step jitter?

---

Linked theory: research/07-pi0.md

In [1]:

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import warnings
warnings.filterwarnings('ignore')

torch.manual_seed(0)
np.random.seed(42)
plt.rcParams.update({
    'figure.dpi': 110,
    'font.size': 11,
    'axes.titlesize': 12,
    'axes.labelsize': 11,
    'figure.facecolor': 'white',
    'axes.facecolor': '#f8f9fa',
    'axes.grid': True,
    'grid.alpha': 0.35,
})


def ring_targets(n, modes=6, radius=2.5, spread=0.18):
    ang = 2 * np.pi * np.random.randint(0, modes, n) / modes
    c = np.stack([radius * np.cos(ang), radius * np.sin(ang)], 1)
    return c + spread * np.random.randn(n, 2)


print('Setup complete.')

import numpy as np import matplotlib.pyplot as plt import torch import torch.nn as nn import warnings warnings.filterwarnings('ignore') torch.manual_seed(0) np.random.seed(42) plt.rcParams.update({ 'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11, 'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa', 'axes.grid': True, 'grid.alpha': 0.35, }) def ring_targets(n, modes=6, radius=2.5, spread=0.18): ang = 2 * np.pi * np.random.randint(0, modes, n) / modes c = np.stack([radius * np.cos(ang), radius * np.sin(ang)], 1) return c + spread * np.random.randn(n, 2) print('Setup complete.')

Setup complete.

---

Experiment 1: Flow matching¶

Question: π0's action expert learns a velocity field that transports Gaussian noise to the action distribution. How is it trained, and does integrating it actually reach the target?

We train a small velocity field $v_\theta(x,\tau)$ by regressing onto the straight-line target $x_1-x_0$, then sample by Euler-integrating from noise. The target action distribution is a ring of 6 modes (six valid actions).

In [2]:

torch.manual_seed(1)


class VelocityField(nn.Module):
    def __init__(self, h=128):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(3, h), nn.SiLU(), nn.Linear(h, h), nn.SiLU(), nn.Linear(h, 2))

    def forward(self, x, tau):
        return self.net(torch.cat([x, tau], dim=1))


v = VelocityField()
opt = torch.optim.Adam(v.parameters(), lr=2e-3)
data = torch.tensor(ring_targets(6000), dtype=torch.float32)
losses = []
for it in range(2500):
    idx = torch.randint(0, len(data), (512,))
    x1 = data[idx]
    x0 = torch.randn_like(x1)
    tau = torch.rand(512, 1)
    x_tau = (1 - tau) * x0 + tau * x1
    target = x1 - x0                       # straight-line conditional velocity
    loss = ((v(x_tau, tau) - target) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()
    losses.append(loss.item())


@torch.no_grad()
def sample(n, steps, return_traj=False):
    x = torch.randn(n, 2)
    traj = [x.numpy().copy()]
    for s in range(steps):
        tau = torch.full((n, 1), s / steps)
        x = x + v(x, tau) / steps          # Euler step of dx = v dtau
        traj.append(x.numpy().copy())
    return (x.numpy(), traj) if return_traj else x.numpy()


gen, traj = sample(1500, 30, return_traj=True)
print(f'final flow-matching loss: {np.mean(losses[-100:]):.4f}')

torch.manual_seed(1) class VelocityField(nn.Module): def __init__(self, h=128): super().__init__() self.net = nn.Sequential(nn.Linear(3, h), nn.SiLU(), nn.Linear(h, h), nn.SiLU(), nn.Linear(h, 2)) def forward(self, x, tau): return self.net(torch.cat([x, tau], dim=1)) v = VelocityField() opt = torch.optim.Adam(v.parameters(), lr=2e-3) data = torch.tensor(ring_targets(6000), dtype=torch.float32) losses = [] for it in range(2500): idx = torch.randint(0, len(data), (512,)) x1 = data[idx] x0 = torch.randn_like(x1) tau = torch.rand(512, 1) x_tau = (1 - tau) * x0 + tau * x1 target = x1 - x0 # straight-line conditional velocity loss = ((v(x_tau, tau) - target) ** 2).mean() opt.zero_grad(); loss.backward(); opt.step() losses.append(loss.item()) @torch.no_grad() def sample(n, steps, return_traj=False): x = torch.randn(n, 2) traj = [x.numpy().copy()] for s in range(steps): tau = torch.full((n, 1), s / steps) x = x + v(x, tau) / steps # Euler step of dx = v dtau traj.append(x.numpy().copy()) return (x.numpy(), traj) if return_traj else x.numpy() gen, traj = sample(1500, 30, return_traj=True) print(f'final flow-matching loss: {np.mean(losses[-100:]):.4f}')

final flow-matching loss: 2.3051

In [3]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5.5))
fig.suptitle('Flow matching carries noise to the action distribution', fontsize=13, fontweight='bold')

ax = axes[0]
ax.scatter(data[:, 0], data[:, 1], s=6, alpha=0.2, color='lightgray', label='target actions')
ax.scatter(gen[:, 0], gen[:, 1], s=8, alpha=0.5, color='seagreen', label='generated (flow)')
ax.set_xlabel('action dim 0'); ax.set_ylabel('action dim 1')
ax.set_title('Generated samples match the 6 modes')
ax.legend(loc='upper right', fontsize=9); ax.set_aspect('equal')

ax2 = axes[1]
traj = np.array(traj)
for i in range(0, 1500, 60):
    ax2.plot(traj[:, i, 0], traj[:, i, 1], color='steelblue', alpha=0.3, lw=0.7)
ax2.scatter(traj[0, ::20, 0], traj[0, ::20, 1], s=12, color='crimson', label='start (noise)')
ax2.scatter(traj[-1, ::20, 0], traj[-1, ::20, 1], s=12, color='seagreen', label='end (action)')
ax2.set_xlabel('action dim 0'); ax2.set_ylabel('action dim 1')
ax2.set_title('Integration paths from noise to actions')
ax2.legend(loc='upper right', fontsize=9); ax2.set_aspect('equal')

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5.5)) fig.suptitle('Flow matching carries noise to the action distribution', fontsize=13, fontweight='bold') ax = axes[0] ax.scatter(data[:, 0], data[:, 1], s=6, alpha=0.2, color='lightgray', label='target actions') ax.scatter(gen[:, 0], gen[:, 1], s=8, alpha=0.5, color='seagreen', label='generated (flow)') ax.set_xlabel('action dim 0'); ax.set_ylabel('action dim 1') ax.set_title('Generated samples match the 6 modes') ax.legend(loc='upper right', fontsize=9); ax.set_aspect('equal') ax2 = axes[1] traj = np.array(traj) for i in range(0, 1500, 60): ax2.plot(traj[:, i, 0], traj[:, i, 1], color='steelblue', alpha=0.3, lw=0.7) ax2.scatter(traj[0, ::20, 0], traj[0, ::20, 1], s=12, color='crimson', label='start (noise)') ax2.scatter(traj[-1, ::20, 0], traj[-1, ::20, 1], s=12, color='seagreen', label='end (action)') ax2.set_xlabel('action dim 0'); ax2.set_ylabel('action dim 1') ax2.set_title('Integration paths from noise to actions') ax2.legend(loc='upper right', fontsize=9); ax2.set_aspect('equal') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The generated green samples (left) sit right on the six target modes. The right panel traces integration paths: points start as a Gaussian blob (red) and flow outward along the learned field to land on the action modes (green).
2. Why it looks this way: Training regresses the velocity onto the straight line from noise to data, so integrating that velocity from a noise sample carries it to a valid action. The field naturally fans samples out to whichever mode they are closest to.
3. Key takeaway: Flow matching is a generative action head — it learns to transport noise to actions, the mechanism π0 uses instead of discretizing into tokens.

---

Experiment 2: Multimodal actions¶

Question: When several actions are valid, a regression head outputs their mean. Does the flow keep every mode?

We compare the flow-generated action distribution against the single point a mean-squared-error regressor would output.

In [4]:

gen2 = sample(3000, 30)
regression_point = data.numpy().mean(0)   # what an MSE head learns: the global mean
print(f'regression head output: {regression_point.round(2)}  (centre of the ring -> no valid action there)')

gen2 = sample(3000, 30) regression_point = data.numpy().mean(0) # what an MSE head learns: the global mean print(f'regression head output: {regression_point.round(2)} (centre of the ring -> no valid action there)')

regression head output: [-0.   -0.03]  (centre of the ring -> no valid action there)

In [5]:

fig, ax = plt.subplots(figsize=(6.6, 6.2))
fig.suptitle('Flow keeps all valid actions; regression collapses to the centre', fontsize=13, fontweight='bold')

ax.scatter(gen2[:, 0], gen2[:, 1], s=8, alpha=0.4, color='seagreen', label='flow samples (multimodal)')
ax.scatter(*regression_point, s=300, color='crimson', marker='X', zorder=5, label='regression output (invalid)')
ax.set_xlabel('action dim 0'); ax.set_ylabel('action dim 1')
ax.set_title('6 valid actions vs 1 averaged action')
ax.legend(loc='upper right', fontsize=9); ax.set_aspect('equal')

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(6.6, 6.2)) fig.suptitle('Flow keeps all valid actions; regression collapses to the centre', fontsize=13, fontweight='bold') ax.scatter(gen2[:, 0], gen2[:, 1], s=8, alpha=0.4, color='seagreen', label='flow samples (multimodal)') ax.scatter(*regression_point, s=300, color='crimson', marker='X', zorder=5, label='regression output (invalid)') ax.set_xlabel('action dim 0'); ax.set_ylabel('action dim 1') ax.set_title('6 valid actions vs 1 averaged action') ax.legend(loc='upper right', fontsize=9); ax.set_aspect('equal') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The flow samples (green) populate all six action modes around the ring, while the regression output (red X) sits at the dead centre — a point that is not a valid action at all.
2. Why it looks this way: A generative flow draws from the full action distribution, so each sample is one valid mode. An MSE regressor must output a single value and minimizes error by predicting the mean of the modes, which lands in the empty middle.
3. Key takeaway: Like discrete action tokens, flow matching handles multimodal actions — but in continuous space, giving smooth, precise, and diverse actions at once.

---

Experiment 3: Integration steps¶

Question: Sampling integrates an ODE — more steps cost more compute. How few Euler steps does flow matching need, and why do straight paths help?

We sample with increasing numbers of Euler steps and measure how close samples land to a true mode.

In [6]:

modes = 2.5 * np.stack([np.cos(2 * np.pi * np.arange(6) / 6), np.sin(2 * np.pi * np.arange(6) / 6)], 1)


def dist_to_mode(x):
    return np.min(np.linalg.norm(x[:, None, :] - modes[None], axis=2), axis=1).mean()


step_counts = [1, 2, 3, 5, 10, 20, 50]
samples_by_steps = {s: sample(1200, s) for s in step_counts}
errs = [dist_to_mode(samples_by_steps[s]) for s in step_counts]
print('steps     :', step_counts)
print('mean dist :', np.round(errs, 3))

modes = 2.5 * np.stack([np.cos(2 * np.pi * np.arange(6) / 6), np.sin(2 * np.pi * np.arange(6) / 6)], 1) def dist_to_mode(x): return np.min(np.linalg.norm(x[:, None, :] - modes[None], axis=2), axis=1).mean() step_counts = [1, 2, 3, 5, 10, 20, 50] samples_by_steps = {s: sample(1200, s) for s in step_counts} errs = [dist_to_mode(samples_by_steps[s]) for s in step_counts] print('steps :', step_counts) print('mean dist :', np.round(errs, 3))

steps     : [1, 2, 3, 5, 10, 20, 50]
mean dist : [2.163 1.233 0.759 0.436 0.309 0.276 0.254]

In [7]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('Straight-path flow matching needs only a few integration steps', fontsize=13, fontweight='bold')

ax = axes[0]
ax.plot(step_counts, errs, '-o', color='seagreen', lw=2)
ax.set_xlabel('Euler integration steps'); ax.set_ylabel('mean distance to nearest action mode')
ax.set_xscale('log'); ax.set_title('Sample quality vs compute')

ax2 = axes[1]
for s, color in [(1, 'crimson'), (3, 'orange'), (20, 'seagreen')]:
    g = samples_by_steps[s]
    ax2.scatter(g[:, 0], g[:, 1], s=8, alpha=0.4, color=color, label=f'{s} step(s)')
ax2.scatter(modes[:, 0], modes[:, 1], s=160, marker='*', color='black', zorder=5, label='modes')
ax2.set_xlabel('action dim 0'); ax2.set_ylabel('action dim 1')
ax2.set_title('Few steps already land near the modes')
ax2.legend(fontsize=8); ax2.set_aspect('equal')

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5)) fig.suptitle('Straight-path flow matching needs only a few integration steps', fontsize=13, fontweight='bold') ax = axes[0] ax.plot(step_counts, errs, '-o', color='seagreen', lw=2) ax.set_xlabel('Euler integration steps'); ax.set_ylabel('mean distance to nearest action mode') ax.set_xscale('log'); ax.set_title('Sample quality vs compute') ax2 = axes[1] for s, color in [(1, 'crimson'), (3, 'orange'), (20, 'seagreen')]: g = samples_by_steps[s] ax2.scatter(g[:, 0], g[:, 1], s=8, alpha=0.4, color=color, label=f'{s} step(s)') ax2.scatter(modes[:, 0], modes[:, 1], s=160, marker='*', color='black', zorder=5, label='modes') ax2.set_xlabel('action dim 0'); ax2.set_ylabel('action dim 1') ax2.set_title('Few steps already land near the modes') ax2.legend(fontsize=8); ax2.set_aspect('equal') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The mean distance to a mode drops quickly and flattens by ~5–10 steps; even 1–3 steps land most samples near the modes (right panel).
2. Why it looks this way: Flow matching trains on straight-line paths from noise to data, so the velocity field is close to constant along each path — a coarse Euler integration follows it accurately with few steps. Curved diffusion paths would need many more steps.
3. Key takeaway: Cheap few-step sampling is why π0 can run flow matching fast enough for 50Hz real-time control.

---

Experiment 4: Action chunking¶

Question: π0 predicts a chunk of future actions at once rather than one action per inference. How does that smooth control and reduce drift?

We compare a per-step policy (a fresh noisy action each step) against a chunked policy (commit to a smooth multi-step plan), tracking the executed trajectory.

In [8]:

np.random.seed(4)

T = 60
ref = np.stack([np.linspace(0, 6, T), np.sin(np.linspace(0, 6, T))], 1)   # smooth intended path

# per-step: each step adds independent prediction noise (re-decide every step)
per_step = ref + 0.25 * np.random.randn(T, 2)

# chunked: predict in chunks of H; noise is shared/smooth within a chunk
H = 12
chunked = ref.copy()
for s in range(0, T, H):
    e = min(s + H, T)
    chunked[s:e] = ref[s:e] + 0.25 * np.random.randn(1, 2)   # one offset per chunk -> smooth

jitter = lambda p: np.abs(np.diff(p, axis=0)).mean()
print(f'frame-to-frame jitter  per-step: {jitter(per_step):.3f}   chunked: {jitter(chunked):.3f}')

np.random.seed(4) T = 60 ref = np.stack([np.linspace(0, 6, T), np.sin(np.linspace(0, 6, T))], 1) # smooth intended path # per-step: each step adds independent prediction noise (re-decide every step) per_step = ref + 0.25 * np.random.randn(T, 2) # chunked: predict in chunks of H; noise is shared/smooth within a chunk H = 12 chunked = ref.copy() for s in range(0, T, H): e = min(s + H, T) chunked[s:e] = ref[s:e] + 0.25 * np.random.randn(1, 2) # one offset per chunk -> smooth jitter = lambda p: np.abs(np.diff(p, axis=0)).mean() print(f'frame-to-frame jitter per-step: {jitter(per_step):.3f} chunked: {jitter(chunked):.3f}')

frame-to-frame jitter  per-step: 0.272   chunked: 0.099

In [9]:

fig, ax = plt.subplots(figsize=(9, 5.4))
fig.suptitle('Action chunking yields smoother control than per-step decisions', fontsize=13, fontweight='bold')

ax.plot(ref[:, 0], ref[:, 1], color='black', lw=2.5, label='intended path')
ax.plot(per_step[:, 0], per_step[:, 1], '-o', ms=3, color='crimson', alpha=0.7,
        label=f'per-step (jitter {jitter(per_step):.2f})')
ax.plot(chunked[:, 0], chunked[:, 1], '-s', ms=3, color='seagreen',
        label=f'chunked H={H} (jitter {jitter(chunked):.2f})')
ax.set_xlabel('x'); ax.set_ylabel('y')
ax.set_title('Executed trajectory')
ax.legend(loc='upper left', fontsize=9)

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(9, 5.4)) fig.suptitle('Action chunking yields smoother control than per-step decisions', fontsize=13, fontweight='bold') ax.plot(ref[:, 0], ref[:, 1], color='black', lw=2.5, label='intended path') ax.plot(per_step[:, 0], per_step[:, 1], '-o', ms=3, color='crimson', alpha=0.7, label=f'per-step (jitter {jitter(per_step):.2f})') ax.plot(chunked[:, 0], chunked[:, 1], '-s', ms=3, color='seagreen', label=f'chunked H={H} (jitter {jitter(chunked):.2f})') ax.set_xlabel('x'); ax.set_ylabel('y') ax.set_title('Executed trajectory') ax.legend(loc='upper left', fontsize=9) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The per-step trajectory (red) jitters around the intended black path because each step re-decides independently; the chunked trajectory (green) is much smoother, following the intended path with far lower frame-to-frame jitter.
2. Why it looks this way: Committing to a coherent chunk of actions means the noise is shared across the chunk instead of resampled every step, so consecutive actions stay consistent and the motion is smooth.
3. Key takeaway: Predicting action chunks (here ~12 steps, in π0 ~50) gives temporally coherent, low-jitter control and reduces the compounding error of step-by-step replanning.

---

Summary¶

Experiment	Key insight
1. Flow matching	A velocity field trained on straight noise→data paths transports samples to valid continuous actions.
2. Multimodal actions	The generative flow keeps every valid action mode; a regressor averages them into an invalid action.
3. Integration steps	Straight paths integrate accurately in a few Euler steps — fast enough for 50Hz control.
4. Action chunking	Committing to a chunk of actions gives smooth, low-jitter control vs per-step replanning.

Connection to the broader theory¶

π0 is the continuous-action counterpoint to OpenVLA: instead of discretizing actions into tokens, it generates them with flow matching — a relative of normalizing flows. Together they bracket the design space a VLA ModelAdapter must cover. The tier now turns to its anchor model, LeWM, a large-scale embodied world model with a 3DGS-based latent that unifies world modelling and action.