Diffusion Policy: sample the action, do not summarise it

Instead of predicting one action, learn to turn noise into an action, and every roll of the dice lands on a different valid answer instead of the average of all of them.

Where you are. You know that a squared-error loss regresses the conditional mean, and that on multimodal demonstrations the mean is an action nobody took. This lesson replaces the head that summarises with one that samples.

Static, resolving

Open any image generator and ask it for a mug on a table. Watch the preview instead of the result.

The first frame is television static. The second is static with a suggestion of a horizon in it. Somewhere around the twentieth there is a mug, with a handle, sitting on a table with a shadow.

Now press go again. You get a different mug. Different handle, different table, different light. What you never get is a grey smear - which is exactly what the average of every mug in the training set would look like.

An action is a much smaller object than an image. Sixty-four timesteps of seven joint angles is 448 numbers; even a small photograph runs to tens of thousands. Nothing in the procedure cares which of the two you hand it.

The idea in one paragraph

A diffusion policy never predicts an action. It learns to remove noise from an action, and you run it as a loop: start with a block of random numbers the same shape as an action chunk, ask the network which noise is sitting on it, subtract a little of that, and repeat about a hundred times. What falls out is one sample drawn from the distribution of actions the demonstrators took in this situation - not their average. Change the random numbers you started from and you get a different valid sample. So when half your demonstrators went left and half went right, this policy goes left or right; it never goes straight into the trolley, because straight is not a place any of the demonstrations lived and the loop has nothing pulling it there.

Two objectives wearing the same architecture

The network in a diffusion policy is not exotic. It can be the same stack of linear layers you used for behaviour cloning, or the same 1D convolutional net over time, or the same transformer. What changed is the question you ask it.

	Regression head	Diffusion head
Input	the observation	the observation, a noisy action chunk, and how noisy it is
Output	an action chunk	a guess at the noise
Called	once per decision	once per denoising step, so ~100 times per decision
Answer	always the same for a given observation	a fresh sample each time

The forward process is arithmetic, not learning

Pick a real action chunk $a^0$ out of the dataset. Define a fixed recipe for spoiling it by a controlled amount:

a^{k} = \sqrt{\bar\alpha_k}\;a^{0} + \sqrt{1 - \bar\alpha_k}\;\varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)

Here $k$ runs from 0 to about 100, $\varepsilon$ is a block of standard normal noise the same shape as the chunk, and $\bar\alpha_k$ is a number that slides from nearly 1 down to nearly 0 as $k$ grows. In words: at $k=0$ you have the human’s action untouched, at $k=100$ you have pure noise, and every step in between is a known blend of the two.

The training loss is mean squared error, and that is not a contradiction

Training is three lines. Draw a real chunk from the dataset. Pick a random click $k$ and a random noise block $\varepsilon$ , and build $a^k$ with the formula above. Then ask the network to name the noise:

\mathcal{L} = \mathbb{E}\left[\left\lVert\, \varepsilon - \varepsilon_\theta\!\left(a^{k},\, k,\, o\right) \right\rVert^{2}\right]

Training a diffusion policy: mix known noise into a real action chunk, ask the network to name that noise, and score it against the noise you drew

Wider than the screen; scroll it sideways.

Stop on that for a second, because the last lesson said mean squared error regresses a mean and that regressing a mean is the whole problem. This objective is mean squared error.

What changed is the label. The target is no longer the action, which had two valid values; it is $\varepsilon$ , the noise you drew yourself a moment ago and therefore know exactly. Averaging over the dataset is now harmless, because the quantity being averaged is a well-defined function of inputs you fully control.

The ambiguity has not been solved so much as moved. It now lives in the random block you start sampling from, and you draw a fresh one every time you ask the policy for an action.

Sampling is a loop with a dash of noise added back

To get an action, start from pure noise $a^{K} \sim \mathcal{N}(0, I)$ and walk $k$ backwards:

a^{k-1} = \frac{1}{\sqrt{\alpha_k}}\left(a^{k} - \frac{\beta_k}{\sqrt{1 - \bar\alpha_k}}\;\varepsilon_\theta\!\left(a^{k}, k, o\right)\right) + \sqrt{\beta_k}\;z

In plain terms: subtract a fraction of the noise the network claims to see, then add back a smaller dash of fresh noise $z$ , and repeat. The final step drops the added noise and returns the chunk.

Sampling an action chunk by denoising: start from random numbers, ask the network which noise is present, subtract some, and repeat, with the observation fed in at every pass

Wider than the screen; scroll it sideways.

Watch it split

The obstacle task from the last lesson, in miniature: one observation - how far ahead the obstacle is - and one action, the sideways-and-forward step to take. Half the demonstrations swerve left, half right, nothing in between. Two models, same data, same 2,000 gradient steps.

Demonstrated actions, the single action a regression head predicts, and eight hundred samples from a diffusion head, at one fixed observation

Real example

Measured at an obstacle distance of 0.5, over 2,000 samples, on a laptop CPU. “Straight” is a dead band around zero sideways motion - the answer no demonstrator gave.

	went left	went straight	went right
the demonstrations	49%	0%	51%
regression head	0%	100%	0%
diffusion head	51%	1%	48%

The regression head is not slightly wrong. It is wrong on every single sample, and its mean sideways motion, $-0.004$ , is an almost perfect estimate of the mean of a distribution that has no mass there.

Snapshot the sample cloud partway through the loop and you can see the decision being made. Early on it is a fuzzy ball centred on nothing in particular. By the last twenty steps it has pulled apart into two lumps, and each individual sample has already committed.

Six hundred samples partway through denoising, starting as one cloud of noise and separating into the two demonstrated modes

What Diffusion Policy adds on top of “diffusion”

Four design choices turn “diffusion” into a policy you can run on a robot. The first two are the paper’s argument; the last two are what the reference implementations actually ship.

Diffuse a sequence, not a single action. The thing being denoised is a chunk of $H$ future actions, so the same machinery that fixes multimodality also gives you action chunking for free.
Condition on the observation rather than diffusing it. The image encoder runs once per decision; only the small action-denoising network runs a hundred times. Diffusing images and actions jointly would be a hundred forward passes of the image encoder per control step.
Receding horizon. Predict a long chunk, execute a shorter prefix of it, then re-plan from a fresh observation.
A separate visual encoder per camera, with a spatial-softmax layer that reduces a feature map to a handful of keypoint coordinates rather than flattening it.

The bill: one decision costs a hundred forward passes

That multiplier is the entire reason this part of the field is moving. But it is a setting, not a law, and the trade is much better than the default suggests:

Real example

The same trained model, sampled with fewer steps, on the same laptop.

inference steps	went left	went straight	went right	microseconds per action
100	50%	1%	49%	80
50	52%	1%	47%	41
20	52%	1%	47%	18
10	51%	2%	47%	8
5	50%	4%	46%	4

Ten steps is ten times faster than a hundred and still puts 98% of its samples on a demonstrated mode. Five starts to leak into the gap. On this task the deployable setting is around ten, and finding your own number is worth doing once per task - the default of a hundred is chosen to be safe, not to be fast. The microsecond column is the one that will not reproduce exactly on your machine; the percentages will.

Field state

LeRobot v0.6.1’s diffusion policy defaults: n_obs_steps=2, horizon=64, n_action_steps=32, DDPM with num_train_timesteps=100, beta_schedule="squaredcos_cap_v2", prediction_type="epsilon", a ResNet-18 vision encoder with 32 spatial-softmax keypoints, optimizer_lr=1e-4.

Two traps in that list. num_inference_steps=None defaults to the full 100 denoising steps at inference - that is the latency number, and it is a knob, not a law. And horizon and n_action_steps were 16 and 8 in v0.5.1, four times shorter in both. Numbers from before and after that boundary are not comparable, and neither are numbers from before the v0.6.0 fix to loss normalisation on padded actions. Pin the version in anything you publish.

Flow matching has largely replaced this exact sampler in new work, for the latency reason and no other. The 2026 literature says so directly: generative heads are dominant for visuomotor policies, and iterative denoising is the part that hurts. The argument in this lesson - sample the distribution, do not summarise it - is what survived, and it now sits inside the large vision-language-action models. The specific hundred-step sampler is the part being replaced.

Check yourself

1. The training loss is mean squared error, which the previous lesson called the source of the problem. Why is it safe here?

Because the label changed. Regressing the action is dangerous when several actions are valid, since the minimiser is their mean. Here the label is $\varepsilon$ , the noise you drew yourself before corrupting the chunk, and given the noisy chunk, the step index and the observation there is a single well-defined thing to predict. The ambiguity has moved into the random block you start sampling from, which you redraw for every decision.

2. Two robots are given exactly the same observation. Why do they not necessarily do the same thing, and is that a bug?

They start their sampling loops from different random noise, so they can land in different valleys. It is not a bug; it is the point. Any policy that is guaranteed to output the same action for the same observation is a function of the observation alone, and can therefore only ever report a single summary of what the demonstrators did. Determinism is precisely what you gave up in order to stop averaging.

3. Where does the compute go at inference, and what is the one knob that changes it?

Into repeating the same forward pass. One action chunk costs num_inference_steps passes of the denoising network - 100 by default in LeRobot v0.6.1 as of August 2026. The knob is that step count. Lowering it is a direct trade of sample quality for latency, and it is the first thing to try when a diffusion policy misses its control deadline. The image encoder is not the problem: it runs once per decision, because the policy conditions on the observation rather than diffusing it.

4. Why does Diffusion Policy denoise a whole chunk of future actions rather than one action at a time?

Three reasons stack. It reuses the mechanism you already need for action chunking, so one design buys both. It keeps a sampled chunk internally consistent - the whole plan comes out of one draw, so it commits to one route rather than re-flipping the coin every timestep. And it amortises the hundred denoising passes over $H$ executed actions instead of paying them per action.

5. A colleague reports 65% success for Diffusion Policy on a benchmark and compares it against a number from a paper published eight months earlier. Name two reasons the comparison may be invalid.

The configuration defaults moved: LeRobot’s horizon and n_action_steps quadrupled between v0.5.1 and v0.6.x, so “the default diffusion policy” is a different policy on either side of that line. And a loss-normalisation bug affecting padded actions was fixed in v0.6.0, so training numbers straddling that release are not measuring the same objective. Either alone is enough to break the comparison. Publish the version and the config, not just the percentage.

6. In the sampling loop you subtract the predicted noise and then add a smaller amount of fresh noise back. Why add any back at all?

Because a purely downhill walk from a given starting point is deterministic, and would slide into whichever valley the starting point happened to sit nearest, with the transition points between valleys decided entirely by the initial draw. The added noise keeps the walk exploratory in the early steps, when the sample has not yet committed, and it shrinks as $k$ falls so the last steps settle cleanly. The final step drops it entirely.

Do this

Two runs and one experiment, around twenty minutes on CPU.

cd module-03-robot-learning/solutions
python diffusion_toy.py

It trains both heads on the obstacle task and prints the left / straight / right table above. Then open code/diffusion_toy.py, which leaves the two halves of diffusion for you to write:

The forward process. Fill in q_sample(a0, k, eps) so it returns $\sqrt{\bar\alpha_k}a^0 + \sqrt{1-\bar\alpha_k}\varepsilon$ . Sanity-check it: at $k=0$ the output should equal the input, and at $k=99$ its standard deviation should be close to 1 regardless of the action.
The sampling loop. Fill in sample(). If you get a cloud that never separates, print the mean and standard deviation of the samples at each $k$ and find the step where it stops shrinking.
Cut the step count. python diffusion_toy.py --sweep runs the sampler at 100, 50, 20, 10 and 5 inference steps and reports the split and the time per action at each. Find the smallest number of steps at which the two modes are still clean. That is the number you would deploy with, and it is task-specific. Note that skipping steps needs the deterministic update rather than the ancestral one; the solution shows both.
Break it on purpose. Change the coin flip in actions_for so every demonstration swerves the same way, and retrain. Both heads should now land on the demonstrated action - the diffusion head has no advantage when there is nothing to disambiguate, and it still costs a hundred times more to run. Knowing when not to reach for it is worth as much as knowing how it works.

What you can now do

You can write down the forward diffusion process and explain that it involves no learning at all. You can state the training objective, and defend the apparent contradiction of using mean squared error to escape a mean-squared-error problem by pointing at the label. You can trace the sampling loop and say which line makes two demonstrators’ routes stay separate. And you can price it: one decision costs as many network passes as you set inference steps, which is why the field kept the argument and is replacing the sampler.