The problem
Every policy so far answers a question of the form "given this observation, what action?" with one deterministic action, fit by pushing the prediction toward the demonstrated action under a squared or absolute penalty. That penalty has a precise and unavoidable optimum: the conditional mean action. When the demonstrator was consistent, the mean is a fine answer. When the demonstrator (or the task) admits several good actions from the same state, the mean is a point between them — and between two ways around an obstacle is into it.
You cannot fix this by making the network bigger or training it longer; a sharper fit to the mean is still the mean. You have to change what the network is asked to produce: not the average action, but a draw from the distribution of good actions. A model that can sample "left" half the time and "right" half the time — and crucially, never "straight into the middle" — is a generative model. Diffusion is the one we build here.
Build
diffusion.py is one file, about 450 lines, in seven regions: setup, core, toy,
data, train, eval, report. The core region is the diffusion method itself — a
noise schedule, the forward noising, an epsilon-prediction denoiser, and the
reverse sampling loop — shared by the toy and the policy. Nothing is imported to
hide it: no diffusers, no einops, no sampler you cannot read.
The core: noise in, sample out
The forward process $q(x_t \mid x_0)$ jumps straight to any noise level in one shot, and the reverse loop steps back along the DDPM posterior mean $\tilde{\mu}_t$ — and here they are in code:
# The diffusion machinery, shared by toy and policy. Nothing from a diffusion# library — the schedule, the noising, and the reverse sampler ARE the method. def sinusoidal_embed(t: torch.Tensor, dim: int) -> torch.Tensor: """Integer timesteps (B,) -> (B, dim) sinusoidal features (the DDPM/transformer time embedding). Lets one network condition smoothly on WHICH noise level it is undoing, without a separate weight per step.""" half = dim // 2 freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device) / half) ang = t.float()[:, None] * freqs[None] return torch.cat([ang.sin(), ang.cos()], dim=1) def make_schedule(steps: int, device: torch.device, broken: bool = False) -> dict: """Precompute the forward-process constants. `acp` is alpha-bar (cumprod of 1-beta): the fraction of ORIGINAL signal still present at each step. A good schedule drives acp from ~1 (clean) to ~0 (pure noise) so sampling can honestly start from N(0, I); the cosine schedule (Nichol & Dhariwal) reaches ~0 even at modest step counts. `broken` is the Break It: near-zero betas leave acp~1, so the model only ever removes a whisper of noise and never learns the trip back from pure noise.""" if broken: betas = torch.linspace(1e-6, 1e-4, steps, device=device) else: u = torch.linspace(0, steps, steps + 1, device=device) / steps acp_full = torch.cos((u + 0.008) / 1.008 * math.pi / 2) ** 2 acp_full = acp_full / acp_full[0].clone() betas = (1 - acp_full[1:] / acp_full[:-1]).clamp(1e-8, 0.999) alphas = 1.0 - betas acp = torch.cumprod(alphas, dim=0) # acp_prev (acp shifted one step) feeds the reverse posterior; post_sigma is # the noise injected per reverse step — the true posterior variance beta-tilde, # smaller than beta_t, which samples noticeably cleaner than beta_t. acp_prev = torch.cat([torch.ones(1, device=device), acp[:-1]]) post_var = betas * (1.0 - acp_prev) / (1.0 - acp) return {"steps": steps, "betas": betas, "alphas": alphas, "acp": acp, "acp_prev": acp_prev, "sqrt_acp": acp.sqrt(), "sqrt_one_minus_acp": (1.0 - acp).sqrt(), "post_sigma": post_var.clamp_min(1e-20).sqrt()} def q_sample(x0: torch.Tensor, t: torch.Tensor, noise: torch.Tensor, sch: dict) -> torch.Tensor: """Forward process: jump straight to noise level t in one shot (no loop needed). x_t = sqrt(acp_t) * x0 + sqrt(1-acp_t) * noise.""" return sch["sqrt_acp"][t][:, None] * x0 + sch["sqrt_one_minus_acp"][t][:, None] * noise def diffusion_loss(model: nn.Module, x0: torch.Tensor, cond, sch: dict) -> torch.Tensor: """The entire training objective. Pick a random noise level per sample, noise x0 to it, and ask the network to PREDICT THE NOISE it added (epsilon-prediction). Plain MSE on that noise — but because the target is the noise, not the action, minimizing it never averages two actions into a third one.""" t = torch.randint(0, sch["steps"], (len(x0),), generator=gen) noise = torch.randn(x0.shape, generator=gen) x_t = q_sample(x0, t.to(device), noise.to(device), sch) return F.mse_loss(model(x_t, t.to(device), cond), noise.to(device)) @torch.no_grad()def p_sample_loop(model: nn.Module, shape, cond, sch: dict, log=None) -> torch.Tensor: """Reverse process: start from pure noise and walk it back to a sample, one denoising step at a time (ancestral DDPM). Each step subtracts the model's predicted noise, rescales, and re-injects a little fresh noise (except the last step). `log` records the shrinking cloud to rerun so you can watch noise become structure.""" x = randn(shape) for step in reversed(range(sch["steps"])): t = torch.full((shape[0],), step, dtype=torch.long, device=device) eps = model(x, t, cond) # Recover the predicted CLEAN sample x0 from the noise and CLIP it: a tiny # model can predict a wild x0 early on, and clamping to the data's scale # keeps the trajectory on the manifold (Diffusion Policy does this — here it # is the step from a policy that fails to one that works). Then step to # x_{t-1} via the DDPM posterior mean of (x0, x). acp, acp_prev = sch["acp"][step], sch["acp_prev"][step] x0 = ((x - sch["sqrt_one_minus_acp"][step] * eps) / sch["sqrt_acp"][step]).clamp(-X0_CLIP, X0_CLIP) mean = (sch["betas"][step] * acp_prev.sqrt() / (1.0 - acp) * x0 + (1.0 - acp_prev) * sch["alphas"][step].sqrt() / (1.0 - acp) * x) x = mean + (sch["post_sigma"][step] * randn(shape) if step > 0 else 0.0) if log is not None: rr.set_time("denoise_step", sequence=sch["steps"] - step) rr.log(log, rr.Points3D(_xy0(x))) return x class Denoiser(nn.Module): """The noise-predictor, used for BOTH the toy and the policy. In: a noisy sample x_t, its noise level t, and (optionally) a conditioning vector (the obs for the policy, nothing for the toy). Out: the predicted noise, x_t's shape. Just an MLP — the lesson is the OBJECTIVE and the SAMPLER, not the architecture (real Diffusion Policy uses a U-Net; see 'What we cut'). When conditioned, obs normalization lives inside as buffers so a checkpoint carries its own stats.""" def __init__(self, x_dim: int, cond_dim: int, hidden: int, obs_min=None, obs_range=None): super().__init__() self.cond_dim = cond_dim self.net = nn.Sequential( nn.Linear(x_dim + TIME_DIM + cond_dim, hidden), nn.SiLU(), nn.Linear(hidden, hidden), nn.SiLU(), nn.Linear(hidden, x_dim), ) if cond_dim and obs_min is not None: self.register_buffer("obs_min", torch.from_numpy(obs_min)) self.register_buffer("obs_range", torch.from_numpy(obs_range)) def forward(self, x_t: torch.Tensor, t: torch.Tensor, cond) -> torch.Tensor: parts = [x_t, sinusoidal_embed(t, TIME_DIM)] if self.cond_dim: # normalize obs to [-1, 1] inside the model, then condition on it parts.append((2.0 * (cond - self.obs_min) / self.obs_range - 1.0).clamp(-1.0, 1.0)) return self.net(torch.cat(parts, dim=1)) def _xy0(a) -> np.ndarray: """(N, >=2) points -> (N, 3) with z=0 for rerun's 3D point cloud.""" a = a.detach().cpu().numpy() if hasattr(a, "detach") else np.asarray(a) return np.column_stack([a[:, :2], np.zeros(len(a), dtype=np.float32)]).astype(np.float32) schedule = make_schedule(args.denoising_steps, device, broken=BROKEN_SCHEDULE)Diffusion has two processes. The forward process takes a clean sample and adds
Gaussian noise in graded steps until nothing is left but noise; q_sample jumps
straight to any noise level t in one shot, mixing signal and noise so their
powers always sum to one (that is why the coefficients are sqrt(acp) and
sqrt(1-acp) — exercise 3 breaks exactly this). The schedule that controls how
fast the signal fades is a cosine schedule, chosen because it drives the surviving
signal to ~0 by the last step even at modest step counts, so sampling can honestly
begin from pure N(0, I).
The reverse process is what we actually want: start from noise and walk back
to a sample. p_sample_loop runs it one step at a time. At each step the denoiser
predicts the noise in the current point; from that we recover a guess of the clean
sample x0, clip it to a sane scale (a small model can hallucinate a wild x0
early on, and clamping keeps the trajectory on the data manifold — this one line is
the difference between a policy that works and one that doesn't), and step to a
slightly-less-noisy point via the DDPM posterior mean. The denoiser itself
(Denoiser) is deliberately just an MLP: a sinusoidal embedding of the timestep so
one network can handle every noise level, an optional conditioning vector, three
linear layers. The lesson is the objective and the sampler, not the architecture.
The toy: seeing multimodality
# The aha, before any robot. Target: a unit ring — every angle an equally-good# "mode", the center the single point NO sample occupies, exactly where an# averaging regressor lands. The ch1.1 left-vs-right collapse in 2D, where we SEE it.angle = torch.rand(N_TOY, generator=gen) * 2.0 * math.piring = torch.stack([angle.cos(), angle.sin()], dim=1) + 0.03 * torch.randn(N_TOY, 2, generator=gen)ring = ring.to(device) # (a) Diffusion: train the denoiser to predict the noise on ring points.toy = Denoiser(2, 0, args.model_dim).to(device)toy_opt = torch.optim.Adam(toy.parameters(), lr=args.lr)toy_loss = float("nan")for it in range(TOY_ITERS): idx = torch.randint(0, N_TOY, (256,), generator=gen) loss = diffusion_loss(toy, ring[idx], None, schedule) toy_opt.zero_grad() loss.backward() toy_opt.step() toy_loss = loss.item() # (b) Regression baseline: a same-WIDTH MLP (smaller than the timestep-conditioned# denoiser — which only helps the point), trained to map noise -> a point with MSE.# Fresh noise says nothing about WHICH ring point, so the MSE-optimal output is the# average of them all — the empty center. Objective-driven, robust to capacity.regress = nn.Sequential(nn.Linear(2, args.model_dim), nn.SiLU(), nn.Linear(args.model_dim, 2)).to(device)reg_opt = torch.optim.Adam(regress.parameters(), lr=args.lr)for it in range(TOY_ITERS): idx = torch.randint(0, N_TOY, (256,), generator=gen) loss = F.mse_loss(regress(randn((256, 2))), ring[idx]) reg_opt.zero_grad() loss.backward() reg_opt.step() # Sample from both and MEASURE multimodality. `modes_covered` bins on-ring samples# into 8 angular sectors: diffusion should light up all 8, regression none.diff_samples = p_sample_loop(toy, (N_TOY, 2), None, schedule, log="toy/diffusion" if args.rerun else None).cpu().numpy()with torch.no_grad(): reg_samples = regress(randn((N_TOY, 2))).cpu().numpy() def ring_stats(samples: np.ndarray) -> tuple[float, float, int]: r = np.hypot(samples[:, 0], samples[:, 1]) on_ring = np.abs(r - 1.0) < 0.2 sector = (np.floor((np.arctan2(samples[:, 1], samples[:, 0]) + math.pi) / (2 * math.pi) * 8).astype(int) % 8) return float(r.mean()), float(on_ring.mean()), int(len(np.unique(sector[on_ring]))) diff_r, diff_hit, diff_modes = ring_stats(diff_samples)reg_r, reg_hit, reg_modes = ring_stats(reg_samples)print(f"toy multimodality [measured]: diffusion mean_radius {diff_r:.2f} ring_hit {diff_hit:.2f} modes {diff_modes}/8" f" | regression mean_radius {reg_r:.2f} ring_hit {reg_hit:.2f} modes {reg_modes}/8")if args.rerun: # forward: watch structured data melt into noise (negative time = before the reverse pass) for step in range(0, schedule["steps"], max(1, schedule["steps"] // 6)): t = torch.full((N_TOY,), step, dtype=torch.long, device=device) rr.set_time("denoise_step", sequence=step - schedule["steps"]) rr.log("toy/forward", rr.Points3D(_xy0(q_sample(ring, t, randn((N_TOY, 2)), schedule)))) rr.log("toy/target", rr.Points3D(_xy0(ring), colors=(140, 140, 150)), static=True) rr.log("toy/regression", rr.Points3D(_xy0(reg_samples), colors=(230, 102, 90)), static=True)This is the aha, and it is worth running before you read on. We sample the ring,
train the denoiser on it with diffusion_loss (noise a point, predict the noise —
that's the entire objective), and separately train a same-width MLP as a one-shot
regressor. Then we sample both and measure. modes_covered bins the on-ring
samples into eight angular sectors; mean_radius says how far from the center they
sit. Diffusion: 8/8 modes, radius 0.87. Regression: 0/8, radius
0.06 — the dead center. Same width, same data, opposite outcome,
because one learned to sample and the other learned to average. Open the recording
and scrub the denoise_step timeline to watch a cloud of pure noise resolve into
the ring.
The policy: conditioning on the observation
# The policy trains on the SAME scripted PushT demos as ch1.1 — the only change# from BC is the head (a denoiser) and the loss (predict noise, not the action).if args.data is None: # Regenerate every run (never reuse a leftover dir): a cache from a different # --seed/--num_demos would silently train on the wrong data. gen_demos is # deterministic, so same args -> bit-identical demos, built or rebuilt. args.data = args.out / "demos" if args.data.exists(): shutil.rmtree(args.data) gen_demos.main(["--episodes", str(args.num_demos), "--seed", str(args.seed), "--out", str(args.data), "--no-video"])if not (args.data / "meta" / "info.json").is_file(): sys.exit(f"no dataset at {args.data} — record one in ch0.4, or generate demos:\n" f" python curriculum/common/envs/pusht/gen_demos.py " f"--episodes 500 --seed 0 --out {args.data} --no-video") from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: E402 (heavy import — after cheap failures) frames = LeRobotDataset("local/pusht-demos", root=args.data).hf_dataset.with_format("numpy")obs = np.stack(frames["observation.state"]).astype(np.float32) # (N, 10) — layout in pusht_env.pyactions = np.stack(frames["action"]).astype(np.float32) # (N, 2) — already clipped to [-1, 1] # The sampler starts from N(0, I), so the data must sit at ~unit scale or the# reverse process fights a scale mismatch: these expert actions have std ~0.3, and# standardizing them to zero-mean/unit-std is what takes this policy from 0% to# working (measured). So we STANDARDIZE the diffused actions and MIN-MAX normalize# the conditioning obs (constant obs dims get range 1 -> a constant, not a divide-by-zero).obs_min = obs.min(0)obs_range = np.where(obs.max(0) - obs_min < 1e-4, np.float32(1.0), obs.max(0) - obs_min)act_mean = actions.mean(0)act_std = np.where(actions.std(0) < 1e-4, np.float32(1.0), actions.std(0))obs_t = torch.from_numpy(obs).to(device)act_t = torch.from_numpy((actions - act_mean) / act_std).to(device)act_mean_t = torch.from_numpy(act_mean).to(device)act_std_t = torch.from_numpy(act_std).to(device)print(f"dataset: {len(np.unique(frames['episode_index']))} episodes / {len(obs)} frames, " f"denoising_steps={args.denoising_steps}, model_dim={args.model_dim}")The policy trains on the same scripted PushT demos as chapter 1.1. The only new
data step is a normalization the sampler forces on us: the reverse process starts
from N(0, I), so the actions we diffuse must live at unit scale too. These expert
velocities have a standard deviation around 0.3, and leaving them there strands the
sampler in a scale mismatch — it pins the policy at 0% success until we
standardize the actions to zero mean and unit variance. So we standardize the
actions we diffuse and min-max normalize the observation we condition on.
# Chapter 1.1's loop, one line changed: the loss is diffusion_loss (predict the# noise on a noised action, conditioned on the obs) instead of MSE on the action.# Unlike BC/ACT we DON'T cosine-decay the lr — the denoiser must keep fitting noise# at all levels; decaying it undertrains (measured: it drops below the baseline).torch.manual_seed(args.seed) # policy init + training-noise stream reproducible, independent of the toygen.manual_seed(args.seed) # give the policy a fresh noise stream (the toy drained the shared gen above)policy = Denoiser(ACT_DIM, OBS_DIM, args.model_dim, obs_min, obs_range).to(device)optimizer = torch.optim.Adam(policy.parameters(), lr=args.lr)shuffle = torch.Generator().manual_seed(args.seed + 1) # torch-side RNG for batch ordertrain_loss, global_step = float("nan"), 0for epoch in range(args.epochs): epoch_loss, num_batches = 0.0, 0 for batch in torch.randperm(len(obs_t), generator=shuffle).split(args.batch_size): loss = diffusion_loss(policy, act_t[batch], obs_t[batch], schedule) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss, num_batches = epoch_loss + loss.item(), num_batches + 1 if args.rerun: rr.set_time("step", sequence=global_step) rr.log("policy/loss/train", rr.Scalars([loss.item()])) global_step += 1 train_loss = epoch_loss / num_batches if epoch % 50 == 0 or epoch == args.epochs - 1: print(f"epoch {epoch:4d} diffusion_mse {train_loss:.5f}")The training loop is chapter 1.1's with a single line swapped: the loss is
diffusion_loss — noise the demonstrated action to a random level, condition the
denoiser on the observation, predict the noise — instead of MSE on the action.
There is one deliberate difference from BC and ACT worth flagging: we do not
cosine-decay the learning rate. A denoiser has to keep fitting noise across every
level, and decaying the rate to zero undertrains it: a decayed run actually finishes
below the untrained baseline, while a constant rate clears it across seeds. It is a
small thing that quietly decides whether the chapter works.
# Loss measured denoising on dataset states; rollouts measure the task. At every# env step we SAMPLE an action: start from noise, run the reverse loop conditioned# on the current obs. --break few_steps runs this at denoising_steps=2 (set in# setup): too few steps can't resolve the action, it comes out under-denoised, and# the pusher wanders — a measured drop, not a crash.def rollout(net: Denoiser, seed: int, tag: str, episode: int) -> tuple[bool, float]: env = PushTEnv() obs_now = env.reset(seed) gen.manual_seed(seed) # seed the sampler from the episode seed: reproducible AND order-independent, so # baseline and trained see the same per-episode noise (a fair, seeded comparison) done, episode_return, info = False, 0.0, {} while not done: cond = torch.from_numpy(obs_now).to(device).unsqueeze(0) # (10,) -> (1, 10) sample = p_sample_loop(net, (1, ACT_DIM), cond, schedule) # in standardized action space action = (sample * act_std_t + act_mean_t).clamp(-1.0, 1.0)[0].cpu().numpy() obs_now, reward, done, info = env.step(action) episode_return += reward if args.rerun: rr.set_time("sim_time", duration=episode * (PushTEnv.MAX_STEPS / PushTEnv.CONTROL_HZ) + env.data.time) rr.log(f"eval/{tag}/action", rr.Scalars(action.astype(np.float64))) rr.log(f"eval/{tag}/pos_err", rr.Scalars([info["pos_err"]])) return bool(info["success"]), episode_return def evaluate(net: Denoiser, tag: str) -> tuple[float, float]: # 10_000 + offset: held out from demo seeds (0..num_demos) by construction. outcomes = [rollout(net, 10_000 + args.seed + ep, tag, ep) for ep in range(args.eval_episodes)] success_rate = float(np.mean([s for s, _ in outcomes])) mean_return = float(np.mean([r for _, r in outcomes])) if args.rerun: rr.log(f"eval/{tag}/success_rate", rr.Scalars([success_rate])) print(f"eval[{tag:9s}]: success {success_rate:.2f} mean_return {mean_return:.3f}") return success_rate, mean_return torch.manual_seed(args.seed + 1) # a fixed random-init reference, independent of eval-time RNGbaseline = Denoiser(ACT_DIM, OBS_DIM, args.model_dim, obs_min, obs_range).to(device)baseline_success, baseline_return = evaluate(baseline, "untrained")success_rate, mean_return = evaluate(policy, "trained")At every environment step we sample an action: start from noise, run the reverse loop conditioned on the current observation, un-standardize, clip to the action range, step. The sampler is seeded from each episode's reset seed, so evaluation is reproducible and the trained-vs-baseline comparison is fair (both see the same per-episode sampling noise).
Run it
python curriculum/phase1_imitation/ch1.4_diffusion/diffusion.py --seed 0 --device cpu
| cpu-laptop | expected wall-clock on cpu-laptop: ~1.65 min (measured) | measured |
| mps | wall-clock on mps: not yet measured | pending |
| t4 | expected wall-clock on t4: ~13.88 min (measured) | measured |
| 4090 | wall-clock on 4090: not yet measured | pending |
The result at the default config (seed 0, CPU):
| held-out success | mean return | |
|---|---|---|
| untrained (random-init denoiser) | 0% | −120 |
| trained diffusion policy | 25% | −112 |
The trained policy clears the untrained baseline on both counts, and it holds across seeds. Be honest about the number, though: 25% is below what plain behavior cloning reaches on the same small budget (around 30%). That is not a bug and it is not a disappointment — it is the lesson. Single-action diffusion pays for its generality with sampling noise: every action is a draw, and on a task whose best policy is nearly deterministic, BC's clean conditional mean is simply a steadier hand. The place diffusion's sampling wins is exactly the place BC's mean loses — genuine multimodality — and PushT's scripted expert is too consistent to show it. The ring already showed it; chapter 1.5 sharpens the sampler, and real Diffusion Policy (below) wins by committing a whole horizon of actions per sample instead of one.
rerun outputs/ch1.4-diffusion/diffusion.rrd
What we cut
This is real DDPM, trained the real way, but it is not the full Diffusion Policy, and the gaps matter enough to name:
- We denoise a single action, not an action horizon. Real Diffusion Policy denoises a chunk of the next several actions jointly and executes them receding-horizon. That is where its temporal coherence — and much of its real performance — comes from, and it is why it tolerates sampling noise that hurts our single-step version. Adding a horizon dimension to the diffused vector is the single highest-value next step, and it is the 4090 Scale Lab the map calls for.
- The denoiser is an MLP, not a temporal U-Net. Real Diffusion Policy uses a 1-D convolutional U-Net over the action horizon with FiLM conditioning. Our MLP is legible; a U-Net is what you reach for once there is a horizon to convolve over.
- We condition on the 10-number state, not images. No camera, no ResNet.
- We use plain DDPM sampling. No DDIM (faster deterministic sampling), no EMA of the weights (a standard stabilizer). Both are real, both are omitted for legibility.
None of these is an approximation that silently degrades a number; each is a whole capability left for later, on purpose, so the noise→sample core is readable. The "read the real thing" segment walks Chi et al.'s repo so you can see exactly what these paragraphs left out.
Break it
Two ablations, each a real diffusion misconception with a measured signature (seed 0, CPU).
--break few_steps — "two denoising steps is plenty, it's faster." This forces
denoising_steps = 2 for the whole run. Watch the trap: the training loss goes
down (0.32 → 0.14) because a 2-step schedule is an easier function to fit — and
the samples get worse. The toy drops from 8/8 modes to 7/8 and its radius
falls from 0.87 to 0.70 as points get pulled off the ring; the policy collapses
from 25% to 0%. Denoising steps are not a speed knob you can turn to nothing:
they are how many chances the reverse process gets to walk noise back to the data,
and two is not enough. The low-loss/bad-sample disconnect is the whole point — your
training curve will not warn you.
--break wrong_schedule — "any increasing noise schedule will do." This swaps
in near-zero betas, so the forward process barely adds noise and acp never
approaches zero. The denoiser is then only ever asked to remove a whisper of noise,
and it never learns the trip back from the pure N(0, I) that sampling starts from.
The signature is unmistakable: training loss sits near 1.0 (the model cannot fit
the task at all), the toy overshoots the ring to radius 1.21, and the
policy is at 0%. The schedule is not decoration; it defines the very distribution
the sampler begins in.
The transferable lesson: in diffusion the sampler and the schedule are part of
the model, not afterthoughts. A perfect denoiser with too few steps or a broken
schedule samples garbage, and — the few_steps trap again — the training loss is
the last thing that will tell you.
Read the real thing
"What we cut" named the gaps; this is where you see them in Chi et al.'s code. We
read real-stanford/diffusion_policy pinned at commit 5ba07ac. Three focus
points, each our version set beside theirs — teaching floor against production, not
worse against better.
The reverse sampler and the denoiser. Ours is p_sample_loop in the core
region: a hand-written loop that, at each step, predicts the noise with our MLP
Denoiser, recovers and clips x0, and steps via the DDPM posterior — the whole
method on the page. The real version is conditional_sample in
diffusion_policy/policy/diffusion_unet_lowdim_policy.py (lines 59–96 @ 5ba07ac).
It is the same loop — start from torch.randn, iterate over timesteps, predict,
step — but each denoising step is delegated to diffusers' DDPMScheduler.step,
and the network it calls is a ConditionalUnet1D
(diffusion_policy/model/diffusion/conditional_unet1d.py): a temporal 1-D U-Net
with FiLM conditioning, not an MLP. What they add is a library scheduler you can
swap DDPM for DDIM without touching the loop, and a convolutional net that shares
structure across time. We cut both — our own loop, our own MLP — because there is no
time axis to convolve over and no library worth hiding, which is exactly why reading
it here lands.
The action horizon. Our eval region samples one action per environment step:
p_sample_loop(net, (1, ACT_DIM), ...). The real predict_action (same file, lines
99+) samples a whole trajectory of shape (B, horizon, action_dim) in one
denoise, then executes only a slice — action_pred[:, start:end], n_action_steps
of it — before re-planning (receding-horizon control). This is the single biggest
thing we dropped, and the one that most explains why our 25% trails BC while real
Diffusion Policy leads: denoising a coherent chunk both smooths the trajectory and
dilutes the per-action sampling noise that punishes our single-step sampler.
EMA. We evaluate the raw trained weights straight out of the train region.
Real Diffusion Policy keeps a shadow copy — EMAModel in
diffusion_policy/model/diffusion/ema_model.py — whose step (lines 56–88) blends
each parameter toward the live one under a warmed-up decay (get_decay, climbing
toward 0.9999), and evaluates that. It is a standard stabilizer: cheap, no effect
on the objective you just learned, a few real points at the end. We omit it so the
training loop stays the handful of lines you can read in one breath.
None of these makes our version wrong; each is production hardening wrapped around
the identical noise→sample core. Read next, in order: conditional_sample first
(you already know the loop it runs), then predict_action for the horizon slice,
then conditional_unet1d.py for the U-Net, and ema_model.py last.
Exercises
Four, in exercises/. Two ask you to commit to a prediction before the run is
allowed to answer — the ring's diffusion-vs-regression modes, and what two denoising
steps do to the ring. One is a bug-hunt in the forward-noising equation (one wrong
coefficient, loss still falls, samples still bad). One has you implement the reverse
posterior-mean step from its definition — the one line the whole sampler turns on.
What's next
You now have a policy that samples actions instead of averaging them, and a ring that shows you why that matters. But DDPM's reverse loop is a long, slightly fussy walk — dozens of steps, a schedule to get right, a sampler that injects noise you then have to fight. Chapter 1.5 keeps the exact same file structure and the exact same task and changes about sixty lines: it replaces the denoising objective with flow matching, the pi0-style objective that learns a straight-line path from noise to data. Same generative idea, a cleaner and faster sampler — and the site renders the two files as a diff so you can see precisely what changed.