Contractive Diffusion Policies: Robust Action Diffusion via Contractive Score-Based Sampling with Differential Equations

TL;DR

The paper addresses instability and data-efficiency challenges in score-based diffusion policies for offline control. It introduces Contractive Diffusion Policies (CDPs), which add a contraction-regularized training objective that couples the forward diffusion drift with a learned score, and enforces a negative-spectral-radius condition on the score Jacobian via efficient proxies like power iteration or a Frobenius surrogate. The authors establish theoretical links between contraction in the reverse diffusion and robustness to solver and score-matching errors, including bounded action variance and seed sensitivity. Empirically, CDPs improve offline policy performance on D4RL and Robomimic benchmarks, particularly under limited data, and demonstrate feasible integration with existing diffusion backbones and offline learning frameworks, with real-world robotic experiments validating practical impact.

Abstract

Diffusion policies have emerged as powerful generative models for offline policy learning, whose sampling process can be rigorously characterized by a score function guiding a Stochastic Differential Equation (SDE). However, the same score-based SDE modeling that grants diffusion policies the flexibility to learn diverse behavior also incurs solver and score-matching errors, large data requirements, and inconsistencies in action generation. While less critical in image generation, these inaccuracies compound and lead to failure in continuous control settings. We introduce Contractive Diffusion Policies (CDPs) to induce contractive behavior in the diffusion sampling dynamics. Contraction pulls nearby flows closer to enhance robustness against solver and score-matching errors while reducing unwanted action variance. We develop an in-depth theoretical analysis along with a practical implementation recipe to incorporate CDPs into existing diffusion policy architectures with minimal modification and computational cost. We evaluate CDPs for offline learning by conducting extensive experiments in simulation and real-world settings. Across benchmarks, CDPs often outperform baseline policies, with pronounced benefits under data scarcity.

Figures (21)

• Figure 1: CDP's performance compared to diffusion policy baselines in offline and tasks.
• Figure 2: Contraction in diffusion sampling. We observe that contraction pulls nearby diffusion flows closer. Contraction plays a critical role in diffusion sampling for offline learning: it dampens solver and score-matching errors while reducing unwanted variance in the generated actions.
• Figure 3: 2D toy experiments. As we train and a standard policy, both methods produce increasingly accurate actions. However, the actions generated by tend to concentrate near the mean of distinct action modes, and show signs of mitigating solver and score-matching drifts.
• Figure 4: Methodology overview. The policy is trained on offline data to minimize contraction and diffusion losses. For each batch of data, the score Jacobian $J_{{\epsilon_\bm{\theta}}}$, is efficiently computed for all denoising steps, and is then penalized with the contraction loss. At deployment, the diffusion policy is frozen, and the sampling process generates the actions given observations.
• Figure 5: Experiments on reduced datasets. We report the average episode return across different random seeds. When training on only 10% of the original dataset size, we observe that decisively outperforms the baselines. This improvement stems from the ability of contraction to dampen score-matching errors, which are amplified in low-data regimes.

Theorems & Definitions (7)

• Theorem 2.1
• Definition 2.1
• Theorem 3.1: Interplay of score Jacobian and contractive sampling
• Corollary 3.1.1: Bounded action variance under contraction
• Lemma 3.1: Power iteration austin2024power_iters
• proof
• proof