Skip to content

Conditioning Matters: Training Diffusion Policies is Faster Than You Think

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Diffusion-based VLA models suffer from "loss collapse": when vision-language conditions are hard to distinguish, the conditional flow matching objective degenerates into modeling the marginal (unconditional) action distribution, causing the policy to ignore conditioning inputs. Cocos (condition-conditioned source distribution) fixes this by anchoring the flow matching noise source around semantics extracted from each condition, achieving 2.14× faster convergence and higher success rates without changing the policy network architecture.

Problem

Standard conditional flow matching in VLA training decouples the source noise distribution from the condition: q(z) = q(x₁, c)q(x₀), where x₀ is sampled as pure Gaussian noise independent of condition c. When the policy network struggles to differentiate between conditions, Theorem 1 (Gradient Contraction) shows the gradient update becomes condition-agnostic, collapsing into an unconditional objective that minimizes the expectation over all c. This creates a destructive feedback loop—the weaker the condition utilization, the more the objective marginalizes over c, further reinforcing condition omission.

Method

Cocos replaces the standard source distribution with a condition-dependent Gaussian:

q(x₀|c) = N(x₀; αFφ(E(c)), β²I)

where E is the frozen vision-language encoder and Fφ is a lightweight single-layer Transformer autoencoder trained to compress E(c) into action-space dimensionality via cosine similarity reconstruction loss. During training, x₀ is sampled from this condition-anchored distribution instead of N(0, I); inference follows the same conditional sampling. The scalar α controls prior strength and β (default 0.2) controls uncertainty; α=0, β=1 recovers the baseline. Training proceeds in two stages—pretrain Fφ then freeze it—or jointly via EMA on a target encoder Fφ⁻ to stabilize the evolving x₀ distribution. The rest of the flow matching objective (linear interpolation schedule, RDT policy backbone ~40M params, DINOv2-Base + T5-Base encoders) is unchanged.

Key Contributions

  • Formal characterization of loss collapse in conditional flow matching via Theorem 1 (gradient contraction bound under independent vs. conditional sampling measure)
  • Cocos: a plug-and-play source distribution modification that provably prevents loss collapse by making q(x₀|c) depend on c
  • Empirical evidence that condition omission manifests as high cosine similarity and low norm-shift of hidden states pre/post condition injection, directly correlated with task failure
  • Comprehensive benchmark: 40 LIBERO tasks, 30 MetaWorld tasks, 10 real-world SO-100 tasks, 10 real-world xArm tasks

Results

  • LIBERO average: Cocos 94.8% vs. DP-DINOv2 86.5% (+8.3%) vs. DP-scratch 64.4%; matches fine-tuned π0 and OpenVLA-OFT with only 30K gradient steps vs. π0's 100–250K and OpenVLA-OFT's 50–150K
  • LIBERO Long suite: Cocos 88.7% vs. DP-DINOv2 82.7%, outperforming Seer (a larger non-end-to-end model)
  • MetaWorld average: 74.8% (Cocos) vs. 59.5% (DP-DINOv2) vs. 33.6% (scratch), a +25.7% absolute gain
  • Training speed: 2.14× faster convergence to high performance (>95% on Spatial and Object suites within 5–10K steps)
  • Real world: Consistent improvements on both SO-100 and xArm across all task suites; failure mode analysis shows reduced spatial-reference confusion and object localization errors

Limitations

  • Source distribution instantiated only as fixed-variance Gaussian; more expressive forms (flow-based, attention-guided) not explored
  • β < 0.2 (e.g., 0.1) substantially degrades performance, requiring at least light hyperparameter selection
  • Evaluated exclusively on multi-task imitation learning; scaling to large-scale VLA pretraining not validated
  • Real-world evaluation is limited to 10 trials per task, which may underestimate variance

Relevance to Vision-Language Models

This paper directly probes how VL representations are consumed by generative action models, revealing that standard architectures routinely discard condition inputs rather than integrating them—a failure mode that likely affects any VLM-conditioned generative system, not just robot policies. The theoretical result (Theorem 1) generalizes to any conditional flow matching setup using independent source sampling, making it relevant to image and video generation VLMs as well. For VLM researchers tracking embodied AI, Cocos demonstrates that condition utilization quality—not encoder quality or model scale—can be the dominant bottleneck, and that it can be addressed with a lightweight architectural change to the training distribution rather than to the model itself.