Diffusion Policy, Transformer explained

Almost every model worth studying is an encoder and a decoder. The encoder turns the world into a small set of things worth referring to. The decoder answers by asking that set questions.

So the ceiling is the encoder. Attention can only be as good as the thing it attends to.

We took apart Cheng Chi's Diffusion Policy, transformer variant, and trained it on our arm. Here it is.

01What is actually being denoised

The policy emits a tensor of shape (1, 48, 6). Forty-eight future timesteps, six joint angles each. At 30 frames per second that is 1.6 seconds of motion committed in one forward pass. This is action chunking, and it is the reason imitation learning works at all.

Six numbers, on an arm with five degrees of freedom. Both are true.

The sixth number is the jaw, not a direction. Five joints place the gripper. The sixth only opens it, which is why a six-number action vector describes a five-DOF arm.

And the five you get are not five free directions. shoulder_lift, elbow_flex and wrist_flex all pitch about exactly parallel axes, measured at |axis·axis| = 1.000000. Three joints, one plane. So shoulder_pan yaws to pick a vertical plane, the three pitch joints work inside it, and wrist_roll spins the tool. There is no wrist yaw anywhere in the chain.

Over 200,000 random poses the gripper's approach direction never left the plane shoulder_pan chose, to within 0.43°. You cannot hold the gripper still and re-aim it. Pose the arm below and try.

Drag to move the arm · tap a joint to learn what it does

One row of the action chunk

	joint	value	limit
1	shoulder_pan	−4.0°	±110°
2	shoulder_lift	−29.0°	±100°
3	elbow_flex	+38.2°	±96.8°
4	wrist_flex	+28.4°	±95°
5	wrist_roll	−61.9°	−157/+163°
6	gripper	+1.2°	−10/+100°

1–5 place the gripper. 6 only opens it.
Values are the median action over all 120 episodes of phi_so101_cubes_cylinder_v1. Limits from the URDF.

During training we do not ask the network to produce that trajectory from nothing. We take the real trajectory from a human demonstration, deliberately corrupt it, and ask the network to identify the corruption.

Four lines, and they are the entire training objective:

trajectory	the real 48-step chunk from a demo
eps	Gaussian noise, same shape
k	a random integer in 0 to 99
noisy	√āk · trajectory + √(1−āk) · eps

The network sees noisy and k, and predicts eps. That is it. The loss is mean squared error against the noise we added. Input and output have identical shape, because the model is not generating a trajectory, it is performing a repair operation on one.

The āk is a precomputed lookup table. One number per noise level, 100 entries, fixed before training starts. Where it comes from explains a real problem in our data.

If diffusion itself is new to you, stop here and watch this explainer first. For the full derivation, Stanford's diffusion lecture series is the best free treatment we have found.

02Why the square roots are there

The naive way to corrupt something is to keep adding noise: xₖ = xₖ₋₁ + √βₖ · εₖ. Do that and the variance grows without bound, one step at a time. After a hundred steps the numbers are enormous and nothing in the network is scaled for them.

So instead you shrink the signal by exactly as much as you grow the noise:

xₖ = √αₖ · xₖ₋₁ + √(1−αₖ) · εₖ

The square roots are there because variance scales with the square of a multiplier. To shrink variance by a factor α, you multiply the values by √α. Then the two terms add to

Var(xₖ) = αₖ · Var(xₖ₋₁) + (1 − αₖ)

and if Var(xₖ₋₁) = 1, that is αₖ + 1 − αₖ = 1. Variance conserved. This is what "variance preserving" means, and it is why every activation in the network sees roughly the same scale no matter which noise level it is handed.

Chain the recursion and the whole thing collapses into a single jump. Two steps in, you get one x₀ term and two independent noise terms, and because independent zero-mean Gaussians combine by adding variances, the two noise terms merge back into one. Repeat, and:

xₖ = √āₖ · x₀ + √(1−āₖ) · ε, where āₖ = ∏αᵢ

That product is the lookup table. It is why training can sample a random k and jump straight there in one line, instead of simulating a hundred steps.

03The assumption that is not true on our robot

All of that rests on Var(x₀) = 1. Our actions are normalised min-max to [−1, 1], and it is very easy to read that as unit variance.

It is not. We measured it.

measured variance of normalised actions: 0.313

Min-max puts the extremes at ±1. It says nothing about the bulk. On our pick-and-place data the actions cluster hard in the middle and a handful of outlier frames set the range. One joint came out at variance 0.043, which is a distribution squeezed into a tenth of the interval it was given.

So the schedule is not variance preserving on our data. It is variance ramping: the trajectory starts at 0.313 and climbs to 1.0 as noise takes over.

This is not fatal. LayerNorm inside the network absorbs a constant scale difference, which is most of what this is. But it is a genuine deviation from the assumption the schedule was designed around, and it has a matched pair: clip_sample_range=1.0 is only correct because min-max guarantees the data lives in that interval. Change one and you must change the other.

The clean fix is quantile normalisation, mapping the 1st and 99th percentiles to ±1 instead of the extremes, so outlier frames stop defining the scale for everyone. It is a knob we have not turned yet.

04Six numbers into 256

Now the trajectory becomes tokens. A single linear layer, 6 → 256, applied to each of the 48 timesteps independently. 1,792 parameters.

The obvious objection is that this cannot create information. Six numbers projected into 256 dimensions still lie on a 6-dimensional plane inside that space. Nothing was added.

True, and beside the point. The width is not there to hold the input. It is there to hold everything the token will accumulate later.

Attention scores are dot products. With a 6-dimensional token you can ask roughly six independent questions of it before the directions start colliding. With 256 you have room for hundreds, and more importantly you have room for the eight decoder layers to keep writing context into the same token: where this timestep sits in the chunk, how noisy the input is, what the wrist camera saw, what the neighbouring timesteps are doing. The 6 dimensions are the seed. The other 250 are workspace.

05Why positions have to be told

Attention is permutation invariant. Shuffle the 48 tokens and self-attention produces the same set of outputs, reordered. It has no idea that token 3 comes before token 4. (This is the clearest explanation of position embeddings we know of.)

For a trajectory that is fatal, so we add a position embedding: a learned table of shape (1, 48, 256), 12,288 parameters, initialised from a normal with standard deviation 0.02 and trained like everything else.

Learned, not sinusoidal. A formula earns its place when you need to generalise to positions you never trained on. Here the horizon is always 48. Every entry in the table is visited in every batch. When the vocabulary is small, fixed, and densely trained, a free lookup table is strictly more expressive than any formula, because it can represent any function of position including ones no sinusoid can. The cost is 0.14% of the model.

That is the action path complete: (1,48,6) to (1,48,256), 48 tokens carrying a corrupted trajectory and knowing where they sit in it. They are the queries. Now we need something for them to query.

06How a camera becomes 64 numbers

Three cameras, wrist and front and top, each with its own ResNet-18. Not shared. 11,197,088 parameters each, 33,591,264 total, on the bet that a wrist view and an overhead view deserve different filters. ACT bets the opposite and shares one backbone. We have never A/B'd it.

One camera's journey:

step	shape
raw frame	(3, 480, 640)
resize, then crop	(3, 216, 288)
ResNet-18 body, stride 32	(512, 7, 9)
Conv2d 512 → 32, kernel 1×1	(32, 7, 9)
spatial softmax	(32, 2)
flatten, Linear 64 → 64	64

Look at that fourth row. The feature map is 7 × 9, which is 63 cells. Each cell summarises a 64 × 64 pixel region of the original frame. That is the entire spatial resolution the policy has to work with.

Then the pivotal operation. It is a swap, not a compression.

The 1×1 convolution mixes across channels only, touching no neighbours. ResNet gave 512 learned detectors at each of 63 cells; this asks which 32 combinations are worth locating. Then softmax runs over the 63 cells, so each of the 32 channels independently becomes a probability map over positions. Finally you take the expectation against a fixed coordinate grid: sum of probability times coordinate.

Before: channel k says how strongly feature k fires, at every location. After: channel k says where feature k is. The height and width axes are consumed and replaced by an axis of size two.

Because it is a weighted mean and not an argmax, a keypoint can land between cells. Precision comes from how peaked the softmax is, not from the 7×9 spacing. That matters enormously at 63 cells.

between cells. precision comes from peakedness, not the grid. Before: channel k says how strongly feature k fires, at every cell. After: channel k says where feature k is. The output is not features-at-locations. It is locations.

Three cameras at 64 each is 192, plus the 6 joint angles, gives 198 numbers per observation step. We feed two consecutive steps, t−1 and t, so the decoder can infer velocity. That is the (1, 2, 198) the memory block receives.

A 33-million-parameter vision stack, on three cameras, exists to produce 198 numbers. Everything the policy will ever know about the world passes through that gap.

07The memory is three tokens

Two very different things get assembled into what cross-attention reads.

The observation goes through Linear(198 → 256), 50,944 parameters, giving two tokens.

The diffusion step k becomes the third. It is a single integer, and it gets a sinusoidal embedding: 128 frequencies swept geometrically from 1.0 down to 1/10000, each dimension oscillating at its own rate, sine and cosine concatenated to 256.

Two measured properties earn the sinusoid here, and neither is the reason usually given. Adjacent noise levels get embeddings with cosine similarity 0.9723, so the network can interpolate behaviour across k rather than memorising 100 unrelated cases. And every embedding has the same norm, so the conditioning signal never grows or shrinks with k. Smoothness and constant norm. Not, as we originally wrote down and later had to correct, anything to do with sampling non-integer timesteps.

Then the three are concatenated along the sequence axis, not the feature axis. (1,1,256) and (1,2,256) become (1,3,256). Nothing is summed. The widths already match at 256, which is precisely why both branches projected to 256 first. Add a learned position embedding of 768 parameters so the decoder can tell the three apart, and the memory is:

token	carries
0	how noisy this trajectory is
1	what I saw one step ago
2	what I see now

One strange detail. That block is called the encoder, but at our configuration it is a two-layer MLP, 256 → 1024 → 256, applied to each token independently. There is no attention in it. The noise-level token and the two observation tokens never mix. They only meet inside the decoder's cross-attention, filtered through whatever each action token chooses to ask for.

That MLP costs 525,568 parameters, 5.8% of the model, for a position-wise map that does no mixing. It is the same size as one decoder feed-forward. Turning it into real attention layers is a one-line config change we have not tested.

08The decoder: 48 questions, 3 answers

Eight identical layers. Each has three sublayers, each wrapped in a residual with the normalisation on the way in rather than the way out, which is the standard cure for deep transformers diverging.

Self-attention lets the 48 action tokens look at each other. Queries, keys and values all come from the trajectory, packed into one (768, 256) matrix so a single matmul produces all three. Four heads, 64 dimensions each. The score matrix is 48 × 48. This is where smoothness comes from: timestep 20 can see what timesteps 19 and 21 are doing and avoid proposing a discontinuity.

Cross-attention is where the world enters. Queries still come from the 48 action tokens, but keys and values come from the 3 memory tokens. The score matrix is 48 × 3. Every action token, independently, decides how much to weight the noise level against the previous frame against the current frame.

Forty-eight questions, three answers, each question weighting them differently. That asymmetry is the design.

Feed-forward is 256 → 1024 → 256 with a GELU between, applied per token with no mixing. It is the largest of the three at 525,568 parameters.

After eight of those, a final LayerNorm and a Linear(256 → 6) project every action token back down to six joint values. Output shape (1, 48, 6), identical to the input. That is the predicted noise.

At inference we run the decoder 16 times, subtracting a little of the predicted noise each pass. The memory is computed once and reused across all 16, which is why the vision cost is paid a single time per chunk rather than sixteen times.

09Why LayerNorm is in there at all

Each of those three sublayers is wrapped the same way:

x = x + sublayer(LayerNorm(x))

Note the order. The normalisation happens on the way in, not on the way out. That is pre-norm, and it is the standard cure for deep transformers diverging during training: the residual path from input to output stays clean, so a gradient at layer 8 can reach layer 1 without being rescaled eight times.

But what does the normalisation itself do? It works across the 256 features inside one token, independently for every token. Not across the batch. Not across the 48 timesteps.

Take a real activation tensor, (1, 48, 256). Token 0 might arrive with mean 5.08, token 7 with mean 4.21, token 47 with mean 5.12. After LayerNorm every one of them has mean 0 and standard deviation 1. Forty-eight separate normalisations, each looking only at its own 256 numbers.

Two consequences follow, and both matter here.

It is batch-independent. The function computed at inference with batch 1 is identical to the function computed during training at batch 64. BatchNorm cannot promise that. For a policy that runs one sample at a time, sixteen times per chunk, this is not a nicety.

It absorbs scale. Which is exactly why chapter 03 was survivable. Our actions arrive at variance 0.313 instead of 1.0, and the first LayerNorm rescales that away. The schedule's assumption is violated and the network mostly does not care, because normalisation is standing right there.

Then two learned vectors, both of length 256, put a chosen scale and offset back:

γ · (x − μ) / σ + β

γ initialises to 1 and β to 0, so the layer starts as pure normalisation and learns its way out. If a feature genuinely wants a large mean, it can have one. It just has to ask.

10The whole model, and where the money goes

block	params
input embedding, Linear 6 → 256	1,792
action position embedding, 48 × 256	12,288
observation embedding, Linear 198 → 256	50,944
memory position embedding, 3 × 256	768
encoder MLP, 256 → 1024 → 256	525,568
decoder, 8 × 1,053,440	8,427,520
final norm and head	2,054
denoiser total	9,020,934
three vision encoders, on top	33,591,264

That ledger closes exactly, which is the strongest evidence we have that the port is wired correctly. It also matches the parameter count in the Diffusion Policy paper's own table, and widening the embedding to 768 gives 80,540,166 which matches their next row.

Now read the last two lines together. The vision encoders are 3.7× the size of the entire denoiser. And widening the transformer from 256 to 768, an 8.9× increase in denoiser parameters, cost us 0.5% in wall-clock time. Vision dominates completely.

Which brings the post back to where it started.

11Why we went looking at the encoder next

The decoder is not doing anything exotic. Forty-eight tokens ask three tokens some questions, eight times. It closes on the paper's numbers, and it is not where the ceiling is.

The ceiling is the 64. Three cameras, 33 million parameters, 921,600 numbers per frame, all funnelled into 32 keypoint locations per view. Colour is gone. Texture is gone. Which of two similar cubes you are looking at is gone. If the task needs any of that, no decoder can recover it, because it was never in the memory to attend to.

That is a bet, not a mistake. Geometry may be all a pick-and-place policy needs, and 63 cells is a strong inductive bias toward exactly that. But it is testable, and it is where we think the interesting work is.

So that is where we went. SmolVLA keeps a pretrained vision-language backbone instead of training three ResNets from scratch, so the encoder arrives already knowing what objects are. Patch Policy keeps dense vision-transformer patches rather than pooling them to keypoints at all, and beats a 7.6-billion-parameter model at 0.7% the size.

Different answers to the same question: what should the encoder hand over. Everything downstream is just attention pointed at whatever survived.

Every number in this post is measured from our own trained checkpoint on an SO-ARM101, not quoted from a paper. The full derivations, the parameter ledger that closes to the digit, and the measurements we got wrong and later corrected are in the phi repo under docs/theory/.