Read this next to 1.4
Chapter 1.4 solved mode-collapse by learning to denoise: add Gaussian noise to an action in graded steps, train a network to predict the noise, then walk pure noise back to a sample with the DDPM reverse loop. It works — the ring came out multimodal, the policy beat its baseline. But the machinery is heavy: a noise schedule you have to get right, dozens of reverse steps, and fresh noise injected at every one that you then fight back down. This chapter keeps the idea — turn noise into a sample, so the model commits to one mode instead of averaging two — and replaces the machinery with something you can hold in your head.
Open flow.py beside diffusion.py. Same seven regions, the same 2D ring toy, the
same PushT policy, the same eval. The mechanism differs in just two places: the
objective (predict a velocity, not the noise) and the
sampler (integrate an ODE, don't run the DDPM posterior). Everything else — the
data pipeline, the training loop, the eval rollouts, the ONNX export — is the ch1.4
code, unchanged. This is the chapter: a single, legible swap, measured.
The swap, in one picture
Diffusion defines a stochastic process — a schedule of betas, an alpha-bar, a
posterior variance — and learns to invert it. Flow matching throws all of that away
and draws a straight line. Pick a data point, pick a noise point, and connect
them:
x_t = (1 - t) * noise + t * data for t in [0, 1]
At t = 0 you are at the noise the sampler starts from; at t = 1 you are at the
data. The velocity of a point sliding along that line is constant — it is just
data - noise, the same everywhere on the segment. So the whole learning problem
becomes: predict that velocity. Put a point on its line at a random t, ask the
network for the line's velocity, penalize the squared error. No schedule, no
posterior, no injected noise — the geometry does the work the schedule used to.
Build
flow.py is one file in the same seven regions as ch1.4: setup, core, toy, data,
train, eval, report. The core region is where diffusion and flow diverge — and
notice what is missing from it versus ch1.4: no make_schedule, no betas/
alphas, no posterior variance. A straight line needs none of them.
The core: velocity in, sample out
The objective places a sample on its straight noise→data line at a random time and regresses the network onto that line's constant velocity $x_0 - \epsilon$ — and here it is in code:
# The flow-matching machinery, shared by toy and policy. Nothing from a generative# library — the interpolation, the loss, and the ODE sampler ARE the method. Note# what is MISSING versus ch1.4: no noise schedule, no betas/alphas, no posterior# variance. A straight line needs none of it. def sinusoidal_embed(t: torch.Tensor, dim: int) -> torch.Tensor: """Continuous flow time (B,) -> (B, dim) sinusoidal features (the same time embedding as ch1.4 / transformers). Lets one network condition smoothly on WHERE along the noise->data path it is, without a separate weight per time.""" half = dim // 2 freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device) / half) ang = t.float()[:, None] * freqs[None] return torch.cat([ang.sin(), ang.cos()], dim=1) def interpolate(x0: torch.Tensor, noise: torch.Tensor, t: torch.Tensor) -> torch.Tensor: """The conditional flow path: a STRAIGHT LINE from a noise point (t=0) to a data point x0 (t=1). x_t = (1-t)*noise + t*x0. This one line replaces ch1.4's entire noise schedule + q_sample — the path is fixed by geometry, not a learned/tuned schedule. Its velocity (below) is the constant x0 - noise, the same everywhere on the line.""" return (1.0 - t)[:, None] * noise + t[:, None] * x0 def flow_matching_loss(model: nn.Module, x0: torch.Tensor, cond) -> torch.Tensor: """The entire training objective (conditional flow matching). Pick a random time t in [0,1] per sample, put the sample on its straight noise->data line at t, and ask the network to PREDICT THAT LINE'S VELOCITY, x0 - noise. Plain MSE on the velocity — but because the target is a velocity toward a SAMPLED data point, not the action itself, minimizing it never averages two actions into a third one. (ch1.4 predicted the noise instead; this is the ~2-line objective swap.)""" t = torch.rand(len(x0), generator=gen) noise = torch.randn(x0.shape, generator=gen).to(device) x_t = interpolate(x0, noise, t.to(device)) target_v = (noise - x0) if BROKEN_TARGET else (x0 - noise) # Break It flips this sign return F.mse_loss(model(x_t, t.to(device), cond), target_v) @torch.no_grad()def ode_sample_loop(model: nn.Module, shape, cond, steps: int, log=None) -> torch.Tensor: """Sampling = integrate the learned velocity field forward, from noise (t=0) to a sample (t=1), with plain forward Euler. Each step just follows the arrow: x <- x + dt * v. No fresh noise, no posterior mean, no clip — the whole ch1.4 reverse step collapses to this. Straight conditional paths mean few steps get you there; `log` records the moving cloud to rerun so you can watch noise slide into data.""" x = randn(shape) # start at pure noise, flow time t = 0 dt = 1.0 / steps for i in range(steps): t = torch.full((shape[0],), i * dt, device=device) x = x + dt * model(x, t, cond) # forward Euler along the velocity field if log is not None: rr.set_time("flow_step", sequence=i + 1) rr.log(log, rr.Points3D(_xy0(x))) return x class VelocityNet(nn.Module): """The velocity predictor, used for BOTH the toy and the policy. In: a point x_t on a noise->data line, its flow time t, and (optionally) a conditioning vector (the obs for the policy, nothing for the toy). Out: the predicted velocity, x_t's shape. Just an MLP — the lesson is the OBJECTIVE and the SAMPLER, not the architecture (real flow policies use a U-Net / transformer; see 'What we cut'). Structurally identical to ch1.4's Denoiser — only the name and what it predicts change. When conditioned, obs normalization lives inside as buffers so a checkpoint carries its own stats.""" def __init__(self, x_dim: int, cond_dim: int, hidden: int, obs_min=None, obs_range=None): super().__init__() self.cond_dim = cond_dim self.net = nn.Sequential( nn.Linear(x_dim + TIME_DIM + cond_dim, hidden), nn.SiLU(), nn.Linear(hidden, hidden), nn.SiLU(), nn.Linear(hidden, x_dim), ) if cond_dim and obs_min is not None: self.register_buffer("obs_min", torch.from_numpy(obs_min)) self.register_buffer("obs_range", torch.from_numpy(obs_range)) def forward(self, x_t: torch.Tensor, t: torch.Tensor, cond) -> torch.Tensor: parts = [x_t, sinusoidal_embed(t * TIME_SCALE, TIME_DIM)] # scale t in [0,1] up before embedding if self.cond_dim: # normalize obs to [-1, 1] inside the model, then condition on it parts.append((2.0 * (cond - self.obs_min) / self.obs_range - 1.0).clamp(-1.0, 1.0)) return self.net(torch.cat(parts, dim=1)) def _xy0(a) -> np.ndarray: """(N, >=2) points -> (N, 3) with z=0 for rerun's 3D point cloud.""" a = a.detach().cpu().numpy() if hasattr(a, "detach") else np.asarray(a) return np.column_stack([a[:, :2], np.zeros(len(a), dtype=np.float32)]).astype(np.float32)Three functions carry the method. interpolate is the straight noise→data line —
one line of code that replaces ch1.4's entire schedule plus q_sample.
flow_matching_loss is the whole training objective: draw a random time t, place
the sample on its line, and predict the line's velocity data - noise with plain
MSE. (Compare diffusion_loss, which predicted the noise on a scheduled noising —
that is the ~2-line objective swap.) ode_sample_loop is sampling: start at noise,
and take forward-Euler steps x <- x + dt * v up the velocity field to t = 1. The
entire DDPM reverse step — recover x0, clip it, blend via the posterior mean,
re-inject noise — collapses to that one addition. The network itself, VelocityNet,
is structurally identical to ch1.4's Denoiser: the same MLP, the same sinusoidal
time embedding, the same optional obs conditioning. Only its name and what it
predicts change.
One necessary difference from ch1.4: flow time t lives in [0, 1] — a continuous
position along the line, not an integer step index — so we scale it up by TIME_SCALE
before the sinusoidal embedding, or the embedding would barely vary across the whole
interval.
The toy: same aha, plus efficiency
# The aha, before any robot. Target: a unit ring — every angle an equally-good# "mode", the center the single point NO sample occupies, exactly where an# averaging regressor lands. The ch1.1 left-vs-right collapse in 2D, where we SEE it.angle = torch.rand(N_TOY, generator=gen) * 2.0 * math.piring = torch.stack([angle.cos(), angle.sin()], dim=1) + 0.03 * torch.randn(N_TOY, 2, generator=gen)ring = ring.to(device) # (a) Flow matching: train the velocity net to predict x0 - noise on ring points.toy = VelocityNet(2, 0, args.model_dim).to(device)toy_opt = torch.optim.Adam(toy.parameters(), lr=args.lr)toy_loss = float("nan")for it in range(TOY_ITERS): idx = torch.randint(0, N_TOY, (256,), generator=gen) loss = flow_matching_loss(toy, ring[idx], None) toy_opt.zero_grad() loss.backward() toy_opt.step() toy_loss = loss.item() # (b) Regression baseline: a same-WIDTH MLP (smaller than the time-conditioned# velocity net — which only helps the point), trained to map noise -> a point with# MSE. Fresh noise says nothing about WHICH ring point, so the MSE-optimal output is# the average of them all — the empty center. Objective-driven, robust to capacity.# IDENTICAL to ch1.4's baseline: the collapse is about the loss, not the generator.regress = nn.Sequential(nn.Linear(2, args.model_dim), nn.SiLU(), nn.Linear(args.model_dim, 2)).to(device)reg_opt = torch.optim.Adam(regress.parameters(), lr=args.lr)for it in range(TOY_ITERS): idx = torch.randint(0, N_TOY, (256,), generator=gen) loss = F.mse_loss(regress(randn((256, 2))), ring[idx]) reg_opt.zero_grad() loss.backward() reg_opt.step() # Sample from both and MEASURE multimodality. `modes_covered` bins on-ring samples# into 8 angular sectors: flow should light up all 8, regression none.flow_samples = ode_sample_loop(toy, (N_TOY, 2), None, args.flow_steps, log="toy/flow" if args.rerun else None).cpu().numpy()with torch.no_grad(): reg_samples = regress(randn((N_TOY, 2))).cpu().numpy() def ring_stats(samples: np.ndarray) -> tuple[float, float, int]: r = np.hypot(samples[:, 0], samples[:, 1]) on_ring = np.abs(r - 1.0) < 0.2 sector = (np.floor((np.arctan2(samples[:, 1], samples[:, 0]) + math.pi) / (2 * math.pi) * 8).astype(int) % 8) return float(r.mean()), float(on_ring.mean()), int(len(np.unique(sector[on_ring]))) flow_r, flow_hit, flow_modes = ring_stats(flow_samples)reg_r, reg_hit, reg_modes = ring_stats(reg_samples)# The chapter's flow-specific measurement: re-sample the SAME trained net with only# EFF_STEPS Euler steps. Flow's straight paths still land the modes at a step budget# where ch1.4's DDPM reverse loop is still a blur — the efficiency claim, measured.lowstep_samples = ode_sample_loop(toy, (N_TOY, 2), None, EFF_STEPS).cpu().numpy()_, lowstep_hit, lowstep_modes = ring_stats(lowstep_samples)print(f"toy multimodality [measured]: flow mean_radius {flow_r:.2f} ring_hit {flow_hit:.2f} modes {flow_modes}/8" f" | regression mean_radius {reg_r:.2f} ring_hit {reg_hit:.2f} modes {reg_modes}/8")print(f"toy step efficiency [measured]: flow covers {lowstep_modes}/8 modes with only {EFF_STEPS} Euler steps " f"(vs {flow_modes}/8 at {args.flow_steps}) — straight paths need few steps")if args.rerun: # forward: watch the ring slide into noise along the straight paths (negative time = before sampling) for k in range(7): t_k = torch.full((N_TOY,), k / 6.0, device=device) rr.set_time("flow_step", sequence=k - 7) rr.log("toy/forward", rr.Points3D(_xy0(interpolate(ring, randn((N_TOY, 2)), t_k)))) rr.log("toy/target", rr.Points3D(_xy0(ring), colors=(140, 140, 150)), static=True) rr.log("toy/regression", rr.Points3D(_xy0(reg_samples), colors=(230, 102, 90)), static=True)This is the ch1.4 toy with two lines swapped (flow_matching_loss for
diffusion_loss, ode_sample_loop for p_sample_loop) — and it makes the same
point. Flow covers 8/8 angular modes of the ring at radius 0.94; the
same-width MSE regressor — identical code to ch1.4's baseline — covers 0/8 at
radius 0.06, the dead center. The collapse was never about diffusion; it is about
sampling vs averaging, and flow samples.
What is new here is the last measurement. We re-sample the same trained net at only
5 Euler steps and it still covers 8/8 modes. Sweep the step count and the
pattern is clean: full mode coverage appears by 3 steps, ring quality keeps
tightening out to ~20 and then saturates, and only 2 steps is too coarse for
Euler to resolve the ring (that is the few_steps break). The point is not just "few
steps work" — it is that sampling steps are decoupled from training. Diffusion's
step count is welded to its schedule; change it and you retrain. Flow's step count is
a free choice you make at sampling time, from one trained network. That is what
"straighter and faster" actually cashes out to, and it is measured, not asserted.
The policy: conditioning on the observation
# The policy trains on the SAME scripted PushT demos as ch1.1 / ch1.4 — the only# change from ch1.4 is the objective (predict velocity) and the sampler (integrate).if args.data is None: # Regenerate every run (never reuse a leftover dir): a cache from a different # --seed/--num_demos would silently train on the wrong data. gen_demos is # deterministic, so same args -> bit-identical demos, built or rebuilt. args.data = args.out / "demos" if args.data.exists(): shutil.rmtree(args.data) gen_demos.main(["--episodes", str(args.num_demos), "--seed", str(args.seed), "--out", str(args.data), "--no-video"])if not (args.data / "meta" / "info.json").is_file(): sys.exit(f"no dataset at {args.data} — record one in ch0.4, or generate demos:\n" f" python curriculum/common/envs/pusht/gen_demos.py " f"--episodes 500 --seed 0 --out {args.data} --no-video") from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: E402 (heavy import — after cheap failures) frames = LeRobotDataset("local/pusht-demos", root=args.data).hf_dataset.with_format("numpy")obs = np.stack(frames["observation.state"]).astype(np.float32) # (N, 10) — layout in pusht_env.pyactions = np.stack(frames["action"]).astype(np.float32) # (N, 2) — already clipped to [-1, 1] # The ODE starts from N(0, I), so the data must sit at ~unit scale or the sampler# fights a scale mismatch: these expert actions have std ~0.3, and standardizing# them to zero-mean/unit-std is what takes this policy from 0% to working (measured,# same as ch1.4). So we STANDARDIZE the flowed actions and MIN-MAX normalize the# conditioning obs (constant obs dims get range 1 -> a constant, not a divide-by-zero).obs_min = obs.min(0)obs_range = np.where(obs.max(0) - obs_min < 1e-4, np.float32(1.0), obs.max(0) - obs_min)act_mean = actions.mean(0)act_std = np.where(actions.std(0) < 1e-4, np.float32(1.0), actions.std(0))obs_t = torch.from_numpy(obs).to(device)act_t = torch.from_numpy((actions - act_mean) / act_std).to(device)act_mean_t = torch.from_numpy(act_mean).to(device)act_std_t = torch.from_numpy(act_std).to(device)print(f"dataset: {len(np.unique(frames['episode_index']))} episodes / {len(obs)} frames, " f"flow_steps={args.flow_steps}, model_dim={args.model_dim}")The data region is ch1.4's, unchanged: the same scripted PushT demos, the same
normalization. The ODE starts from N(0, I) exactly as the DDPM sampler did, so the
actions we flow must sit at unit scale too — standardizing them is what takes this
policy from 0% to working, the same lesson as ch1.4.
# Chapter 1.4's loop, one line changed: the loss is flow_matching_loss (predict the# velocity of a noised action's straight path, conditioned on the obs) instead of the# diffusion noise-prediction loss. As in ch1.4 we DON'T cosine-decay the lr — the net# must keep fitting the velocity at all times t; decaying it undertrains (measured).torch.manual_seed(args.seed) # policy init + training-noise stream reproducible, independent of the toygen.manual_seed(args.seed) # give the policy a fresh noise stream (the toy drained the shared gen above)policy = VelocityNet(ACT_DIM, OBS_DIM, args.model_dim, obs_min, obs_range).to(device)optimizer = torch.optim.Adam(policy.parameters(), lr=args.lr)shuffle = torch.Generator().manual_seed(args.seed + 1) # torch-side RNG for batch ordertrain_loss, global_step = float("nan"), 0for epoch in range(args.epochs): epoch_loss, num_batches = 0.0, 0 for batch in torch.randperm(len(obs_t), generator=shuffle).split(args.batch_size): loss = flow_matching_loss(policy, act_t[batch], obs_t[batch]) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss, num_batches = epoch_loss + loss.item(), num_batches + 1 if args.rerun: rr.set_time("step", sequence=global_step) rr.log("policy/loss/train", rr.Scalars([loss.item()])) global_step += 1 train_loss = epoch_loss / num_batches if epoch % 50 == 0 or epoch == args.epochs - 1: print(f"epoch {epoch:4d} flow_mse {train_loss:.5f}")The training loop is ch1.4's with one line changed: the loss is flow_matching_loss.
As in ch1.4 we do not cosine-decay the learning rate — the velocity net must keep
fitting the field at every time t, and decaying it undertrains the net.
# Loss measured on velocity fits; rollouts measure the task. At every env step we# SAMPLE an action: start from noise, integrate the ODE conditioned on the current# obs. --break few_steps runs this at flow_steps=2 (set in setup): too few steps# can't follow the curved MARGINAL field, the action lands off the data, and the# pusher wanders — a measured drop, not a crash.def rollout(net: VelocityNet, seed: int, tag: str, episode: int) -> tuple[bool, float]: env = PushTEnv() obs_now = env.reset(seed) gen.manual_seed(seed) # seed the sampler from the episode seed: reproducible AND order-independent, so # baseline and trained see the same per-episode noise (a fair, seeded comparison) done, episode_return, info = False, 0.0, {} while not done: cond = torch.from_numpy(obs_now).to(device).unsqueeze(0) # (10,) -> (1, 10) sample = ode_sample_loop(net, (1, ACT_DIM), cond, args.flow_steps) # in standardized action space action = (sample * act_std_t + act_mean_t).clamp(-1.0, 1.0)[0].cpu().numpy() obs_now, reward, done, info = env.step(action) episode_return += reward if args.rerun: rr.set_time("sim_time", duration=episode * (PushTEnv.MAX_STEPS / PushTEnv.CONTROL_HZ) + env.data.time) rr.log(f"eval/{tag}/action", rr.Scalars(action.astype(np.float64))) rr.log(f"eval/{tag}/pos_err", rr.Scalars([info["pos_err"]])) return bool(info["success"]), episode_return def evaluate(net: VelocityNet, tag: str) -> tuple[float, float]: # 10_000 + offset: held out from demo seeds (0..num_demos) by construction. outcomes = [rollout(net, 10_000 + args.seed + ep, tag, ep) for ep in range(args.eval_episodes)] success_rate = float(np.mean([s for s, _ in outcomes])) mean_return = float(np.mean([r for _, r in outcomes])) if args.rerun: rr.log(f"eval/{tag}/success_rate", rr.Scalars([success_rate])) print(f"eval[{tag:9s}]: success {success_rate:.2f} mean_return {mean_return:.3f}") return success_rate, mean_return torch.manual_seed(args.seed + 1) # a fixed random-init reference, independent of eval-time RNGbaseline = VelocityNet(ACT_DIM, OBS_DIM, args.model_dim, obs_min, obs_range).to(device)baseline_success, baseline_return = evaluate(baseline, "untrained")success_rate, mean_return = evaluate(policy, "trained")At every environment step we sample an action by integrating the ODE conditioned on the current observation, un-standardize, clip, step. The trained policy reaches 40% held-out success versus 0% for the untrained net, and it holds across seeds (0.40, 0.40, 0.55). Now be precise about what that 40% means, because this is where flow matching is easy to oversell.
It edges chapter 1.4's diffusion policy — ~25% at the same 100-demo budget — a modest but real gap, plausibly because straight-line integration is a steadier hand than DDPM's noisy reverse walk on a task whose good policy is nearly deterministic. It does not, however, beat chapter 1.1's behavior cloning: BC at its full 500-demo default reaches 62%, well above flow here, and only at a matched 100-demo budget does BC fall back into flow's neighborhood. Flow matching does not buy you a better policy — it buys you cheaper sampling at the same quality, which is a different and more honest claim. And this is 20 episodes (noisy — chapter 1.6 is about exactly this), so treat even the diffusion gap as a directional win, not a headline. The rock-solid results are the toy and the step-efficiency above.
Run it
python curriculum/phase1_imitation/ch1.5_flow/flow.py --seed 0 --device cpu
| cpu-laptop | expected wall-clock on cpu-laptop: ~1.18 min (measured) | measured |
| mps | wall-clock on mps: not yet measured | pending |
| t4 | expected wall-clock on t4: ~8.94 min (measured) | measured |
| 4090 | wall-clock on 4090: not yet measured | pending |
| held-out success | mean return | |
|---|---|---|
| untrained (random-init velocity net) | 0% | −106 |
| trained flow policy | 40% | −89 |
rerun outputs/ch1.5-flow/flow.rrd
What we cut
This is real conditional flow matching, trained the real way, but it is not a full flow-matching policy like pi0, and the gaps are the same shape as ch1.4's:
- We flow a single action, not an action horizon. Real flow policies flow a chunk of future actions jointly and execute them receding-horizon — the source of temporal coherence and much of the real performance.
- The velocity net is an MLP, not a temporal U-Net / transformer. Legible here; a bigger backbone is what you reach for with a horizon and images.
- We condition on the 10-number state, not images or language. pi0 conditions a flow-matching action head on a VLM backbone — that is chapter 1.8.
- We integrate with plain forward Euler. Real implementations use a higher-order or adaptive ODE solver (fewer, better steps). Euler is the readable floor; the toy already shows 3–5 steps suffice here.
None of these silently degrades a number; each is a whole capability left for later so the velocity→sample core stays readable. The "read the real thing" segment walks a production flow-matching repo so you can see exactly what these paragraphs left out.
Break it
Two ablations, each a real flow-matching misconception with a measured signature.
--break few_steps — "two Euler steps is plenty." This forces flow_steps = 2
at sampling time only. Watch two things. First, the toy degrades — 8/8 modes to
3/8, radius 0.94 to 0.51 — because two Euler steps overshoot the curved
marginal field the net actually learned (each conditional path is straight, but the
field averaged over which endpoint is not). Second, and this is the honest twist: the
policy is unharmed (45%, if anything up from 40), because the PushT action
distribution is nearly unimodal and even a crude two-step integration lands close
enough. The same two-step regime that collapsed ch1.4's diffusion policy to 0%
leaves flow's policy standing. And note the training loss: it is bit-identical to
the full run (0.951594), because flow_steps is a sampler knob that never touches
the objective. That is the decoupling, made unmistakable.
--break wrong_target — "the sign of the velocity can't matter that much." This
flips the target to noise - data, so the net learns to flow away from the data.
The ODE then integrates outward: the toy explodes off the ring to radius 4.12, the
policy drops to 0% — and the training loss stays low (0.935), because the net
fits the flipped target perfectly well. Low loss, wrong direction: the sign of the
velocity is the direction of the flow, and the loss curve will not warn you. (The
ch1.4 wrong_schedule trap, in flow's language.)
Read the real thing
The paper that named this objective ships a reference library — facebookresearch/flow_matching, pinned here at commit 11568d3. It is not a policy; it is the objective and the sampler you just built, generalized and hardened. Read it in three passes, against your core region.
The path and the velocity target. Your interpolate hard-codes one straight
line, (1-t)*noise + t*x0, and flow_matching_loss regresses its constant velocity
x0 - noise. The library's AffineProbPath.sample() in
flow_matching/path/affine.py builds the same point as x_t = sigma_t*x_0 + alpha_t*x_1 and returns its velocity as dx_t = d_sigma_t*x_0 + d_alpha_t*x_1
inside a PathSample (flow_matching/path/path_sample.py); the training loss is one
line in that file's docstring — mse_loss(path_sample.dx_t, model(x_t, t)), your
loss exactly. Your straight line is their CondOTProbPath, whose CondOTScheduler
(flow_matching/path/scheduler/scheduler.py) sets alpha_t=t, sigma_t=1-t, d_alpha_t=1, d_sigma_t=-1. What the library adds is choice: a scheduler object so
the path can be any affine schedule (there is a PolynomialConvexScheduler for
curved ones), plus a full set of conversions — target_to_velocity,
epsilon_to_velocity, velocity_to_epsilon — so you can train in data-, noise-, or
velocity-space and move between them. You cut every bit of it because this chapter is
about the one path where none of it is needed.
The sampler. Your ode_sample_loop is the whole integrator: x <- x + dt*v,
steps times. That is precisely the method="euler" branch of ODESolver.sample()
in flow_matching/solver/ode_solver.py, which hands the velocity field to
torchdiffeq.odeint. What production adds is the rest of the menu — dopri5,
midpoint, heun3: adaptive and higher-order solvers that reach the data in fewer,
better steps than Euler — plus a compute_likelihood() that integrates the field's
divergence for an exact log-likelihood, a capability your loop does not have. The toy
already showed 3–5 Euler steps suffice here; a stiffer field is where those solvers
earn their keep.
The horizon and the VLM — honestly, not here. This is the cut to be precise
about: this repo has no action horizon and no VLM, because it is a
generative-modeling library, not a robot policy. Those pieces live one repo further
out — in a flow-matching policy like pi0 / openpi, built on this exact objective —
and that is chapter 1.8. What this repo does show of the production direction is
examples/image/: the same velocity loss wired to a real U-Net (models/unet.py),
EMA (models/ema.py), and distributed training — the scaffolding, minus the robot.
Read these, in order. examples/2d_flow_matching.ipynb first — your ring toy in
their hands. Then flow_matching/path/affine.py and
flow_matching/solver/ode_solver.py — your three core functions, generalized.
Then examples/image/train.py, to watch the same loss survive a real backbone. The
action head that flows a chunk of actions on a VLM backbone is where 1.8 picks up.
Exercises
Four, in exercises/. Two ask you to commit to a prediction before the run answers —
the ring's flow-vs-regression modes (the same result as ch1.4, which is the point),
and flow-vs-diffusion in the few-step regime. One is a bug-hunt in the interpolation
(the two time coefficients are swapped, so the path runs data→noise; the loss still
falls). One has you implement the Euler ODE sampler from its definition — the one
update the whole method turns on.
What's next
You now have two generative policies that sample actions — diffusion and flow — and
a ring that shows why sampling beats averaging, plus a step-efficiency result that
says flow gets there in a handful of straight steps. Chapter 1.8 puts this exact
velocity objective to work as the action head of a tiny vision-language-action
model: the same data - noise target, now conditioned on a VLM backbone instead of a
ten-number state, flowing a chunk of actions instead of one. The sampler you just
wrote is the last piece before the VLA.