Flow Matching from Absolute Zero

What You'll Master

Where We Are

Same Flappy Bird environment as Problem 1. Same 4-D observation, same 20-step action chunk, same hard mode with alternating single and double-gap pipes.

What's different: the algorithm. Instead of training an MLP to regress the expert's action via MSE (Problem 1's approach), you'll train a generative model that learns the entire distribution of expert actions at each state. At inference time, you sample from that distribution to get an action chunk.

The deliverable for Problem 2:

1. Implement FlowMatchingSchedule.interpolate — sample noise and build the noisy training input.
2. Implement FlowMatchingSchedule.sample — integrate the learned ODE from noise to a clean action.
3. Implement flow_matching_loss — MSE on velocity prediction.
4. Run on hard mode — report mean and std episode length, explain the result in 2-3 sentences.

By the end of this guide you'll know:

• What a generative model is and how it differs from a regressor.
• What a vector field is and why it can transport noise to data.
• The flow matching training objective: regress velocity along straight-line interpolations.
• How Euler integration turns a learned vector field into a sampling procedure.
• Why flow matching solves the multimodality problem that killed P1's MSE regression.
• How flow matching relates to diffusion (it's a strict simplification).

Why Flow Matching Fixes P1

Recall the climax of Problem 1: hard mode has multimodal experts. Two valid expert actions per state — gap 1 (top) or gap 2 (bottom). MSE regression converges to the conditional mean of expert actions, which is the midpoint between the gaps — the wall.

The fundamental issue was that MSE assumes a unimodal target distribution. We want a method that can represent multiple modes simultaneously: at this state, the expert sometimes does X, sometimes Y, and the policy should be able to sample one or the other rather than averaging.

The two paths to fix multimodality

Two orthogonal approaches:

1. Make the model richer — train a generative model that learns the full conditional distribution p(a | s), not just its mean. Sample from this distribution at inference. This is Problem 2: flow matching.
2. Make the data unimodal — use a deterministic expert that always picks the same gap. The data is now single-modal and MSE works fine. (Problem 3: DAgger.)

Both work. Modern robot learning systems often use both: a generative policy (flow matching, diffusion) trained on the cleanest data possible.

Generative Models for BC

Step back from flow matching for a moment. What does it mean to "model the conditional distribution p(a | s)" rather than just its mean?

What a regressor does

The MSE-trained policy from Problem 1:

Pass in a state, get back one action. The same state always gives the same action.

What a generative model does

A generative policy is sampled, not evaluated:

Pass in a state, get back one sample from the conditional distribution. Same state, different samples each time. If the distribution has two modes (gap 1 and gap 2), then sometimes the sample lands at gap 1, sometimes at gap 2.

For the Flappy Bird problem, this is exactly what you want. Each individual rollout commits to a coherent path (sample gap 1, target gap 1 for all 20 future steps). Different rollouts may sample different gaps. None of them takes the middle.

The hard part: how do you parameterize a distribution?

For the Gaussian case, easy: a network outputs mean and standard deviation. Sample from the resulting Gaussian. But that only gives unimodal distributions.

For arbitrary distributions, harder. There are several common approaches:

Diffusion and flow matching are the dominant choices in modern robot learning. They share the same overall structure (noise → data via iterative refinement) but differ in implementation details.

The Flow Matching Idea

Flow matching's central insight: generate samples from a complex distribution by continuously transporting samples from a simple distribution along a learned vector field.

The setup

Two distributions of interest:

• Source distribution: a Gaussian N(0, I). Easy to sample from.
• Target distribution: the expert's conditional action distribution p(a | s). Hard to sample from directly — that's what we want to learn.

We want a way to take a sample from the source and transform it into a sample from the target.

The vector field idea

Imagine the action space (a 20-D box, but think of it as a 2-D plane for visualization). At each point in this space, draw an arrow pointing toward where samples should "flow" if we want to move from noise to data.

What we want from the field

If we sample x₀ ~ N(0, I) at time τ = 0 and follow the velocity arrows for time τ from 0 to 1, we should end up at x₁ ~ p(a | s). Mathematically:

This is an ordinary differential equation. The velocity v_θ(x, s, τ) is a neural network parameterized by θ. Different network outputs give different vector fields, which produce different transport behaviors.

The training problem

We have data — samples from p(a | s) — and noise — samples from N(0, I). We want to train v_θ so that integrating from noise to data along the field gives the right answer.

The key trick: we don't need to integrate during training. We can derive the target velocity at any (x, s, τ) directly. That's what makes flow matching simple to train.

The next chapter shows how.

Training: Velocity Regression

The flow matching training objective is breathtakingly simple. Three steps.

Step 1: pick a clean action and a noise sample

Take a₁ from the expert dataset (a clean action chunk, shape [20]). Sample a₀ ~ N(0, I) independently of a₁ (also shape [20]).

Step 2: pick a random time and interpolate

Sample τ uniformly from [0, 1]. Compute the linear interpolation:

At τ = 0, this is pure noise a₀. At τ = 1, this is the clean action a₁. At τ = 0.5, it's a 50/50 mix. As τ varies, a_τ traces a straight line from noise to data.

Step 3: the target velocity is the straight-line slope

Differentiate a_τ with respect to τ:

This is the velocity that, if the network produced it everywhere along the line from a₀ to a₁, would correctly transport a₀ to a₁ over time τ ∈ [0, 1].

The training loss: regress the network's output v_θ(a_τ, s, τ) toward this target velocity:

Why this works (the surprising part)

You might worry: we're training the network on a specific noise sample a₀, but at inference time we'll sample a fresh different noise sample. The network has only seen one specific path from a₀ to a₁; how does it know what to do for arbitrary noise inputs?

Answer: averaging across many noise samples and many data samples, the network learns the conditional expectation of the velocity at each (a_τ, s, τ) point:

Across all (a₀, a₁) pairs that pass through point x at time τ, the average target velocity is the value the network learns. This expected velocity, integrated over time, transforms the source distribution N(0, I) into the target distribution p(a | s). That's the theorem of Conditional Optimal Transport Flow Matching (Lipman et al., 2023).

The algorithm

Steps 1-5 happen in your FlowMatchingSchedule.interpolate. Steps 6-8 happen in your flow_matching_loss.

Sampling: Euler Integration

Training is done. Now we have a learned velocity field v_θ(x, s, τ). To generate an action chunk for state s, we integrate the ODE from τ = 0 to τ = 1.

The ODE

We want x(1), starting from x(0). There's no closed form — we have to integrate numerically.

Euler's method

The simplest numerical ODE solver. Discretize τ into n steps of size h = 1/n. At each step, take a small move in the direction of the velocity:

"From the current position, take a step of size h in the direction of the velocity at the current position and time." Repeat n times to get from τ = 0 to τ = 1.

For this homework, num_steps = 20, so h = 1/20 = 0.05.

The full sampling procedure

Why this works (intuition)

The trained velocity field, at any point along an interpolation path, points toward the data distribution. Following the field's arrows traces out a curve that, over time, transports the noise sample toward a data sample.

For a multimodal target like the gap-1-or-gap-2 distribution, samples from N(0, I) are partitioned by the field: those near the "gap 1 attractor" flow toward gap 1; those near the "gap 2 attractor" flow toward gap 2. The exact path depends on the initial noise sample, so different runs of the sampler give different modes — that's how flow matching captures multimodality.

The trade-off in num_steps

Flow matching with conditional optimal transport (the variant in this homework) tends to need fewer integration steps than diffusion, because the paths between noise and data are straight lines rather than curved. That's one of the practical advantages of FM over diffusion.

FM vs Diffusion

Flow matching and diffusion are the two dominant generative paradigms in modern robotics (Diffusion Policy, RT-2, π₀, …). They share the same core idea but differ in implementation.

Why FM has straight paths

In diffusion, you add small Gaussian increments many times. The cumulative trajectory from data to noise is a random walk, which is curved. To generate data, you reverse this curved process — needing many small steps to follow the curve accurately.

In flow matching, you define the path as a straight line from data to noise. The velocity along this straight line is constant (it's just a₁ − a₀). The network only has to learn average velocities along straight paths — an easier function class. Sampling traces the straight paths back, which can be done with fewer Euler steps.

Are they really different?

Mathematically, no — they're both ways of learning a transport map between two distributions, and there are formal connections between them. Practically, yes — the implementation details differ. Flow matching's simplicity is what makes it the default choice for new code in 2024+ robot learning papers.

The U-Net Architecture

For Problem 1, the policy was a 3-layer MLP — you wrote it. For Problem 2, the velocity network is much fancier: a 1D temporal U-Net with conditional residual blocks. The starter code provides this entire architecture; you don't write any of it. But understanding what it does makes the rest of the homework click.

What the U-Net does

It implements v_θ(noisy_action, state, timestep):

Why a U-Net (vs another MLP)?

For the BC regressor in P1, an MLP was fine. For flow matching, the function we're learning is much richer: it depends on three inputs (noisy action, state, time) and has to handle the entire dynamics of pushing noise toward data. An MLP can do this in principle, but a 1D U-Net is the standard architecture for diffusion/FM in time-series spaces because:

• The action chunk has a temporal structure (20 future actions). Convolutions across the time axis exploit this.
• U-Nets have a multi-scale structure (downsample to coarse features, upsample with skip connections). This helps capture both high-frequency details and long-range structure.
• The conditional residual blocks let the network modulate its computation by the timestep and state (via FiLM-style scale-and-bias). This is how the network "knows" what time it is.

Black-box view (you don't need to understand the internals)

For the purposes of Problem 2, you can treat the model as a function:

model = policy.model       # the U-Net
schedule = policy.schedule  # the FlowMatchingSchedule (you implement)

# During training: predict velocity at one (a_tau, s, tau) point
v_pred = model(a_tau, state, tau)        # shape [B, 20]

# During inference: integrate the ODE
sampled_action = schedule.sample(model, state)

The FlowMatchingPolicy wrapper class bundles both together so you can pass policy around as a single object.

Code Tour

Same files as Problem 1. Different blanks.

Already provided

• SinusoidalPosEmb, Conv1dBlock, ConditionalResidualBlock1D, ConditionalUnet1D — the U-Net building blocks.
• TemporalNoisePredictor — wraps the U-Net to match the (noisy_action, state, timestep) signature.
• FlowMatchingPolicy — bundles the model and schedule together.

What you fill in

The training loop (already done)

Same generic loop as Problem 1, just with different policy and loss:

policy = FlowMatchingPolicy(...)
optimizer = optim.Adam(policy.parameters(), lr=1e-3)

for epoch in range(num_epochs):
    for s_batch, a_batch in loader:
        loss = flow_matching_loss(policy, s_batch, a_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

The framework is identical to P1. Your flow_matching_loss swaps in for mse_loss; everything else stays the same.

The inference loop (also already done)

From visualization.py, paraphrased:

def act(state):
    state_tensor = torch.as_tensor(state)
    actions = policy.schedule.sample(policy.model, state_tensor)
    return actions[0]   # single action

Your schedule.sample drives the policy at inference time.

Your Three Changes, Decoded

Per-line annotations. This is the centerpiece chapter.

def interpolate(self, x1, t):
    eps = torch.randn_like(x1)                # shape (B, action_dim)
    t_expanded = t.view(-1, 1)               # shape (B, 1) for broadcasting
    x_t = t_expanded * x1 + (1.0 − t_expanded) * eps
    velocity = x1 − eps
    return x_t, velocity

@torch.no_grad()
def sample(self, model, state):
    B = state.shape[0]
    x = torch.randn(B, self.action_dim, device=self.device)
    dt = 1.0 / self.num_steps

    for i in range(self.num_steps):
        t = torch.full((B,), i * dt, device=self.device)
        v = model(x, state, t)
        x = x + dt * v

    return x.clamp(0.0, 1.0)

def flow_matching_loss(policy, s_batch, a_batch):
    B = a_batch.shape[0]
    t = torch.rand(B, device=a_batch.device)
    x_t, target_velocity = policy.schedule.interpolate(a_batch, t)
    pred_velocity = policy.model(x_t, s_batch, t)
    return nn.functional.mse_loss(pred_velocity, target_velocity)

Putting it all together

Three blanks, totaling ~15 lines of code. The training loop calls flow_matching_loss per minibatch. Inside, your loss function calls schedule.interpolate. At inference, the policy wrapper calls schedule.sample. Your three pieces interlock to form the full algorithm.

Running It

The command

python main.py --method bc_flow --env hard

This:

1. Collects expert demos on hard mode (same as P1).
2. Builds a FlowMatchingPolicy (the U-Net + your schedule).
3. Trains for some number of epochs using flow_matching_loss.
4. Evaluates on 50 episodes via schedule.sample for action selection.
5. Saves results to bc_flow_hard.txt.

Expected results

You should see a substantial improvement over Problem 1 on hard mode. If episode length stays low, something is off — check the gotchas in Chapter 10.

How long does training take?

Longer than P1's MLP. The U-Net has more parameters, and each training step does an interpolation + forward pass. Expect roughly 5-15 minutes on CPU, much faster on GPU.

What healthy training looks like

The deliverables

Per the PDF:

1. Table of mean ± std episode length on hard mode (50 evaluation episodes).
2. 2-3 sentence explanation of why flow matching performs the way it does on hard mode.

Your one-paragraph writeup template

Cheat Sheet & Self-Quiz

Equations to memorize

API reference

Self-quiz

1. What's the difference between a regressor and a generative model?
2. Why does flow matching solve the multimodality problem from P1?
3. What is a vector field and why is it useful here?
4. What does the network v_θ(x, s, τ) predict, in one sentence?
5. Why is the target velocity a₁ − a₀ constant in τ?
6. Why do we sample τ uniformly during training instead of using a fixed value?
7. What does the Euler step x ← x + (1/n) · v compute?
8. Why do we clamp to [0, 1] only at the end of sampling, not after each step?
9. What's the connection between flow matching and diffusion?
10. Why does FM typically need fewer sampling steps than diffusion?
11. What's the difference between torch.randn and torch.rand?
12. If you set num_steps = 1 in the sampler, what does flow matching reduce to?

Implementation order

1. FlowMatchingSchedule.interpolate — ~2 minutes. Five lines.
2. flow_matching_loss — ~1 minute. Five lines. Verify training loss decreases.
3. FlowMatchingSchedule.sample — ~3 minutes. Eight lines. Hardest because of the loop and device handling.

Total: ~6 minutes of typing. Run on hard mode, see episode length jump from ~300 (BC regression) to ~800+ (flow matching). Write the explanation.

You can now teach this

Three big ideas, in order of importance:

1. Generative models preserve multimodality. A regressor collapses multimodal targets to their mean. A generative model samples from the distribution, preserving each mode. For multimodal experts, this is the difference between crashing and succeeding.
2. Flow matching is MSE on velocity along straight-line paths. The training procedure is: sample data, sample noise, sample time, build linear interpolation, regress network output toward the constant difference (data − noise). The math is breathtakingly simple.
3. Sampling is Euler integration. Start from noise. At each step, query the network for velocity at the current point and time. Step forward by 1/n in that direction. Repeat n times. Clamp. Done.

If a friend asks: "Why does flow matching work where MSE regression failed?" — you say: "MSE regression learns the conditional mean of expert actions, which collapses bimodal data into the wrong middle ground. Flow matching learns a vector field that maps noise to data; different noise samples flow to different modes, so the policy can sample either mode rather than averaging them. Mathematically, both are MSE losses, but flow matching applies MSE to velocity predictions along learned paths, not to raw actions. The richer prediction target preserves multimodality."

You can teach this. On to Problem 3.