Conditioning Matters: Training Diffusion Policies is Faster Than You Think

TL;DR

This paper introduces Cocos, a simple method to improve the training efficiency of diffusion policies for vision-language-action models by preventing loss collapse through condition-dependent source distributions, leading to faster convergence and better performance.

Contribution

The paper proposes Cocos, a novel approach that modifies source distributions in diffusion policy training to enhance condition integration and training efficiency.

Findings

Cocos accelerates diffusion policy training convergence.

Cocos achieves higher success rates with fewer parameters.

The method is compatible with various policy architectures.

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.