Conditioning Matters: Training Diffusion Policies is Faster Than You Think

TL;DR

This work tackles the inefficiency of diffusion policies for vision-language-action robotics by identifying a loss-collapse phenomenon when conditioning signals are difficult to distinguish. It introduces Cocos, a simple, plug-in modification that grounds the source distribution in condition semantics via a condition-conditioned Gaussian, supported by theory showing prevention of collapse and extensive experiments across simulation and real robots. Results show faster convergence and higher success, with compact models rivaling large pretrained VLAs, validating Cocos as a general, lightweight improvement for diffusion policy training. The approach emphasizes that conditioning-aware priors, not just larger models, can substantially boost performance and robustness in embodied AI.

Abstract

Diffusion policies have emerged as a mainstream paradigm for building vision-language-action (VLA) models. Although they demonstrate strong robot control capabilities, their training efficiency remains suboptimal. In this work, we identify a fundamental challenge in conditional diffusion policy training: when generative conditions are hard to distinguish, the training objective degenerates into modeling the marginal action distribution, a phenomenon we term loss collapse. To overcome this, we propose Cocos, a simple yet general solution that modifies the source distribution in the conditional flow matching to be condition-dependent. By anchoring the source distribution around semantics extracted from condition inputs, Cocos encourages stronger condition integration and prevents the loss collapse. We provide theoretical justification and extensive empirical results across simulation and real-world benchmarks. Our method achieves faster convergence and higher success rates than existing approaches, matching the performance of large-scale pre-trained VLAs using significantly fewer gradient steps and parameters. Cocos is lightweight, easy to implement, and compatible with diverse policy architectures, offering a general-purpose improvement to diffusion policy training.

Figures (11)

• Figure 1: Fusing generative condition into the source distribution greatly simplifies diffusion policy training. Diffusion policy trained with our method achieves $\pi_0$ performance on the LIBERO benchmarks with only 30K gradient steps, which is 2.14x faster than the vanilla model. We also show the cosine similarity and the norm scale change between the policy hidden states before and after injecting condition information, demonstrating that our method fundamentally compels the policy network to utilize condition information, rather than simply embedding conditions into the source distribution.
• Figure 2: Diffusion Policy w/ Cocos. Our approach requires only replacing the standard Gaussian source distribution with a condition-conditioned Gaussian. An autoencoder compresses condition representations to match dimensionality, providing the mean while maintaining a fixed standard deviation.
• Figure 3: Loss collapse causes the policy to degrade to an average. The policy omits language instructions ('Move it to the left' or 'Move it to the right') and yields actions based on their frequency in the training data. Cocos prevents loss collapse, enabling the policy to produce distinct actions corresponding to 'Left' and 'Right'.
• Figure 4: Learning curve on LIBERO benchmark. Dashed line scores are reported by kim2025openvlaoft.
• Figure 5: Evaluation results on SO100 and xArm platforms. We collect 4 task suites or 10 tasks for each robot platform. Each task provides 20 demonstrations and is tested over 10 trials.

Theorems & Definitions (4)

• Lemma 1
• Theorem 1: Gradient Contraction under Independent Sampling
• proof : Proof of \ref{['thm:1']}
• proof : Proof of \ref{['thm:2']}