Training a diffusion policy, and a fair fight against ACT

Train three heads on one dataset at one budget, roll them out, and discover that the policy with the worse validation loss is the one that reaches the goal.

Where you are. You know why regressing a mean fails on multimodal demonstrations and how a diffusion head fixes it. Now you make the two compete on the same data, under the same budget, and report the result honestly.

The marking scheme

Twenty scripts from one exam, one invigilated room. Everybody sat the same paper for the same two hours with the same books on the desk, which is the easy half of fairness and the half people remember to do.

The marking scheme awards a point for each line that agrees with the model answer. It is quick, it is consistent, and two markers using it reach the same total.

One candidate proves the result a different way. Shorter, correct, and sharing almost nothing with the model answer line by line. The scheme gives it four out of ten. Another reproduces the model answer almost exactly and then stops one step before the conclusion, which the scheme barely notices, because nine of its ten lines matched.

Now mark the same twenty scripts a second way: does the proof work, yes or no. The order comes out different, and the script the scheme buried is at the top, where it was the whole time.

The room was fair and the clock was fair. The scheme was the part that was measuring something else.

The idea in one paragraph

A policy comparison is only worth reading if exactly one thing differs between the arms, and if the number you report is the thing you actually care about. Matching the budget is the first half: same dataset, same split, same gradient steps, same batch size, same evaluation start states, so the only variable left is the head. The second half is harder and catches more people: held-out action error is not a proxy for success, and it is structurally biased against any policy that samples. A generative head that picks one of two valid routes is scored against the one route that happens to be in the validation file, and is punished for the half of the time it picks the other. Rank by rollouts, or you will rank backwards.

One demonstration set feeding three heads at an identical budget, each rolled out on the same start states and reported as a rate with an interval

Wider than the screen; scroll it sideways.

What has to be held fixed

	Why, or what to do instead
Dataset and split	Obvious, and still the most common leak. Build the split once and reuse the indices.
Gradient steps and batch size	The budget. A head that saw twice the updates is not a better head.
Chunk length	Chunking changes success on its own. Vary it in a separate experiment.
Evaluation start states and seeds	Every head faces the same two hundred episodes, not two hundred fresh ones.
Number of trials	Fixed in advance. Never “run more until it looks decided”.
Parameter count - cannot match	Get within a factor of two and say so. Forcing them equal distorts the architectures.
Inference cost - cannot match	Report it as a result. It is one of the things you are choosing between.
Tuning effort - cannot match	The one nobody reports. Say how much search each arm got.

The run

The task is the corridor: drive a point past an obstacle to a goal, forty steps, chunks of eight executed open-loop. Fifty demonstrations, half of which go left round the obstacle and half right. Three heads, 2,500 gradient steps each, batch 256, all on a laptop CPU.

mlp regresses the action chunk under mean squared error.
cvae is the latent trick shrunk to fit, the one that sits inside ACT: a small latent, inferred during training from the chunk the demonstrator actually produced, then fixed at zero when the policy runs. A second row samples the latent from its prior instead, so you can see what it is worth.
diffusion denoises the chunk out of noise, 100 steps.

head	parameters	training	held-out action error	reached the goal	95% interval	hit the obstacle
mlp	136,464	6.9 s	0.00006	96 / 200	41% - 55%	104
cvae, latent set to zero	213,280	11.7 s	0.00006	65 / 200	26% - 39%	135
cvae, latent sampled	213,280	11.7 s	0.00007	79 / 200	33% - 46%	121
diffusion, 100 steps	141,072	7.5 s	0.00017	151 / 200	69% - 81%	36

Three heads, one dataset, 2,500 gradient steps each

n = 200 trials · corridor past an obstacle, chunks of 8 · 2026-08-09

Sixty rollouts from each head: green paths reached the goal, red ones hit the obstacle, grey ones ran out of steps

Read the intervals before the point estimates. The diffusion head’s 69-81% and the regression head’s 41-55% do not overlap, so that difference is real at this trial count. The two CVAE rows overlap each other, so this run does not establish that sampling the latent beats zeroing it - it only suggests it. Two hundred trials is enough to separate 48% from 76% and not enough to separate 33% from 40%.

Two costs are worth pulling out. Training is cheap and roughly comparable: 7 to 12 seconds for 2,500 steps on a four-core laptop CPU from 2018, and note that the diffusion head is not the expensive one to train - the CVAE is, because it runs two networks per step. Treat those seconds as soft - running the same script while the machine was busy with other work stretched them to between 65 and 280 seconds, which is a useful reminder about every wall-clock number you will read in a paper. Inference is where the methods genuinely differ, and it is structural: with two hundred episodes stepped together, the diffusion head took 0.31 milliseconds per chunk against 0.01 for the others. That factor of roughly thirty is the hundred network passes, amortised by batching over episodes, and no faster laptop removes it.

The regression head is also not failing the way the last lesson predicted. It does not drive neatly into the obstacle every time; it hits 104 times out of 200 and muddles through on the rest. The reason is worth seeing: heading up the middle takes it into the corridor between the two demonstrated routes, which is a region no demonstration ever visited, so what it does next has no supervision behind it at all. Mode-averaging pushes it off-distribution and covariate shift decides the rest. The two failures compound, and that is what the path plot shows: a red bundle heading straight up into the obstacle, and green paths sweeping round both sides. The diffusion head has its own small failure mode visible there too - a handful of grey paths that swing too wide and run out of steps.

The result was partly about my tuning, not the method

The table above uses a latent-divergence weight of 10, because that is the number LeRobot’s ACT configuration ships. It was the wrong number here - the two losses are scaled differently, so the value does not carry across - and I had no reason to expect that until I checked. Sweeping it, at the same budget and the same 200 start states:

divergence weight	latent at inference	held-out action error	reached the goal	95% interval
0.1	zero	0.00022	0 / 200	0% - 2%
0.1	sampled	0.00042	33 / 200	12% - 22%
1.0	zero	0.00007	164 / 200	76% - 87%
1.0	sampled	0.00012	166 / 200	77% - 88%
10.0	zero	0.00006	65 / 200	26% - 39%
10.0	sampled	0.00006	73 / 200	30% - 43%

The latent head is extremely sensitive to one number

n = 200 trials · corridor, CVAE head only · 2026-08-09

Both ends fail for opposite reasons. At 0.1 the divergence term is too weak, the decoder leans on a latent that ranges freely, and feeding it a zero at inference asks for a chunk it was never trained to produce - every single episode hit the obstacle. At 10 the term is strong enough that the latent carries almost nothing and the head degenerates towards the mean. At 1.0 it works, and its interval, 76-87%, overlaps the diffusion head’s 69-81%.

Note also that at a divergence weight of 1.0 the two latent rows are barely distinguishable: zeroing the latent scores 164 and sampling it scores 166. On this task the decoder found a route and stuck to it, rather than using the latent to encode which side to take. Whether that holds on a real task with real demonstrations is exactly the kind of thing you check rather than assume.

Why the loss ranked them backwards

Held-out action error next to success rate for the same three heads, ranked in opposite orders

Wider than the screen; scroll it sideways.

Held-out action error asks: given this observation from a held-out episode, how close is the policy’s chunk to the chunk that episode contains? For a policy that samples, the answer is bad by construction. Half the time it picks the mode the validation episode did not use, and gets charged the full distance between two perfectly valid routes.

So a policy sampling correctly from a two-mode distribution scores worse on this metric than a policy that outputs the midpoint of the two - which is the action that fails.

This is the last honest thing a loss curve can tell you, and it is why the next lesson is entirely about evaluation.

The bug that cost me an afternoon

My first version of this comparison put the diffusion head at 8 out of 200, dead last, with a held-out error eight times the regression head’s. The obvious reading was that diffusion does not work at this scale.

The cause was not the algorithm. Actions in this task are steps of about 0.045 units, so the chunk had a standard deviation of 0.029. The forward process mixes in noise of unit scale whatever your data looks like. Work out the ratio of surviving signal to added noise at the last step of the schedule: $\sqrt{\bar\alpha_{99}}\,(0.029) \big/ \sqrt{1-\bar\alpha_{99}}$ , which with this schedule’s $\bar\alpha_{99} = 0.364$ is 0.022. The network was being shown near-pure static and asked to name it, across most of the schedule.

Real example

The same three heads, trained identically, with and without normalising the action chunks to zero mean and unit spread:

head	without normalising	with
mlp	23 / 200	96 / 200
diffusion, 100 steps	8 / 200	151 / 200

Not one line of either model changed. The regression head improves fourfold because tiny targets make tiny gradients and it was simply under-trained. The diffusion head improves nineteenfold, because for it the scale mismatch is not a slow-training problem but a destroyed signal.

This is why LeRobot ships a meta/stats.json in every dataset and normalises with it. It is not hygiene; it is the difference between a working policy and a broken one, and it fails silently. You get a trained model, a smooth loss curve, and a robot that does nothing sensible.

Guard against it mechanically rather than by remembering. After building your normaliser, assert that the normalised training data has unit standard deviation, and floor the per-dimension standard deviation so a channel that never varies cannot divide by zero.

The control

Retrain all three on demonstrations that always go the same way round. The distribution now has one lump, and the whole argument should evaporate.

python compare_policies.py --control

head	held-out action error	reached the goal	95% interval
mlp	0.00004	200 / 200	98% - 100%
cvae, latent set to zero	0.00004	200 / 200	98% - 100%
cvae, latent sampled	0.00004	200 / 200	98% - 100%
diffusion, 100 steps	0.00007	193 / 200	93% - 98%

The same three heads when the demonstrations agree

n = 200 trials · corridor, one-sided demonstrations · 2026-08-09

Everything reverses. The regression head goes from 48% to perfect, nobody hits the obstacle, and the diffusion head is now very slightly worse - it loses seven episodes to the randomness that was the whole point of using it. That is the honest cost: a sampling policy pays a small tax for being a sampling policy, and it is only worth paying when there is something to disambiguate.

If your control does not do this, the difference you measured was not about multimodality and you have another bug. A comparison without its control is an anecdote with a table.

Doing this for real

Field state

On real data with real policies, the same comparison is two lerobot-train invocations. The bones of it:

lerobot-train --dataset.repo_id=$HF_USER/my_task --policy.type=act \
  --steps=20000 --batch_size=8 --seed=1000 \
  --policy.scheduler_decay_steps=20000 \
  --output_dir=outputs/act --job_name=act --policy.device=cuda

lerobot-train --dataset.repo_id=$HF_USER/my_task --policy.type=diffusion \
  --steps=20000 --batch_size=8 --seed=1000 \
  --policy.scheduler_decay_steps=20000 \
  --output_dir=outputs/dp --job_name=dp --policy.device=cuda

Four things in there are easy to get wrong. The default is --steps=100000, and LeRobot’s own hardware guide says imitation learning usually converges in 5 to 10 epochs over the dataset, so compute total_frames / batch_size and pick your own number. If you shorten training you must shorten --policy.scheduler_decay_steps to match, or the learning rate never finishes decaying. pip install lerobot stopped including training dependencies in v0.6.0, so you want lerobot[training]. And LeRobot needs Python 3.12 or newer with PyTorch 2.7 or newer, which rules out Intel Macs entirely: the last macOS x86-64 PyTorch wheel is 2.2.2.

Field state

Indicative wall-clock from LeRobot’s hardware guide for 5 epochs on about 50 episodes at 640x480 - their own text calls these numbers approximate to within about 50%:

Machine	ACT	Diffusion
RTX 4090 / 3090	30-60 min	2-4 h
L4 or A10G	1-2 h	not listed
Apple silicon, MPS, batch 4	6-14 h	not listed
CPU only	“Don’t train. Use Colab or rent a GPU.”	same

Where both appear, diffusion costs roughly four times ACT. The guide gives no diffusion figure for Apple silicon and I would not extrapolate one: time a hundred steps on your own machine and multiply. Peak memory is about 2-6 GB for ACT and 8-14 GB for diffusion, so a 24 GB card covers both comfortably. Rental at the time of writing: RunPod community RTX 4090 at $0.34/hour; Hugging Face Jobs a10g-small at $1.00/hour, billed by the minute, integrated into lerobot-train through --job.target. Set --job.timeout explicitly rather than trusting a default, because the two layers disagree: Hugging Face Jobs’ own default is 30 minutes, which kills a longer run silently, and LeRobot’s wrapper overrides it to 2 days, which bills a wedged run for 48 hours.

A full ACT-versus-diffusion comparison with three seeds each is a handful of dollars on a community 4090 and roughly three times that on managed jobs. The seeds are the part worth paying for.

Check yourself

1. Your diffusion policy has three times the held-out action error of your regression policy. What have you learned about which is better?

Almost nothing. Held-out error scores a sampled action against whichever mode the validation episode happened to contain, so a policy sampling correctly from a two-mode distribution is charged the distance between the modes about half the time. The metric structurally prefers the mean, which is the action that fails. The only way to rank them is rollouts.

2. After the divergence-weight sweep, what is the strongest claim these tables actually support, and what did the sweep cost the original headline?

Only this: both generative heads beat plain regression on multimodal demonstrations. The tuned latent head’s interval, 76-87%, overlaps the diffusion head’s 69-81%, so the experiment does not separate the two generative heads. It does separate either of them from the regression head, whose 41-55% overlaps neither. Non-overlapping intervals are a conservative check rather than a proper test, but they are enough to tell you which claims to drop. The sweep also exposes the hidden variable: an arm you tuned and an arm you did not are not competing on equal terms, whatever the table says about steps and batch size.

3. A colleague shortens training from 100,000 steps to 20,000 by passing --steps=20000 and nothing else. What breaks?

The learning-rate schedule still decays over 100,000 steps, so the run finishes while the learning rate is still near its starting value and the policy never anneals into a good minimum. --policy.scheduler_decay_steps has to move with --steps. This looks like “the method does not work at 20k steps” and is really “I changed one of two coupled settings”.

4. Normalising the action chunks moved the diffusion head from 8/200 to 151/200 and the regression head from 23/200 to 96/200. Why is the effect so much larger for diffusion?

Because the forward process adds noise of unit scale regardless of your data. If the actions have a standard deviation of 0.03, then across most of the schedule the surviving signal is a few percent of the noise, and the network is asked to name the noise in a sample that is essentially all noise. The regression head has no such fixed external scale: small targets only mean small gradients, so it under-trains and degrades gracefully rather than being destroyed. That is why normalising helps both and rescues only one.

5. Why is the unimodal control run worth the compute, when you already have a result you believe?

Because the result on its own is consistent with several explanations: a broken regression baseline, an unlucky seed, a task the diffusion head happens to suit. The control changes exactly one property of the data - it removes the second mode - and predicts a specific outcome, namely that the gap collapses. If the gap survives the control, your explanation is wrong and the table is measuring something you have not identified.

6. You have $30 of GPU rental and want the strongest claim you can buy. Spend it.

Not on one longer run. Three seeds per policy at a short honest budget beats one long run of each, because seed-to-seed variation is a separate source of spread from trial-to-trial variation and conflating the two is the most common error in this literature. On a community RTX 4090 at $0.34 an hour, three ACT runs and three diffusion runs at an hour or two each is well inside the budget, and leaves room for the ten-versus-fifty-demonstration data-scaling curve, which is usually the more interesting result anyway.

Do this

About two minutes of compute for the two runs, all on CPU, plus however long you spend on the three fill-ins.

cd module-03-robot-learning/solutions
python compare_policies.py
python compare_policies.py --control

Then take the starter in code/compare_policies.py, which leaves three pieces for you:

The normaliser. Per-dimension mean and standard deviation, with a floor so a constant channel cannot divide by zero. Assert that the normalised training chunks have unit spread before you train anything.
The CVAE’s latent. The reparameterised sample and the divergence term. Then sweep kl_weight over 0.1, 1 and 10 yourself and reproduce the middle table. Predict the shape of the curve before you run it: both ends should fail, for opposite reasons.
The confidence interval. Fill in wilson(). Check it against a case you can reason about: 8 successes out of 10 should come out near 49% to 94%. Sit with how wide that is, because it is the trial count most robot demos are reported at.

Then two experiments of your own:

Sweep the inference steps. Run the diffusion head at 100, 50, 20 and 10 denoising steps and record success rate against milliseconds per chunk. Find the cheapest setting whose interval still overlaps the 100-step result. That is the setting you would deploy.
Sweep the chunk length. Try 1, 4, 8 and 16, holding everything else fixed. A chunk of 1 should hurt the diffusion head badly, and the reason is the whole argument for chunking: with a fresh sample every timestep, the policy re-flips the coin constantly and never commits to a route.

What you can now do

You can set up a policy comparison where exactly one thing differs, name which budget you matched and which you did not, and report a success rate with the interval attached rather than a bare percentage. You can explain why held-out action error ranks a sampling policy below a mean-regressing one, and refuse to use it for model selection. You can spot the two ways a comparison lies to you even when the budget is matched: an unreported normaliser, and unequal tuning effort between the arms. And you can price the real version of this experiment on rented hardware without discovering either of the two job-timeout defaults the hard way.