Conditioning Matters: Training Diffusion Policies is Faster Than You Think

Question 1: Other ways to learn the prior distribution

This is a very insightful question. First, we'd like to categorize image generation and speech synthesis under the umbrella of AIGC. We want to emphasize that the challenges diffusion models face in Embodied AI versus AIGC are likely very different. AIGC currently grapples with generating increasingly complex modal distributions, which places significant demands on a model's generative capacity. In contrast, action chunks are fundamentally low-dimensional, smooth trajectories, meaning the main difficulty for the model lies in accurately understanding the condition $c$.

Therefore, we believe the motivation behind Cocos is different from the AIGC works you're referring to. For AIGC, their prior distributions primarily encode information about the data itself, aiming to improve the realism and fidelity of synthetic data. However, Cocos is designed to enhance the model's ability to differentiate between conditions and to prevent loss collapse.

What about learning a data prior in a supervised learning way? More specifically, using an MSE loss to learn an action prior, and then using a diffusion model to refine this prior distribution? We did experiment with this approach: we removed the autoencoder's decoder and used the encoder as an action prior model, predicting the action directly with an MSE loss. Similar to the autoencoding version, we added Gaussian noise with std=0.2 to the predicted action to represent $q(x_0|c)$. In this case, $x_0$ becomes the "not-so-accurate" action overlaid with Gaussian noise.

Here are the results of this experiment:

We observed that while this approach still improves training efficiency and scores compared to a vanilla Gaussian prior $p(x_0)$, it does not perform as well as autoencoding $c$. This suggests that using an explicit action prior is not necessarily a better choice. We believe that compressing as much information as possible from the condition (e.g., vision-language cues) into $p(x_0|c)$ might be a more critical choice.

---

---

---

Question 2: Section 2 Clarity

We sincerely apologize that Section 2 was confusing to read. We understand the importance of clarity and will thoroughly revise this section. We'll ensure that all notations, especially the definition of $z$, are explained more precisely. Our goal is to make it much easier for readers to smoothly follow the main argument of the paper.

---

---

---

Thank you once again for your insightful suggestions. We'll incorporate all the discussions above into our revision, conduct a broader survey of prior distributions, and delve deeper into their exploration. We look forward to your further feedback and will respond promptly.

Question 1: Improve Theorem quality

WE SINCERELY APOLOGIZE FOR THE CURRENT QUALITY of this section! Thank you for your suggestion; we will certainly improve the rigor and clarity of this content, including more formal statements. Due to NeurIPS policy this year, we cannot provide external links or PDFs, but we will include the revised version in our updated paper.

---

---

---

Question 2: Design choice about $\mu(x_0|c)$

Our design choice for this form primarily stems from our finding that the way to avoid the loss collapse is to sample training data from $q(x_1, c)q(x_0|c)$ instead of $q(x_1, c)q(x_0)$. We elaborated on this reason in the final part of Theorem 1 (Sorry again for any unclear expressions in the Theorem that might confuse. We will improve it.) Therefore, the simplest way to recover $q(x_0|c)$ is to perform an autoencoding, making $x_0$ a latent representation of $c$, which doesn't necessarily have to be an action prior.

We also experimented with making $q(x_0|c)$ directly an action prior using a supervised learning MSE loss. Specifically, we removed the autoencoder's decoder and used the encoder as an action prior model, predicting the action directly with an MSE loss. Similar to the autoencoding version, we added Gaussian noise with $\text{std}=0.2$ to the predicted action to represent $q(x_0|c)$. In this case, $x_0$ becomes the "not-so-accurate" action overlaid with Gaussian noise.

Here are the results of this experiment:

We observed that while this approach still improves training efficiency and scores compared to a vanilla Gaussian prior $p(x_0)$, it does not perform as well as autoencoding $c$. Overall, we believe modeling action distributions for diffusion models is a relatively easy task—they aren't as complex as modalities like images. The current bottleneck for diffusion policies likely lies in more accurately understanding $c$.

---

---

---

Thank you very much for your valuable suggestions. We will incorporate all the new content discussed above into our revision. Our goal is to refine any imprecise descriptions and avoid all potential ambiguities. We look forward to your further feedback and will respond promptly.

Question 1: Lack of Clarity in Analysis Metrics

Thanks for your feedback. We apologize for the oversight regarding clarity; we'll move the explanations of these metrics from the appendix to the main body in the revision. This will ensure readers can follow the main arguments more smoothly.

---

---

---

Question 2: How does the method prevent loss collapse, or does it only reduce the chance of collapse?

We believe this is a misunderstanding stemming from insufficient clarity in our explanation. In the final part of Theorem 1, we theoretically demonstrate that loss collapse is entirely avoided by sampling training data from $q(x_1, c)q(x_0|c)$ instead of $q(x_1, c)q(x_0)$. This approach ensures that the upper bound of the loss gradient's distance can be arbitrarily large, preventing collapse from occurring. This isn't a heuristic to merely reduce the chance of collapse. We'll adjust the relevant section in the revision to highlight this explanation more clearly and prevent any misunderstandings.

---

---

---

Question 3: Potential Limitations in Diversity and Generalization

We agree that demonstrations in MetaWorld are relatively deterministic. However, Libero offers diversity. For instance, all tasks in Libero-Goal are set in nearly identical environments (except for some randomness in object layout). Demonstrations involve multiple tasks performed in the same scene, requiring the model to fully understand language instructions to differentiate them. Additionally, due to layout randomness and variations in human teleoperator habits, demonstrations themselves are diverse and multi-modal. Furthermore, our real-robot experiments also incorporate randomness in object placement and maintain diversity in manipulation actions during teleoperation data collection. Therefore, we believe our current experimental design ensures sufficient diversity, except for MetaWorld.

---

---

---

Question 4: Questionable Semantic Prior Assumption

As with our response to Question 2, we believe this is a misunderstanding. Our core argument is that the source distribution should contain sufficient information about the condition $c$, not necessarily that it must be an action prior. We conducted further experiments to investigate this:

We also experimented with making $q(x_0|c)$ directly an action prior using a supervised learning loss. Specifically, we removed the autoencoder's decoder and used the encoder as an action prior model, predicting the action directly with an MSE loss. Similar to the autoencoding version, we added Gaussian noise with std=0.2 to the predicted action to represent $q(x_0|c)$. In this case, $x_0$ becomes the "not-so-accurate" action overlaid with Gaussian noise.

Here are the results of this experiment:

We observed that while this approach still improves training efficiency and scores compared to a vanilla Gaussian prior $p(x_0)$, it does not perform as well as autoencoding $c$. This suggests that using an action prior is not necessarily a better choice.

In some settings, similar instructions can require significantly different actions due to scene geometry or context.

We don't believe this scenario is a concern. Even though the source distribution contains instruction information, the diffusion model still learns a distribution-to-distribution transform. This allows it to effectively model multi-modal distributions and, consequently, output significantly different actions even for similar instructions when context or scene geometry demands it.

---

---

---

Thank you very much for your valuable suggestions. We will incorporate all the new content discussed above into our revision. Our goal is to refine any imprecise descriptions and avoid all potential ambiguities. We look forward to your further feedback and will respond promptly.

Question 1: Simplify the Writing of Theorem

Thanks for the suggestion! We'll rewrite this section to make the key points clearer, avoid misunderstandings, and simplify the writing to be much easier to follow.

---

---

---

Question 2: Adding More Experiments on Datasets like Robomimic

We didn't initially consider Robomimic primarily because this benchmark seems to be largely solved. Also, since Libero and Robomimic share the same simulation engine, we opted for the more challenging Libero. To address your concern, we've now included some experiments on Robomimic. The tasks seem to be simple, and we don't see a significant difference in performance:

---

---

---

Question 3: Cocos Scale with Model Size

That's an interesting idea! Testing scaling capabilities requires not only increasing model size but also correspondingly increasing data size. Given the limited time for rebuttal, we couldn't undertake large-scale pretraining. Instead, we tried fine-tuning PI0 (a very powerful and popular pre-trained VLA with 3 billion parameters) with Cocos on Libero. Here are the results:

The results show that Cocos can indeed bring benefits when fine-tuning large-sized VLA models.

---

---

---

Question 4: The Method Still Lacks Behind Methods like Seer

The figures might have been misread; DP+Cocos actually surpasses Seer on Libero. However, it's true that it still doesn't reach the levels of OpenVLA-oft and PI0. Rather than attributing this to model size, we believe it's more likely because these VLA models have undergone large-scale pretraining, whereas DP+Cocos is trained from scratch on specific benchmarks.

Our work primarily aims to explore the capability bottlenecks of current diffusion policies, rather than to release a state-of-the-art VLA model. Therefore, we plan to discuss the application of Cocos in large-scale pretraining more extensively in future work.

---

---

---

Question 5: Can this method be used for general tasks like text to image generation?

Theoretically, our method could be beneficial for tasks like text-to-image generation, as our theory doesn't limit specific modalities. However, we want to emphasize that the challenges faced by diffusion models in Embodied AI versus AIGC are likely very different.

AIGC currently grapples with generating increasingly complex modal distributions, posing significant demands on a model's generative capacity. In contrast, action chunks are fundamentally low-dimensional, smooth trajectories, and the model's main difficulty lies in accurately understanding the condition $c$. This is why AIGC models are developing increasingly powerful modal compressors to reduce the complexity of the distributions diffusion models need to fit, making the design of a data prior potentially more crucial for image generation. Cocos, on the other hand, is specifically tailored to the challenges of diffusion policies in Embodied AI, aiming to help the model better differentiate between various conditions.

---

---

---

Thank you very much for your valuable suggestions. We will incorporate all the new content discussed above into our revision. Our goal is to refine any imprecise descriptions and avoid all potential ambiguities. We look forward to your further feedback and will respond promptly.