Title: Conditioning Matters:  Training Diffusion Policies is Faster Than You Think
Venue: NeurIPS 2025
Accepted: yes
Keywords: Embodied AI; Diffusion policies; Imitation learning; Diffusion models; Flow matching; Vision-language action models
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper shows that diffusion-based imitation learning approaches can suffer from loss collapse, meaning that the diffusion process can not differentiate any more between different conditions. The authors show theoretically that loss collapse can be avoided if the prior distribution of the diffusion process is also conditioned. To do so, the mean of the prior is learned via an auto-encoding framework. The authors show the effectiveness of the approach on several simulated (libero and meta-world) and real robot tasks. Comparisons include baselines to fine-tuned VLA models, where the given approach can reach similar performance with less updates. \nThe paper also contains an illustrative study about the effect of the conditioning on different modalities and can reduce modality over-reliance.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n- Good theoretical result (even though I could not check the proofs in detail) which is underpinned with emperical data\n- Rather simple approach that can be used to extend various diffusion / flow based architectures\n- Convincing results in sim and on real robot\n- Exhaustive analysis of the properties of the algorithms and ablations\n\nWeakness:\n- Writing could be clearer at some places\n- Other ways to learn the prior distribution should be discussed (and compared) in more detail.\n\nQUESTIONS:\n- In different applications (image generation - DALL-E style prior) or PriorGrad in SpeechSynthesis have also learned the prior for diffusion processes. It should be discussed why such ways of learning the pior can not be transferred to this diffusion policy setup\n- Could the prior be also learned by supervised learning (i.e. learning a Gaussian log-likelihood with mean and variance given by a DNN)? This would at least be a good baseline\n- For a reader that is not fully familiar with the conditional flow matching literature, Section 2 was quite confusing to read and it got only clearer once I read over the remaining sections. When introducing the variable z, it should be better explained what this variable corresponds to\n\nLIMITATIONS:\nLimitations are appropriately discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

comments: SUMMARY:
This paper proposes a simple idea to avoid conditional collapse in diffusion policy learning. The authors claim that the reason why condition collapse occurs is because the training objective degerates into modeling the marginal action distribution. To prevent this, they use a simple trick where instead of using a simple isotropic Gaussian to model the noise, the authors employ a conditional gaussian which depends on the conditioning variables. Experiments show that the appraoch is very effective.

STRENGTHS AND WEAKNESSES:
Strengths:
1. The method is simple, yet effective. 
2. Cocos improves over not using it in many experiments.

Weakness:
1. The theoretical section is a bit hard to parse and understand. I would encourage the authors to simplify the writing a bit.
2. Adding more experiments on datasets like Robomimic, and applying Cosos  over different models (different sizes, larger models) or methods can make the experimental section more strong.
3. I would have liked to see more analysis on the conditional vectors. For example, how does the \mu vector look for different conditions. Is there a difference in the distributions when instructions change - for example, move left, move right, etc. Such kind of analysis would make it more interesting.
4. The method still lacks behind methods like Seer. I understand the method is smaller and undertrained. But, if trained with larger models, will it help?

QUESTIONS:
1. Does the performance improvement seen with Cocos scale with model size?
2. Can this method be used for general tasks like text to image generation? I wonder if the findings in this paper help across other tasks.

LIMITATIONS:
I feel the idea is simple and effective which is nice. But the experimental section seems somewhat weak.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a modification to conditional flow matching for training diffusion policies. The authors identify a failure mode called loss collapse, where the policy ignores condition inputs and degenerates to modeling the marginal action distribution. To mitigate this, it replaces the standard Gaussian prior with a condition-aware source distribution derived from a vision-language encoder. Extensive experiments on the LIBERO, MetaWorld, and real-world robot benchmarks demonstrate improved training efficiency and performance.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n- Well-identified failure mode: The analysis of loss collapse is insightful. The authors present both empirical and theoretical evidence of a degenerative training loop, wherein weak condition integration leads to diminished conditioning influence, further reinforcing the collapse.\n\n- Simple and intuitive solution: Modifying the source distribution to be condition-dependent is a natural and plausible fix. Anchoring the prior around condition semantics better aligns the training objective with the desired conditional generation behavior.\n\n- Comprehensive experiments: The method is thoroughly evaluated across two simulation benchmarks and 10 real-world tasks. \n\nWeaknesses \n\n- Lack of clarity in analysis metrics: The use of cosine similarity and norm scale metrics to analyze condition utilization is central to the paper’s claims. However, their interpretations are deferred to the appendix, which makes it harder to follow the main argument.\n\n- Incomplete mitigation of root cause: It is questionable whether the method fundamentally resolves the loss collapse issue or merely reduces its likelihood. The method strengthens conditional input but does not seem to prevent collapse.\n\n- Potential limitations in diversity and generalization: The demonstrations used in LIBERO and MetaWorld are relatively deterministic. It remains unclear how well the method scales to environments with highly diverse or multimodal demonstrations.\n\n- Questionable semantic prior assumption: The assumption that vision-language features offer a good prior for action distributions may not always hold. In some settings, similar instructions can require significantly different actions due to scene geometry or context.\n\nQUESTIONS:\n- How does the method prevent loss collapse, or does it only reduce the chance of collapse?\n\n- How does the method perform when applied to datasets with highly diverse demonstrations?\n\n- Can you provide more insight into what would happen when semantic vision-language priors misalign with the actual distribution (e.g., occlusions, ambiguous phrasing)?\n\nLIMITATIONS:\nThe paper briefly discusses limitations regarding the design of the condition-aware prior (e.g., fixed variance Gaussian) and the scope of current evaluations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper uses conditioning-variable dependent priors as means to increase the performance of flow-based policies, particularly by improving both training time and sensitivity to conditioning variables.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n\n(1) This method clearly works very, very well. It's performance is almost uniformly an improvement over the state of art. \n(2) The method is quite modular and easy to use; in indeed, the principle is compatible with any particular choice of prior. Moreover, the VAE method has competitive performance with optimal \\beta choice, which limits additional hyperparameters, and the EMA version has comparable performance\n(3) COCOs drastically enhances policy expressivity, allowing the model to compete with far more complex architectures\n(4) The authors identify context sensitivity as a a key bottleneck to performance. \n(5) Generally, the experiment were very dilligent and comprehensive. The visualizations representing the actions, i find, were very insightful. \n(6) Overall, the writing is rather clear (with the exception of the theorem below)\n(7) the benchmarks, real hardware experiment, and use of text-conditioning as well as observation conditioning make the result very compelling.\n\nWeakness:\n(1) The theorem is incredible informal. The authors do not define certain key terms in the theorem statement, and the phrasing is colloqiual. Moreover, it is not clear to what extent the theorem \"adds something new\" in the sense that, all that is being said is that if your scores are Lipschitz in c in the network, your network will do similar things. This is rather tautological. I understand what it is trying to say, but I would encourage the authors to include a more formal statement if they intend to include a theorem.\n(2) I am a little unclear about how the encoding of the action behaves. What seems strange is that F(E(c)) doesn't really seem like an action prior, in the sense that it is trained purely via an autoencoder loss, rather than, say a Gaussian approximation to a target action distribution. Said otherwise, the distribution \\mu(x_0 | c) is \"separated\" for different c, but the prior doesn't ensure that mu(x_0 | c) is close to the target p(x_1 | c). I think this is ok - in some sense, you just want to separate out the network gradients, but it is counterintuitive and bears some explanation. Moreover, it might make sense if potentially mu(x_0 | c) is related to p(x_1 | c) in some way (but perhaps it doesn't have to be?)\n\nQUESTIONS:\nCan the authors describe their decision choices for learning mu(x_0 | c) purely via an autoencoder loss? Do the authors think performance could be improved if the action prior was somehow closer to the action distribution?\n\nLIMITATIONS:\nYes, the authors adequately describe limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}