Conditioning Matters: Training Diffusion Policies is Faster Than You Think

Zibin Dong (Tianjin University) · Yicheng Liu (Nanjing University) · Yinchuan Li (Huawei Technologies Ltd.) · Hang Zhao (Harbin Institute of Technology) · Jianye Hao (Tianjin University)
condition integrationconditional diffusionconditional flow matchingconvergencediffusion policiesempirical resultsgradient stepsloss collapsepolicy architecturesreal-world benchmarkssimulation benchmarkssource distributionsuccess ratestheoretical justificationtraining efficiencyvision-language-action models

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.