Journal

Embodied AI · 23 Jun 2026 · 7 min read

Diffusion policies, and what breaks when you ship one

Denoising an action sequence is an unreasonably good way to model multimodal behaviour. It is also a sampling loop inside a control loop, which is where the engineering starts.

DiffusionRoboticsInference

The reason diffusion policies work is that they refuse to average incompatible demonstrations. The reason they are hard to deploy is that every decision now costs N forward passes instead of one.

Give a regression policy two equally valid ways to go around an obstacle and it will learn to go straight through it. This is not a training bug; it is the L2 loss doing exactly what it was asked. The conditional mean of a bimodal distribution sits in the valley between the modes, and the valley is where the obstacle is.

Everything interesting about diffusion policies follows from the fact that they do not have this problem.

The formulation#

Instead of predicting an action, you learn to denoise a sequence of actions. Training corrupts a ground-truth chunk with Gaussian noise at a random level and asks the network to predict the noise, conditioned on the observation. At inference you start from pure noise and iterate the learned denoiser down to a clean trajectory.

code
def training_step(model, obs, actions, scheduler):
    """actions: (B, H, D) — a chunk, not a single action."""
    B = actions.shape[0]
    t = torch.randint(0, scheduler.T, (B,), device=actions.device)
    noise = torch.randn_like(actions)
    noisy = scheduler.add_noise(actions, noise, t)

    pred = model(noisy, t, obs)             # predicts the noise
    return F.mse_loss(pred, noise)


@torch.no_grad()
def sample(model, obs, scheduler, shape, steps=10):
    x = torch.randn(shape, device=obs.device)
    for t in scheduler.timesteps(steps):    # e.g. DDIM: 10 of 100
        x = scheduler.step(model(x, t, obs), t, x)
    return x                                # (B, H, D) clean action chunk

The conditioning is the design decision that gets underweighted. Two options dominate: FiLM-style conditioning, where the observation modulates the denoiser's feature maps, and cross-attention, where the denoiser attends to a sequence of observation tokens. Cross-attention is more expressive and noticeably more expensive; FiLM is usually sufficient when the observation is a fixed-size embedding. If you are conditioning on a language instruction that needs to be attended to token-by-token, you want cross-attention and should budget for it.

Diffusion policies model a distribution over trajectories, not a point estimate. That is the entire advantage, and it is only an advantage if your data actually contains multiple valid strategies.

If your demonstrations are all one operator doing the task the same way every time, a diffusion head buys you almost nothing over regression and costs you ten times the inference. We have seen teams reach for it out of fashion and then wonder where their latency went.

Receding horizon, and why you throw most of it away#

The standard deployment pattern predicts a chunk of H actions and executes only the first E, then re-plans. Typical values are H = 16, E = 8 — you discard half of every generated trajectory.

That waste is deliberate. The later actions in a chunk are conditioned on an increasingly stale observation, and they are the ones the denoiser is least confident about. Discarding them costs compute and buys reactivity. Where you set E is the same tradeoff as the chunk-length question in action chunking, with one extra wrinkle: because each re-plan is a fresh sample from a distribution, consecutive plans can pick different modes. The robot goes left, re-plans, decides right, and oscillates in front of the obstacle.

Three fixes, roughly in order of how much we like them:

Warm-start the sample

Initialize the reverse process from the previous plan's tail plus a small amount of noise, rather than from pure noise. The sampler stays in the mode it was already committed to unless the observation strongly disagrees. Cheap and very effective.

Longer execution horizon

Commit further before re-planning. Reduces oscillation by reducing the number of opportunities to change your mind. Costs reactivity.

Explicit mode tracking

Sample K trajectories, cluster them, and prefer the cluster nearest the previous commitment. Principled, and K times the compute. Reserve it for cases where the first two are not enough.

Making the sampling loop affordable#

A 100-step DDPM sampler inside a 10 Hz control loop is a non-starter. The options, from least to most invasive:

| Technique | Steps | Quality | Notes | | --- | --- | --- | --- | | DDPM (baseline) | 50–100 | Reference | Only useful offline | | DDIM | 8–16 | Very close | Deterministic; first thing to try | | DPM-Solver++ | 4–10 | Close | Higher-order solver, drop-in | | Consistency distillation | 1–2 | Slight loss | Needs a distillation stage | | Flow matching | 2–8 | Comparable | Different training objective entirely |

In practice DDIM at 8–10 steps is where most teams land, because it requires changing three lines and no retraining. If that is still too slow, distillation or a flow-matching reformulation are the real answers — and flow matching is increasingly the default for new work, since it reaches diffusion-quality multimodality in a handful of integration steps by construction rather than by distillation.

The other lever is architectural. A U-Net over the action sequence is the classic choice and is heavier than it needs to be for short horizons; a small transformer over H action tokens is usually faster at equal quality when H ≤ 32, and it composes better with cross-attention conditioning.

Batch the denoising steps, not the samples. If you are drawing K candidate trajectories, run all K through each denoising step together — one kernel launch per step instead of K. On an embedded GPU the launch overhead is a disproportionate share of the total.

Normalization is not a detail#

Diffusion assumes its target lives roughly in a unit-scale space. Robot actions do not: a translation in metres, a rotation in radians, and a binary gripper command have wildly different natural scales, and the noise schedule treats them identically.

Normalize per dimension, using percentiles rather than min/max, computed on the training set and frozen. Then treat the gripper separately — it is a categorical variable wearing a continuous disguise, and denoising it produces values like 0.43 that mean nothing. Either predict it with a separate small classification head, or accept that you will threshold it and make sure the threshold is applied consistently between training-time evaluation and deployment.

Rotation deserves the same care. Denoising Euler angles will eventually produce gimbal-lock artefacts, and denoising raw quaternions produces unnormalized garbage. The 6D continuous rotation representation, re-orthogonalized after sampling, is the pragmatic default.

What actually goes wrong in deployment#

Silent distribution drift. The policy samples from a distribution fitted to your demonstrations. When the scene drifts — a new tablecloth, different ambient light — the samples get worse in a way that produces no error, no exception, and no obvious signal. The best cheap detector we know is sample variance: draw a handful of trajectories and measure their spread. High spread in a situation that should be unambiguous means the model is out of its depth. It is not calibrated uncertainty in any formal sense, but it correlates well enough to trigger a slowdown or a handoff.

Nondeterminism in evaluation. Two runs of the same policy on the same scene give different trajectories. This is correct behaviour and it makes A/B testing statistically awkward — you need many more rollouts to detect a difference than you would with a deterministic policy. Fix the sampler seed for evaluation runs, report both seeded and unseeded numbers, and do not let anyone compare a seeded run against an unseeded one.

The gripper problem, again. More deployment bugs trace back to gripper command handling than to anything in the denoiser. It is worth writing down, once, exactly how the gripper value flows from demonstration to training target to sample to hardware command, and checking each hop.

When we reach for one#

Diffusion (or flow matching, increasingly) is the right call when the task has genuinely multiple valid solutions, when your data comes from multiple operators with different styles, or when contact dynamics make the correct action sharply discontinuous in the observation. Those are exactly the cases where a regression head produces a policy that hovers, hesitates, and splits the difference.

It is the wrong call when your control loop cannot absorb an N-step sampler, when your demonstrations are unimodal, or when the team does not yet have an evaluation protocol robust enough to detect the difference. In that last case the honest advice is to build the evaluation harness first, ship a regression baseline, and let the measured failure modes tell you whether multimodality is your actual problem.

Let's build

Building something in this space?

If this is the kind of problem your team is working on, we'd like to hear about it — especially the parts that aren't working yet.