Quantizing Diffusion Policy via Dynamic Scaling and Reweighted Distillation
Background: AgileDP rollouts at W4A4 from our simulation evaluation
Diffusion Policy is a state-of-the-art way to learn robot behaviour, but every action it outputs costs tens to hundreds of passes through the same network. On an on-board computer that is too slow. Quantization, running the network on low-bit integers instead of 32-bit floats, is the standard remedy, yet existing methods for diffusion models stop at 8 bits or quantize only the weights. Which raises the question:
Can a diffusion policy run entirely on 4-bit integers without losing the fine corrections that manipulation depends on?
We propose AgileDP and trace the failure of 4-bit diffusion policies to two dimensions. Across denoising timesteps, activation ranges shift drastically, so a fixed quantization grid either clips or rounds. TADA (Timestep-Aware Delta Adaptation) rescales activations with a tiny timestep-conditioned MLP that is fitted in closed form and then frozen, adapting the quantization step without per-timestep parameters. Along the action sequence, low-magnitude actions carry large relative quantization error. MERD (Magnitude-aware Error Reweighted Distillation) finds those regions and reweights distillation toward them.
Across 7 tasks in 3 benchmarks, AgileDP stays close to full precision at 4-bit weights and activations, where every prior method collapses, and transfers to real Franka Panda and Unitree G1 robots. Custom INT4 Tensor Core kernels deliver up to 9.6× speedup and 7.0× memory savings.
A network stores weights and activations as 32-bit floats. Quantization snaps each value onto a grid of a few integer levels, spaced by a scale δ: divide, round, clip, and multiply back. Fewer bits mean a lighter model and faster integer math, but a coarser grid. Drag the value or pick a matrix entry to see what is lost.
A linear layer Y = XW runs in three steps. Inputs and weights are quantized to integers, the expensive matrix product runs entirely in integer arithmetic, and one multiplication by the two scales brings the result back to real values.
A diffusion policy starts from pure noise and refines an action sequence step by step: one denoiser is called T times for every action chunk the robot executes. Whatever one pass costs, the robot pays it T times, so making each pass cheaper pays off T-fold. Press Generate to run both policies at their measured speed.
With only 16 levels there is no slack. Prior diffusion-quantization work handles the denoising axis; a diffusion policy adds a second one, the action sequence it predicts. Both have to be fixed.
Early in denoising the input is mostly noise and activations spread wide; by the last steps they shrink by an order of magnitude. One grid cannot fit both ends: calibrate it on narrow late-step activations and early steps get clipped; calibrate on wide early ones and late steps fall onto a couple of levels (rounding). Drag the timestep, or press play.
A fixed grid bounds the absolute rounding error by the same constant everywhere. Along a predicted action sequence the magnitudes vary, so the small actions carry the largest relative error. Those small actions are the fine, contact-rich corrections (grasping, insertion, hand-over) that decide whether a task succeeds. Coarsen the grid and watch the small end.
AgileDP is quantization-aware fine-tuning of a pretrained Diffusion Policy. TADA handles the denoising axis by rescaling activations per timestep before they hit a single fixed grid. MERD handles the action axis by distilling hardest where relative error is largest. Below, Algorithm 1 from the paper plays step by step; click any line or dot to jump, or use the arrow keys.
Instead of learning a separate step size for each of the T timesteps, TADA keeps one step size \(\delta_x\) per layer and multiplies the activation by a timestep-dependent factor first: \(\tilde{x}_t = s(\tilde t)\,x_t,\ \hat{x}_t = Q_U(\tilde{x}_t,\delta_x)\), which is the same as quantizing with the dynamic step \(\delta_x / s(\tilde t)\). The measured effect on a real policy:
Each layer's \(s(t)\) comes from a 49-parameter MLP: one input (the normalised timestep), 16 hidden SiLU units with random, frozen weights, and a Softplus output. The targets are \(s^*_t = \sigma_{\mathrm{ref}}/\sigma(t)\) from the activation RMS on a calibration set, so only the output layer is unknown, and it is a ridge regression with a closed-form solution \(\theta^* = (H^\top H + \mu I)^{-1} H^\top y\). This widget runs that exact fit in your browser on the measured activations above.
An average teacher–student loss weights every action step equally in absolute terms, so the large actions dominate the gradient and the small ones are left with large relative error. MERD keeps a global term over the whole prediction and adds a local term on the top-K steps by relative error score:
Train the toy student below with each loss. Uniform distillation lowers the error everywhere but keeps its shape; MERD flattens it.
On real policies, relative quantization error is tightly anti-correlated with predicted action magnitude before MERD (Pearson r ≈ −0.9). After MERD, magnitude no longer explains it.
7 tasks across Robomimic (Lift, Can, Square, Transport, ToolHang), Push-T and Franka Kitchen, with both a U-Net (DP-C) and a Transformer (DP-T) Diffusion Policy, compared with PTQ, PTQD and SVDQuant. Every success rate below is the average over 3 training seeds × 50 initial states × 2 checkpoints; hover a bar for max / average.
Starting from quantization-aware LoRA fine-tuning (QALoRA) alone, add the modules one at a time. TADA is what makes 4-bit feasible at all; MERD recovers the precise, contact-rich tasks such as Transport.
W4A4 AgileDP runs on our custom INT4 Tensor Core kernels with TADA scaling fused into activation quantization; W8A8 PTQ runs on TensorRT. Race one action chunk (all denoising steps, batch 1) in real time.
Each row starts all four W4A4 policies from the same test initial state, so the only difference is how the policy was quantized. Transformer Diffusion Policy (DP-T), state observations. Bars under each clip give that method's average success on the task over 300 rollouts.
Two Franka Panda tasks and one bimanual Unitree G1 task with an onboard Jetson Orin Nano, U-Net Diffusion Policy, 25 trials each. PTQ fails every trial; AgileDP stays close to full precision. Clips play at 2×.
Quantizing a diffusion policy fails along two axes: activation shift across denoising timesteps and large relative error at low-magnitude actions. AgileDP addresses the first with TADA and the second with MERD, giving robust W4A4 quantization across all 7 tasks in 3 benchmarks, on U-Net and Transformer backbones, and on real robots.
@inproceedings{joung2026agiledp,
title = {AgileDP: Quantizing Diffusion Policy via Dynamic Scaling and Reweighted Distillation},
author = {Joung, Jiyeon and Lee, Seungseop and Kim, Namyoon and Park, Jaeheung and Jang, Keunwoo},
booktitle = {Conference on Robot Learning (CoRL)},
year = {2026}
}We dedicate this work to the memory of our late advisor and co-author, Professor Jaeheung Park, whose guidance and insight were instrumental in shaping this research.
This work was supported by the Technology Innovation Program (RS-2025-25450204) funded by the Ministry of Trade, Industry and Resources (MOTIR, Korea), by the KIST Institutional Program (No. 26E0041, 26Z9021), and by Hyundai Motor Company and Kia.