Scaling Diffusion Policy in Transformer to 1 Billion Parameters for Robotic Manipulation

In plain words

A diffusion policy is a robot controller that learns from human demonstrations by generating the next few movements from noise, a bit like an image generator that starts from static. The usual transformer version of this controller stops improving, and sometimes gets worse, when it is made bigger. The authors find that the cause is unstable training: the signal that carries what the robot sees into the network is too large. They change how that signal is injected, and they let each movement in a short planned sequence look at the others, not only at the ones before it. With these two changes the model keeps improving from about 10 million to 1 billion parameters, and it beats the original design on the simulated and real-robot tasks the authors tried.

TL;DR

Diffusion Policy in a transformer (DP-T) does not get better as it grows, and the paper links this to large gradients in the observation-fusion step. ScaleDP swaps cross-attention conditioning for adaptive layer norm and drops the causal mask in action self-attention. Five sizes from 10M to 1B parameters are evaluated on 50 MetaWorld tasks and on 7 real-robot tasks. All numbers are author-reported, and the real-robot evaluation uses 20 trials per task. No official code is linked.

Problem

Diffusion Policy §3 models the distribution of the next action chunk conditioned on observations and proprioception, and the transformer variant (DP-T) is its baseline §3.1. The authors report that adding depth or heads to DP-T does not reliably help on MetaWorld §3.1. With 8 layers, DP-T reaches 80.1% success; with 12 layers 78.4%; with 14 layers 74.6% §1 fig. 1. They also report that the standard deviation of gradient magnitudes grows with depth §3 fig. 1. The paper assumes expert demonstrations with paired visual and proprioceptive observations §3.

Method

Problem setup. The dataset holds demonstration trajectories, each a sequence of visual observations and proprioceptive state §3. The robot is controlled by a 6D pose: position (x, y, z) and rotation (roll, pitch, yaw) §3. How the gripper is commanded is not spelled out in the action definition.

Diffusion objective. The policy is a DDPM over an action chunk A_t = {a_t, ..., a_{t+C}} conditioned on (o_t, x_t) §3. Sampling starts from Gaussian noise and applies the denoising update of eq. 1 until k = 1. The noise predictor ε_θ is trained with an MSE loss on noised action chunks eq. 1 §3. The coefficients α, γ and σ of eq. 1 are called "parameters of the denoising process" and are not given numerically §3.

Change 1: AdaLN instead of cross-attention. DP-T concatenates the diffusion-step embedding k and the observation embedding o into a sequence and attends to it with cross-attention, separate from the action sequence §3.2. ScaleDP removes this path. Instead it takes each standard layer norm and replaces it with adaptive layer norm §3.2 fig. 2. The scale γ(k, o) and shift β(k, o) are regressed from the sum of the embeddings of k and o §3.2, and the layer computes (γ + 1) · x + β eq. 2. The paper motivates this with the normalised-parameter argument used in image generation §1 §3.2.

Change 2: non-causal self-attention. DP-T masks self-attention so that each action token sees only earlier tokens §3.2. ScaleDP removes this mask in the self-attention layers, so every action token in the chunk sees the whole chunk §3.2. The paper argues that this helps because the policy executes only the first action of a longer predicted chunk §1.

Scaling configurations. Following ViT, depth N, width d and the number of heads are scaled together §3.2. Five configurations are defined: Ti, S, B, L and H tab. 1. The head dimension is not stated; it is 64 for Ti, S, B and L and 80 for H when computed from tab. 1 (derived, not stated in the paper).

Output. After the last block, a final AdaLN and a linear layer decode the action tokens into predicted noise §3.2.

Algorithm-level sketch, written from §3 and §3.2. The paper gives no algorithm box, so this is not a verbatim procedure.

# ScaleDP block, one of N stacked blocks  [§3.2] [eq. 2]
# Input:  actions (B, C, d) noisy action tokens
#         cond (B, d) summed embeddings of k and o  [§3.2]
# Output: actions (B, C, d)

def scaledp_block(actions, cond):
    # Regress scale and shift from the condition; the
    # regression network is not stated  [§3.2]
    gamma, beta = regress_scale_shift(cond)    # (B, d) each
    # Modulated layer norm: (gamma + 1) * x + beta  [eq. 2]
    normed = (gamma + 1) * layer_norm(actions) + beta
    # Self-attention with the causal mask removed  [§3.2]
    actions = actions + self_attention(normed, causal=False)
    # Feed-forward sub-layer; whether it is also AdaLN
    # modulated is not stated, this follows §3.2 text  [§3.2]
    normed = (gamma + 1) * layer_norm(actions) + beta
    actions = actions + mlp(normed)
    return actions

# Inference: denoise one action chunk  [§3] [eq. 1]
# Input:  observation o, proprioception x
# Output: action chunk A with C steps

def sample_action_chunk(o, x):
    # Start from Gaussian noise  [§3]
    actions = random_normal((C, action_dim))   # (C, 6)
    for k in range(K, 0, -1):
        # Predict noise and apply one update  [eq. 1]
        noise = model(o, x, actions, k)
        actions = alpha * (actions - gamma * noise)
        actions = actions + random_normal_sigma()
    return actions

# Training step for the denoiser  [§3]
# Input:  o, x, clean action chunk
# Output: scalar MSE loss

def training_step(o, x, clean_actions):
    # Sample a diffusion step and noise  [§3]
    # (schedule and K are not stated)
    k = sample_diffusion_step()
    noise = random_normal(clean_actions.shape)
    noisy = add_noise(clean_actions, noise, k)
    # Predict the noise and compare with MSE  [§3]
    predicted = model(o, x, noisy, k)
    return mean_squared_error(noise, predicted)

Hyperparameters

Model configurations tab. 1:

Other settings:

Results

Author-reported. Real robot: 20 trials per task, so each success rate is a multiple of 5 points. Baselines are DP-T, DP-Unet, Octo, Beso, MDT and ACT §4.1.

MetaWorld (simulation), 50 tasks in four difficulty levels. Text values only: the abstract reports that the largest ScaleDP outperforms DP-T by an average of 21.6% [abstract]. The per-level and per-size curves are in fig. 4 (panels a to f), with no numbers in the text. The paper reports that ScaleDP-Ti beats DP-T on all four levels at similar parameter count [§4.2 Comparison with DP-T] fig. 4a.

Single-arm Franka, four tasks, 20 trials each tab. 2:

Bimanual UR5, three tasks, 20 trials each tab. 3:

The text's headline comparisons follow from these tables. ScaleDP-H exceeds ScaleDP-L by 13.75 points on Franka and 8.33 points on UR5 §4.3. ScaleDP-L exceeds ACT by 7.5 points on Franka and 16.67 points on UR5 [§4.3 result]. The abstract's 36.25% and 75% gains over DP-T on the single-arm and bimanual sets correspond to ScaleDP-H against DP-T (87.50 − 51.25 and 98.33 − 23.33) [abstract] tab. 2 tab. 3.

Non-causal attention ablation, bimanual UR5, 20 trials tab. 4:

Scaling and generalisation, qualitative. The paper reports that larger models gain more from more demonstrations, and that the larger models converge faster on MetaWorld Assembly and on a real task §4.2 fig. 4c–f. It reports appearance, object, light and distractor tests on ScaleDP-L, where the baseline is said to fail on different-coloured targets and to be sensitive to a distractor object §4.4 fig. 5 fig. 6. These tests use one trial each fig. 5 fig. 6, and no numbers are given.

Paper vs code

No official code. Search done: the paper, its abstract and its arXiv record link no code repository. The project page at the abstract's URL links one repository, juruobenruo/DexVLA, a different method whose owner is not an author; the card does not use it. GitHub repository searches for "ScaleDP", "scaling diffusion policy transformer" and "diffusion policy" under the midea-ai organisation returned no match for this paper.

Pitfalls

• Figures are not checkable here. The MetaWorld numbers (21.6%, the curves in fig. 4) appear only as plots or in the abstract. Figure numbers in this card follow source order and were not checked against the compiled PDF.
• Missing training recipe. Optimizer, learning rate, batch size, training steps, diffusion step count K, noise schedule, action chunk length C and the vision encoder are all not stated. An implementer has to choose them and should report them as their own choices, not as the paper's.
• AdaLN is underspecified. eq. 2 gives only the modulation. The paper does not state the regression network, initialisation, whether gating is applied to residual branches, or whether the feed-forward sub-layer is modulated too §3.2. The abstract says the observation feature is "factorize[d] ... into multiple affine layers", and the body describes regressing γ and β from summed embeddings. Implement eq. 2 as written in the body.
• Conditioning changes the sequence length. DP-T puts the conditioning into a separate token sequence for cross-attention. ScaleDP carries it only through the modulation, so the attention sequence is the action chunk alone §3.2.
• Non-causal attention is one paper's ablation. The gain is shown on three bimanual UR5 tasks with 20 trials each tab. 4. The mask question for cross-attention, which DP-T uses, is not discussed §3.2.
• Noise in the real-robot numbers. One trial equals 5 points. Test sets are 20 trials per task, and no confidence intervals appear in the active tables. The source contains commented-out versions with ± values, which are not in the compiled tables.
• Noise coefficient. eq. 1 writes the noise term as N(0, σ²I) at every step, and the training loss adds ε_k directly to A^k without any schedule factor. The paper does not give how noised samples are formed.
• Head dimension at H. H uses 16 heads at width 1280 tab. 1, giving a head size of 80, unlike the 64 of the smaller configs. This is derived from the table, not stated.
• Baseline recipes are not reported. Training settings for DP-T and the other baselines are not stated. Whether they got the same demonstrations, epochs and tuning is not given.
• Uneven demonstration counts. Close Laptop got 40 trajectories and the other real tasks 100 §4.1. The paper does not explain the difference.
• Scaling claims are linked to a fixed protocol. The gains are shown with the paper's tasks, observation setup and demonstration counts, and at single seeds for the real robot.

Lineage

• Builds on: Diffusion Policy, the original DP and its transformer variant DP-T that this paper rebuilds (no card yet).
• Builds on: Diffusion Transformer, the source of the AdaLN conditioning idea, used as in image generation §3.2.
• Compared against: Octo, Beso, MDT, DP-Unet and ACT §4.1 (no cards checked for this card).
• Superseded by: not known from content already in this wiki.