Section 41.1: Diffusion models as planners

A policy that cannot represent uncertainty about where to put the gripper will put it in the average place, which is often wrong in a bimodal world.

A Multimodal Planner

Show the same mug to a skilled robot ten times and a good policy should reach for it ten slightly different ways, yet the single most common training objective quietly forbids exactly that. Action trajectories are multimodal. Faced with a mug on a table, a competent agent has several valid grasps: top handle, side rim, two-finger pinch on the body. A regression policy trained to minimize mean squared error against demonstrations averages those modes, and the average of a left grasp and a right grasp is a collision with the mug. In benchmark bimanual tasks (Chi et al., 2023), an MSE-regression policy typically achieves around 40% success when two grasp modes are equally represented in the data, while a diffusion policy on the same demonstrations typically reaches 85%; the exact gap varies with task and data balance, but the direction is consistent because the diffusion policy commits to one mode instead of averaging both into an infeasible middle. This is the core reason diffusion models entered embodied AI: their iterative denoising naturally represents a multimodal distribution over trajectories rather than a single conditional mean. Figure 41.1A shows this concretely: a score network denoises a trajectory from Gaussian noise conditioned on the scene, and the final plan commits to one of the two valid grasp modes rather than averaging them into a collision.

The practical question is not whether the model looks impressive. The question is which action becomes easier, safer, more data-efficient, or more recoverable when the method is inserted into the loop. For multimodal action spaces, the answer is that the policy stops committing to the average place and starts committing to one coherent mode.

A planner that always predicts the average of its training demonstrations is not a planner; it is a weighted average of human mistakes, confident and wrong in exactly the place it matters most.

Theory

Iterative denoising achieves that commitment to a single coherent mode through two coupled processes: a fixed forward process that corrupts trajectories and a learned reverse process that reconstructs them. The forward and reverse definitions below make the mechanism precise.

A diffusion planner treats a trajectory \(\tau_0\) (a sequence of actions, states, or both) as a sample from a data distribution it must learn to generate. Training defines a fixed forward process that gradually corrupts a clean trajectory into Gaussian noise, and learning fits a reverse process that walks noise back to a plausible trajectory conditioned on context \(c\) (the observation, the goal, or a return target).

Forward corruption and the reverse chain

The forward process adds noise according to a variance schedule \(\beta_1,\dots,\beta_T\). Writing \(\alpha_t = 1-\beta_t\) and \(\bar\alpha_t = \prod_{s\le t}\alpha_s\), the closed form for the noised trajectory at step \(t\) is

$$ q(\tau_t \mid \tau_0) = \mathcal{N}\!\left(\tau_t;\ \sqrt{\bar\alpha_t}\,\tau_0,\ (1-\bar\alpha_t) I\right). $$

As \(t\) grows, \(\bar\alpha_t \to 0\) and the trajectory dissolves into standard normal noise. The reverse process is also Gaussian and is the object we train:

$$ p_\theta(\tau_{t-1} \mid \tau_t, c) = \mathcal{N}\!\left(\tau_{t-1};\ \mu_\theta(\tau_t, t, c),\ \sigma_t^2 I\right). $$

The network learns the mean \(\mu_\theta\) (equivalently, the noise that was added), and sampling chains these reverse steps from \(t=T\) down to \(t=0\). Because each reverse step is stochastic and conditioned on \(c\), repeated sampling from the same observation yields different valid trajectories: this is committing to one coherent mode, exactly the multimodal behavior a regression policy cannot produce. Figure 41.1B lays out this two-way chain end to end, from the clean trajectory through the forward corruption to pure noise, and back along the trained reverse path.

This chain is also where the guidance mechanism used later in this section takes effect: classifier-free guidance (introduced in the algorithm box below, and used throughout section 41.2's Diffuser and Decision Diffuser architectures) works by nudging each reverse step, so the forward/reverse machinery just defined is the substrate that conditioning and guidance operate on, not a separate mechanism layered on top.

The variance schedule matters for embodied AI because robot action spaces mix signals of very different magnitudes: a wrist joint moves in milliradians while a base translates in meters. A schedule that decays \(\bar\alpha_t\) too fast corrupts small-magnitude dimensions into pure noise before the network learns to recover them. Those joints then drift systematically during rollout. To see why this is not a minor rounding error, consider the scales. A wrist joint moving 0.01 radians and a base translating 1 meter sit three orders of magnitude apart on the same input vector. At \(t = T/2\) the wrist signal is already buried in noise while the base signal is still clearly visible, so the denoiser trains almost entirely on the base channel. A schedule that decays too slowly leaves the final noised sample still correlated with the original trajectory. The reverse chain then cannot explore alternative grasp modes and collapses to the dominant training mode. On a physical robot, both failures produce incorrect end-effector poses that no downstream control can compensate.

So which schedule is right for a robot arm whose wrist joint operates in milliradians while its base moves in meters? There is no universal answer, and that tension is precisely what makes schedule design a persistent source of silent failure in deployed systems.

Mechanically, the schedule fixes a sequence of signal-retention values \(\bar\alpha_t \in (0,1)\) that decreases monotonically from near one to near zero. Each training step samples a random \(t\), scales the clean trajectory by \(\sqrt{\bar\alpha_t}\), scales a zero-mean unit-variance noise draw by \(\sqrt{1-\bar\alpha_t}\), and sums the two. The network sees the mixture and must predict the noise component. This reparameterization lets us sample any noise level in one closed-form step, so training never simulates the full chain: each gradient update needs only the noise prediction at a single level.

Worked Example

The probe below makes the forward and reverse processes concrete on a tiny 2D trajectory that should reach the goal at \((1, 0)\). We apply five steps of forward noising under a schedule \(\bar\alpha_t\), then run five reverse steps with a deliberately simple linear denoiser that estimates \(\tau_0\) and takes a DDIM-style update. The diagnostic to watch is the endpoint distance to the goal: it should grow during noising and shrink back toward zero during denoising.

# Forward noising then reverse denoising of a 2D action trajectory.
# Forward:  q(tau_t | tau_0) = N(sqrt(abar_t) tau_0, (1 - abar_t) I)
# Reverse:  estimate tau_0 from tau_t, then take a deterministic step toward tau_{t-1}.
import numpy as np

H = 6                                   # waypoints in the trajectory
goal = np.array([1.0, 0.0])
tau0 = np.stack([np.linspace(0.0, 1.0, H), np.zeros(H)], axis=1)   # clean plan

T = 5
abar = np.array([0.85, 0.65, 0.45, 0.25, 0.08])   # signal retained at each step
rng = np.random.default_rng(7)
noise = rng.normal(size=tau0.shape)               # one fixed noise draw

def endpoint_dist(tau):
    return float(np.linalg.norm(tau[-1] - goal))

print("Forward noising q(tau_t | tau_0):")
forward = []
for t in range(T):
    tau_t = np.sqrt(abar[t]) * tau0 + np.sqrt(1.0 - abar[t]) * noise
    forward.append(tau_t)
    print(f"  t={t+1}  abar={abar[t]:.2f}  dist_to_goal={endpoint_dist(tau_t):.3f}")

def denoise_to_tau0(tau_t):
    # Stand-in for the trained denoiser: pull waypoints back onto the y=0 line to (1,0).
    x = np.clip(tau_t[:, 0], 0.0, 1.0)
    return np.stack([np.linspace(x.min(), 1.0, H), tau_t[:, 1] * 0.1], axis=1)

print("\nReverse denoising p_theta(tau_{t-1} | tau_t, goal):")
tau = forward[-1].copy()
for t in reversed(range(T)):
    tau0_hat = denoise_to_tau0(tau)
    if t > 0:                                       # DDIM-style deterministic step
        eps_hat = (tau - np.sqrt(abar[t]) * tau0_hat) / np.sqrt(1.0 - abar[t])
        tau = np.sqrt(abar[t-1]) * tau0_hat + np.sqrt(1.0 - abar[t-1]) * eps_hat
    else:
        tau = tau0_hat
    print(f"  t={t}  dist_to_goal={endpoint_dist(tau):.3f}")

Read the two columns together. Forward noising monotonically drives the endpoint away from the goal as \(\bar\alpha_t\) shrinks; reverse denoising walks it back to \(0.016\), essentially the goal. A real planner replaces denoise_to_tau0 with a trained network \(\epsilon_\theta(\tau_t, t, c)\), and replaces the single noise draw with fresh Gaussian samples so that repeated runs produce different valid plans. The point is not that this toy "solved planning"; it is that the same forward and reverse machinery scales from this 6-point line to a 16-step manipulation action chunk conditioned on a camera observation.

With the forward and reverse machinery now transparent on a toy trajectory, the remaining question is how to introduce a diffusion planner into a real project without letting its generative flair outrun disciplined evaluation.

Practical Recipe

1. Write the observation, action, horizon, and success metric before choosing a model.
2. Build a baseline that is simple enough to debug by inspection.
3. Add the maintained implementation only after the baseline behavior is understood.
4. Save one artifact containing configuration, seed panel, traces, metrics, and failure labels.
5. Run at least one perturbation test before trusting the result.

A diffusion planner earns its place only when the contract is explicit: observation stream, physical or latent state, action representation, timing budget, and one evaluation artifact, all fixed before any model comparison. That contract answers the four questions an agent must resolve for any method: what changes in the loop, why it should help, how it is measured, and when to reject it.

Diffuser, Decision Diffuser, and Diffusion Policy give three different views: planning trajectories, conditioning decisions on return, and generating robot actions from observations. The skeptical-reader test is simple: a claim about Diffusion models as planners must identify the baseline, the shared seed panel, the horizon or task split, and the failure labels saved in one artifact.

Keep one inspectable probe for the model assumption, then use maintained libraries without changing the artifact schema used for baseline comparison.

1. Write the observation, action, state estimate, success metric, and rejection criterion.
2. Run a deterministic smoke test on one seed and save the complete configuration.
3. Add one perturbation tied to the section topic: delay, noise, horizon length, contact change, distractor object, or generated-scene shift.
4. Compare only methods evaluated by the same script, split, seed panel, and metric definition.
5. Record a postmortem that assigns failures to perception, representation, dynamics, planning, control, data coverage, timing, or evaluation.

When Diffusion models as planners fails, do not collapse the result into a single method verdict. Assign the failure to the interface that broke, rerun one controlled perturbation, and keep the trace next to the metric. That habit turns a disappointing rollout into a reusable diagnostic asset.