Conditioning Matters:  Training Diffusion Policies is Faster Than You Think
Authors: Zibin Dong
Yicheng Liu
Yinchuan Li
Hang Zhao
Jianye HAO
Venue: NeurIPS 2025 poster ()
Keywords: Embodied AI; Diffusion policies; Imitation learning; Diffusion models; Flow matching; Vision-language action models
Primary area: reinforcement_learning

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.

Conditioning Matters: Training Diffusion Policies is Faster Than You Think
Zibin Dong", Yicheng Liu*, Yinchuan Li Hang Zhao * Jianye Hao" 9.
Tianjin University, *Tsinghua University, Huawei Noah's Ark Lab
Abstract
Diffusion policies have emerged mainstream paradigm for building vision- language-action (VLA) models. Although they demonstrate strong robot control capabilities, their training efficiency remains suboptimal. In this work, identify fundamental challenge in conditional diffusion policy training: when generative conditions are hard to distinguish, the training objective degenerates into mod- eling 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-depeni 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 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 general-purpose improvement to diffusion policy training.
By
Noise
LIBERO AverRee
Codine Similarity
Norm Scale (NQ)
Success Rate 186)
Diffusion Policy
90
Gcoaa
40
ERONDD2 OR-ONO12
(iretand)
OpeniiaA OFT (frnetawe))
*Open the pol
Figure Fusing generative condition into the source distribution greatly simplifies diffusson policy training. Diffusson policy trained with our method achieves лa performance on the LIBERO benchmarks with only 30K gradient steps, which is 2.x4x faster than the vanilla model. We also show the cosine similarity the norm scale change between the policy hidden states before and after injecting condition information, demonstrating that our method fundamentally compels the policy network to utilize condition information, rather than simply embedding conditions into the source distribution.
Introduction
Denoising generative models have emerged as a scalable approach for high-dimensional data genera- tion [28, 6. 26, 17, 18, 12, 30, 16], significantly advancing Vision- Language-Actinn (VLA) models in robotics (VLAs") [19, 15, 1, 3, 36, 11]. Mainstream VLA models frame robot control as a conditional generation problem: given vision-language inputs conditions, the model generates appropriate robot action sequences. Despite action sequences being substantially lower-dimensional than typ- ical generative content like images text, training VLAs remains unexpectedly challenging and
"Corresponding authors: Hang Zhao (hangzhao mail. www.lliiiinguuuuu..o Jianye Hao (jianye.hao tju.eddu.en) Unless otherwise specified, VLAs in this paper refer specifically to denoising generative VLA models.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

w(..)
Denolsing
Cosine Simitarity
www.wa..........................
Decoder (Self (Self-Aftention + MLP)
Cos--
X1
Encoder Gelf-Aftention- MLP)
to
ACI
Adal
T5-Base
DINOV2-Base
Cocos (A
Key/Volue
"Place Moka pot cooker."
Conditionall Input
Observations
Action Chunk
Figure 2: Diffusion Policy w/ Cocos. Our approach requires only replacing the standard Gaussian distribution with condition-conditoon Gaussian. An autoencoder compresses condition representations to match dimensionality, providing the mean while maintaining a fixed standard deviation.
resource-intensive. Recent efforts enhance VLA training efficiency have explored various direc- tions, including leveraging more powerful usioonlanguage encoders [9, 1, 24, 10, 33, 32], designing compact and expressive action tokenizers [22, 31], and incorporating richer conditioning strategies [7, 8, 35, 14, 23]. These approaches converge on a critical insight: the core challenge may not lie in action generation itself, but rather in how models interpret and utilize conditional inputs.
To further investigate how condition inputs affect policy training, we empirically compare diffusion policies with visual encoders trained from scratch versus those initialized with DINOv2. We analyze the cosine similarity and norm scale change between policy network hidden states before and after condition information injection, correlating these metrics with LIBERO task success rates Figure 1. Higher cosine similarity and lower norm scale change indicate weaker condition influence on the policy network, directly correlating with degraded performance (as observed comparing DP-DINOv2 versus DP-scratch). The results demonstrate a fundamental issue: policies facing difficult-to tontterrttt conditions actively omit conditional inputs, producing actions disconnected from observations. This phenomenon persists even when using high-quality representations, though to a lesser degree.
While empirical results reveal this condition omission phenomenon, we conduct theoretical analysis to identify its root cause in VLA training. We discover that when policy networks struggle to differen- tiate between generative conditions, the flow-matching objective degrades into an unconditional one that merely models the marginal action distribution. This loss collapse creates a destructive feedback loop: as the policy network begins ignoring conditional inputs, the training objective further degrades, reinforcing the network's 's tendency to discard conditions altogether rather than attempting to interpret them. To prevent this loss collapse, we introduce novel conditional flw-matching approach with a condition conoonnnooooiiiiiodiccoooooeeoed source distribution (Cocos). Rather than adopting a standard Gaussian prior q==), Cocos anchors the source distribution around the semanties of each condition 7(=lc), theoretically preventing training loss collapse and forcing the policy network to remain responsive to condition inputs. As demonstrated in Figure 1, diffusion policies trained with Cocos exhibit 2. 14x faster training and substantially higher success rates. Critically, the lower cosine similarity and increased norm scale change indicate that Cocos fundamentally compels the policy network utilize condition information, rather than simply embedding conditions into the source distribution.
Specifically, implement Cocos by formulating the source distribution Gaussian with fixed standard deviation, where the mean derives from latent representations produced by a lision-language autoencoder (see Figure 2). This approach introduces minimal architectural overhead, making compatible with diverse VIA architectures and model scales. This simplicity, combined with strong theoretical guarantees and empirical benefits, positions Cocos general-purposs solution for preventing condition omission in VI.A training.
To comprehensively validate Cocos, we conduct extensive evaluations across diverse settings: 70 simulation tasks from the LIBERO and MetaWorld benchmarks [13, 34], 10 real-world tasks on the low-cost open-source S0-100 robot platform [2], and 10 tasks the high-performance xArm robot platform. Our results demonstrate significant improvements in both manipulation success rates and learning efficiency. Through detailed case studies, analyze loss collapse manifestations and the mechanisms through which Cocos enhances performance. Our contributions include:
We formulate the mathematical framework flow matching with generative conditions and demonstrate that policy networks omit conditions when they are difficult to distinguish.

---

We introduce Cocos, a simple yet effective source distribution modification that prevents loss collapse and significantly improves diffusion policy training efficiency and performance.
We establish comprehensive evaluation benchmark across diverse settings, including simulation tasks from LIBERO and MetaWorld, real-world tasks on a low-cost open-source robot (SO100), and tasks a tigh-performance robot (xArm). Our results demonstrate consistent performance improvements across these varied platforms, validating Cocos as general-purpose, plug-and-plas solution for enhancing diffusion policy training.
Preliminaries
Problem Formulation. Our goal is to train VLA (e(():H(((o) that maps observation to sequences of future robot actions @1 by imitating expert demonstrations. Here, o typically includes current and historical images as well language instructions, while & represents a pre-trained vision-language encoder. In this paper, we focus on denoising generative model VLAs.
Conditional Flow Matching. Assume a smooth time- varying vector field u [0.1] x R - R defines an ordinary differential equation (ODE) dr uz(z)df, which pushforward samples from source distribution Po to the target distribution pl. The density transported along from time 0 to / is
denoted as pt [0,1] 1] x Rd - R.. Pt and u, satisfy the continuity equation: 0p/0t -V(p mt). Suppose that the marginal probability path (x)) is mixture of probability paths p=(x=]), that is pe(x)) = f pe (z)))) where ((2) is some distribution over the conditioning variable. If pr (x(]a) generated by the vector field ue (x(22) from initial conditions po (zz|z), then the vector field
ut (2)
ue == Egyt))
Pr (x
generates the probability path (r():) Let no [0. 1] x Rd Id be time-dependent vector field parameterized as neural network 6. Define the conditional flow matching (CFM) objective:
CCFM == (x ||ve Φ) ur(a|=)11
(2)
In VLA models, a typical choice sets ((2) = ((zr))((1) corresponding to an independent coupling of the source and target distributions, and the conditional flow is defined by
(((2)2) (=) uy (z|=) = ue (z(1)11,
where the source distribution is a standard Gaussian. By optimizing CFM objective, vo x) can converge the vector field (I)),, with which can solve the ODE and sample action chunks. Note that this is a general formulation covering a wide range of denoising generative models, e.g., Var. Exploding [27], Var. Preserving [6] diffusion models, flow matching models [12], rectified flow models [16], etc. So our following discussion can be easily transferred to various policy backbones.
3 How Generative Conditions Can Lead to Loss Collapse
While conditional flow matching offers general framework for transporting probability distributions, it does not inherently consider condition inputs. Prior works adopt straightforward strategy: injecting the generative condition directly into the policy network [4, 37, 38]. However, the implications of this approach have received little theoretical attention. In this section, we first analyze how the underlying flow formulation changes when generative conditions are introduced. then demonstrate that this seemingly benign modification can lead to a critical failure mode: loss collapse, where the training objective degenerates and the model learns omit the generative condition altogether.
3.1 Flow with Generative Conditions
We consider a conditional vector field (x(ccc that transports the source distribution Pu(z/c) to the target Pı (z/c), where is the generative condition. Let pr(22, 2 c) denote the time-varyine intermediate density under generation condition c and flow condition 2, with its associated vector field uy (2|22 c). Then, the overall conditional velocity be expressed
ue 2 (a)
Ue (x|c) = Ey(z)
(4)
3

---

"Move to the left."
"Move it to the right."
P(a) "Left")
P(a "right")
P(o: "Leff")=P(c P| "" "Lefi or Right")
w/ Cocos: Policy learns to distinguish conditions
w/o Cocos: Policy degrades to an average
Figure 3: Loss collapse causes the policy to degrade to an average. The policy omits language instructions
("Move it to the left' or "Move it to the right') and yields actions based their frequency in the training data. Cocos prevents loss collapse, enabling the policy to produce distinct actions corresponding to 'Left' and 'Right'.
Suppose introduce into the neural estimator (t x, c) of the conditional velocity (x(zc)). The corresponding training objective becomes:
LeFMe (0) == (x [1,,,,
(5)
Lemma 1. (f/c))))) > 0 for all x E Re and t [0, 1]. up to constant independent of 0, objective (xppe( (1, ue (= and (a) (t) (1(21 e)l are equal.
By optimizing Equation (5), we use (t, c) to estimate ux ())cc) and solve the ODE for generation. In practice, common formulation of q(2) in Equation (5) is ((2) = (( ,, c). 2 = (æ1, Fo, c). This setup corresponds to sampling condition-action pair from the dataset and perturbing the action with Gaussian noise during training. Under this formulation, the conditional flow and velocity field are defined by:
Po (0( .cE
and (x(2,e) = {{ We (e((z211
Pa (a(212 c) =
(6)
.c z
c $
This reduces objective to that implemented previous works [4, 37, 38]:
CCFM ==
((.a) oeo 1l (x(<,,,,
3.2
Loss Collapse
While this conditional objective encourages the model to learn condition -endition-dependent vector fields, a critical failure mode can occur: the model fails to distinguish the condition information in (e.g., for any C1, C2 E C, [|ua(t, x, ct) vt x, e2)ll S c), the objective effectively collapses into marginal form on that no longer depends on c.
Theorem (Gradient Contraction under Independent Sampling). Under the independent sampling measure
p(dy) q(xi, g(x0) pr(x)21, 20) drodæ
the difference of the score gradients for any conditions C1, 02 E C admits the bound
VoLcFMe(O.c)) VoLcrme c2) < 2(M+
provided that oooll S M, [[]]| < D, and the model output satisfies [lvo(t, (1))) vo x, 02)1 S with its gradient being nutput-sensitiv Lipschitz of constant K. The velocity estimator uy x, c) tends to converge to a c-independent function v² (1, that minimizes the averaged objective:
of (1, = arg min BeEc Bzop [1]((t, x:) (x(2)21.000)
By contrast, when samples are drawn from the conditional measure
(e(dy) = q(z1, c) g(zolc) (= dændzıdr,
the same gradient difference can become arbitrarily large unlesy additional coupling assumptions on g(zole) c) are imposed.

---

This result reveals a vicious cycle: as the policy network becomes confused by condition inputs, the objective degenerates into an unconditional one. This further reinforces the network's 's tendency to discard difficulttto-nnterye comprehend conditional cues during deployment. As shown in Figure 3, loss collapse causes the policy to average actions. In the next section, we introduce Cocos, simple method to prevent it.
conditions, leading to models that appear to be well-trained but fail to
4 Method
In Theorem 1, we identify that loss collapse arises when becomes indistinguishable to the policy net- work. This collapse exacerbated by the standard choice of sampling z from g(z) = q(x1, c) (xo( which decouples 2o from the condition. If et c) s (o,t, the expectation in the objective effectively marginalizes over c, causing it to degenerate into an unconditional one.
To prevent it, we propose simple yet effective modification: using a condition distribution (Cocos) by sampling from g(=) = g( c))(sole). As shown in Theorem 1, this change directly avoids the loss collapse and enforces condition-awarenes during training. While (( e) can designed in various ways, in this work, we choose to maintain the overall framework's simplicity by setting it Gaussian with fixed standard deviation, with the following objective:
condition-conditione conditioned source
Cooas (0) == Ez. e). |100- - uz 1| where (()ale) = (xoo a (Fs(((()),
(8)
where F6 feature encoder that maps the condition representations the action space dimension- ality. The scalar a controls the strength of the condition prior, and 3 adjusts the uncertainty, Setting 0, 3 = 1 recovers the commonly used formulation with a standard Gaussian prior. Incorporating Cocos requires only replacing the standard Gaussian source distribution with a condition-conditionss one. Algorithm clearly shows the differences in training and inference pipelines.
In practice, we adopt an autoencoding objective on condition embeddings to train Fg. Specifically, we minimize the negative cosine similarity between the original and reconstructed embeddings:
[(Ø) Sim (E)))
where Fg and G& are the encoder and decoder networks, respectively. As default setting, we adopt a two-stage training process: we first train the source distribution (i.e., the autoencoder) and then fix it during policy training. However, in practical scenarios, this two-step pipeline may introduce additional inflexibility. To address this, we also explore a joint training alternative. Simultaneous training of this autoencoder and the policy can be unstable due to the evolving distribution of zg- To mitigate this, we adopt an Exponential Moving Average (EMA) strategy. We maintain target network F--. updated via EMA from F',, and use this EMA copy during policy training, ensuring stability. practice, this approach achieves performance comparable to the two-stage strategy, which we will discuss in detail in the experimental section.
We implement our policy network using a compact Robot Diffusion Transformer (RDT) [15]. The diffusion process follows linear interpolation schedule [12]:
Pe ( xo) = (2; ta + (1 l)xo, )²),, we (2]2<1, wo) = 21 200
(10)
We select these design choices as they have been widely adopted in advanced VLA models [15, I].
Algorithm 1 Training and Inference Pseudocode for Diffusion Policy with Cocos
Require: Policy Network Condition Network &, Pre-trained condition encoder Fe
1: while not coverge do # Training
2: (21, q(21, c), 3, 1~ ((0, 1),20 ~ ((xo(c) Equation (8) zo (O0, (without Cocos)
3:
0 0+0+ LCLME Equation (7)
4: end while
5: while not done do Inference
6: Observe c, 2h g( Equation (8) zo (V0 7) (without Cocos)
7: Solve ODE dx; = uolt, xt df from t to 1.
8: a + 21; Execute action
9:
end while
5
Experiments
We conduct comprehensive experiments to evaluate the effectiveness of Cocos in both simulated and real-world robot manipulation tasks. The experiments designed to answer three core questions:
5

---

Libero goal
Libero spatial
Libero abject
übero-long
Libero- averrge
Rate
70
15
15
30
Te (hretunde) m OpeniLk (Ihe-tune)
Seer
+
DF-C0N00- Cooss
BROND
Figure 4:
Learning curve on LIBERO benchmark. Dashed line scores are reported by Kim et al. [10].
Table Success rate on LIBERO benchmark. Each the four task suites includes 10 tasks, and we evaluate 150 trials for each task. The 1s. 2t highest scores are emphasized with bold and the 3rd with underline. Cocos (0.2 std) is our default configuration, and other Cocos variants serve for an ablation study.
DP-scratch
DP-DINOv2
Cocos (0.2 std)
Cocos (0.1 std)
Cocos (0.4 std)
Cocos (VAE)
Cocos (EMA)
Goal
64.311.3
82.910.
93.8.0.3
50.911.8
92.3.0.9
92.5g0..
90.510.4
LIBERO
Spatial
16.1116
91.420.7
97,510.7
97.140.1
96.740.5
96.120.7
97.510.2
Object
71 71.4
88.911.0
19.1
79.32008
97.6x1.0 1.
98.1a
96.910.5
Long
45.9+1.6
82,7+0.2
88.7+0.1
82.7+0.5
89.5e0..4
88.2+016
87.3+0.6
Average
64.4
86.5
94,8
77.5
94.0
93.8
93.0
(RQ1) Does Cocos improve diffusion policies' training efficiency and final performance?
(RQ2) Does it mitigate the loss collapse phenomenon?
(RQ3) Does it facilitate policy learning in real-world settings across heterogeneous robot platforms?
5.1 Experimental Setup
Diffusion Policy. Our diffusion policy (DP) adopts compact RDT policy network of approximately 40M parameters. The vision-language condition inputs are encoded using a DINOv2- Base [20] and T5-Base [25]. We denote the resulting model DP-DINOv2. To evaluate the effectiveness of Cocos, we compare three key variants: (1) DP-scratch, which is trained without any pre-trained vision encoder; (2) DP-DINOv2, which incorporates frozen DINOv2 features but uses the standard source distribution; and (3) DP-DINOv2 with Cocos, our full method with a condition-condiiione source distribution. The autoencoder used to define the source distribution in Cocos is implemented using a single-layer Transformer for both encoder and decoder components.
Benchmarks. Our simulation evaluations based on the LIBERO and World benchmarks. LIBERO includes 40 tasks in four task suites: Goal, Spatial, Object, and Long, to test different policy generalizations. MetaWorld includes 30 tasks from various difficulty levels. For real- world experiments, we deploy the models on two robot platforms: The SO100 robot (low-cost, open-sourced, equipped with dual RGB cameras) evaluated 10 tasks in four suites: Pick& MoveTo, Wipe, and Unfold; The Arm robot (higher-precision, equipped with one Intel RealSense 1,515 LiDAR camera) evaluated 10 tasks in suites: Pick& Place, Pot, Pour, and Moka. We show detailed task contigurations in Appendix B.
5.2 Overall Comparison (RQ1)
We evaluate the effectiveness of Cocos on two widely-used simulation benchmarks: LIBERO and MetaWorld. For LIBERO. we also include comparisons with fine-tuned To [1]. Open [10], and Seer [29]. To and Open VFA-OFT are SOTA large-scale pre-trained VLA models, featuring flow-matching and next-token prediction, respectively. Seer is not an end-to-end VLA model but predicts future frames to infer actions. We report the results in Table Table 2 and Figure 4.
Cocos significantly improves training efficiency and policy performance. Across all LIBERO suites, Diffusion policy (DP) trained with Cocos outperforms its counterparts without Cocos. Notably, Cocos-enabled models reach high success rates (e.g., >95% on Spatial and Object suites) within only 5-10K training steps, achieving convergence roughly 2.14x faster than DP-DINOv2. On average, Cocos improves LIBERO success rates by 8.3%. Across all World tasks, also outperforms the counterpart by 25.7 % gain, showing consistent performance gains.
Finding 1:

---

Table 2: Success rate on World benchmark. We evaluate 150 trials for each task. The highest scores are emphasized in bold. We take multi-task setting, using language instructions to differentiate each task.
Nettongmes
battor-gmt-ttie
coffesbuatire
deoeegee
devemookk
deve-srieck
drsennooo
drvocrsopen
Average
Cooos
Senl
97.3025
98.3
467713
93,3641
108,000
7486
who Cocos
97.5225
31
99.3ut
97.3
77.3
3.3
67.3
100
108.2n0
59.5 59.525
scrach
48.7811
41.366.
32.7
$2.0a5.
53.3at
50.7467
23.322.5
24.7554
58.7
$7.345.
33.6454
fuxer close
Suued open
hanile press
powss
Wiedow close
Windew open
batde pull
bakkebill
Nepokay
Gaoss
82
108,8464
968633
27
9.543
64.062
who Cozes
NET-38
64 De4
S8
48.00 4
IRUISB
59.1
1:40
swrtch
36.7234
70, 99.3
760022
$3,3254
6077119
2505M0
20616
0.0000
10.0
poll
offiee yonhh
sooep
shelfepicce
stile push
Caoos
19.56622
60.765.
66.042.
767466
72
14,633
39:341 3414
94,0443
who Cozon
72 T
345.
34.0w133
72.741.
0.3425
200654
72.76LB
12 T.lA
smncch
100044
180053
STAL9
Carae Gonecken P9
120
75
85
100
47.1
20
20
Pick&PPlace (120 demos)
MoveTo (40 demos)
wipe (20 demos)
Unfold (20 demos)
50-100 Average
123
86.7
20
50
89.2
79.2
55
43.8
Pick&Place (120 demos)
Pot (40 demos)
Pour (20 demos]
Moka (20 demos)
xArm Average
DN-ORNON2 Cosos
Disorathh
Figure 5: Evaluation results on S0100 and xArm platforms. We collect 4 task suites or 10 tasks for each robot platform. Each task provides 20 demonstrations and tested over 10 trials.
Finding 2:
Cocos enables compact models to rival large-scale pre-trained VLAs. Despite using fewer parameters and fewer gradient steps, DP with Cocos performs comparably to large-scale models like Ty and Open VLA-OFT. As described in Kim et al. [10], while Open VLA-OFT takes 50- 150K gradient steps for onn-ttunig and To takes 100-250K steps, our model reaches similar performance within 30K steps. On the challenging Long suite, Cocos even outperforms Seer, non-end enn--to-end model with a larger parameter count. These results demonstrate that proper condition integration, rather than model scale alone, can be major driver of performance in VLA training.
Finding 3: Cocos encourages deeper condition utilization rather than just injecting priors. To understand the nature of the improvement, we examine internal representations of the policy network before and after condition injection (Figure 1). Specifically, we measure cosine similarity and norm scale changes of hidden states (see Appendix C for details). A lower similarity and a higher norm shift reflect greater responsiveness to conditioning inputs. As shown in Figure 1, these metrics are consistently aligned with higher task success rates. This suggests that the benefit of Cocos does not stem merely from injecting condition priors into the source distribution. Instead, it actively compels the policy network to differentiate between and respond to condition inputs. If the gains were due only to initialization, we would expect weaker condition dependence, contrary to what we observe.
5.3 Real- World Experiments (RQ2)
To assess the practicality of our approach, we deploy diffusion policies trained with and without Cocos on two real- world robotic platforms: S0100, low-cost and open-source arm, and xArm, widely used high-precision industrial robot. The two platforms offer complementary insights: SO100 lowers the entry barrier for research and reproducibility, while xArm reflects deployment scenarios with more demanding task complexity and control fidelity.
As shown in Figure 5 and Figure 6, Cocos consistently improves task success rates across all task suites on both robots. The improvement is particularly pronounced on SO100, where the relatively lightweight structure makes more sensitive to inaccuracies in action generation. We observe that both policy variants, with and without Cocos, are capable of generating smooth and coherent trajectories. However, the main failure modes of the baseline model are closely linked to misinterpretation of condition inputs, for example, confusing spatial references in language, or failing to correctly localize an object in the visual scene. contrast, these types of errors are reduced when using Cocos, indicating improved condition understanding and robustness. These trends align with our theoretical motivation: condition misalignment during training can lead to loss collapse and unreliable behavior at inference time. The fact that Cocos provides consistent improvements across two distinct robotic platforms further highlights its practical value. In the next section, we present detailed case studies to examine these qualitative differences in greater depth.
7

---

Figure 6: Rollout samples of DP with Cocos real-robot manipulation, including MoveTo (Bowl) (Top
Left), Wipe (Top Right), Moka (Bottom Left), and Pour (Bottom Right).
"Open pot
"Pick the sponge and piace it on the right side of the table."
"Open the pol lid.'
"Pick the and the side table."
Figure
Real robot case study. Top left rollout comes from policy w/ Cocos and the others are from baselines.
5.4
Case Studies (RQ3)
To better understand the mechanisms behind Cocos's effectiveness, we conduct qualitative case studies on both simulation and real-world tasks, focusing on common failure patterns and condition sensitivity.
Finding 1: Cocos improves utilization of 3rd-person visual inputs in simulation. In the LIBERO benchmark, we find that policies trained without Cocos tend omit 3rd-person views, relying almost exclusively on the wrist-mounte between randomly sampled visual features extracted before fusion into the policy network. As shown in Figure 8 bottom table, models without Cocos produce highly similar 3rd-person embeddings across diverse scenes, indicating weak discriminability, while Cocos-trained models show significantly greater variation, suggesting stronger visual grounding.
Ist-person camera. To quantify this, we compute cosine similarity
This contrast is clearly illustrated through the rollout comparisons in Figure 8. (Row 1) shows successful execution by the model trained with Cocos: the robot pushes the plate forward, turns smoothly to the left, and finally pushes the plate backward the stove, demonstrating accurate spatial reasoning and multi-step planning. This trajectory serves as reference for the expected behavior. In (Row 2), a failure case from the model without Cocos, the robot accidentally releases the plate mid-way and attempts to recover by locating it again. However, since the plate is no longer visible in the Ist-person view, the policy fails to re-establish contact, even though the 3rd-person view clearly shows the plate's 's new location. This indicates that the model underutilizes the auxiliary camera input. (Row 3) presents similar failure condition, but the plate remains within the first-person field of view after being released. In this case, the robot successfully relocates and completes the task. Comparing Row 2 and Row highlights the model's 's strong reliance on the 1st 1s--person view: the object leaves that view, recovery becomes unlikely. (Row 4) shows another typical failure from the non-Cocos model: after reaching the stove, the robot continues to push the plate beyond the goal region. This happens because the ss--person view offers no clear signal indicating proximity to the stove, while the 3rd-person view does. Again, the model fails to leverage this available information. Taken together, these cases reveal a consistent pattern: policies trained without Cocos tend to over-rely on the prist-mounted camera and struggle to utilize 3rd-person observations. contrast, models trained with Cocos seldom exhibit these failures, suggesting that the policy network learns to better incorporate multi-vien context and more robust to visual uncertainty.
Finding 2: Cocos reduces modality over-reliance and improves contextual awareness in real- world settings. real-robot experiments, we observe that models without Cocos frequently exhibit modality over-reliance, particularly favoring language instructions or arm states over visual cues. As shown Figure? when tasks begin an already-complete configuration (e.g., lid already open), the baseline policy often executes full rollout regardless of the visual scene (left bottom). This indicates

---

that the model has learned to follow the instruction without verifying whether it is still applicable. By contrast, policy trained with Cocos correctly remains idle such states, indicating improved contextual awareness (left top). Similar trends appear in other tasks such as top) and ScissorMoveRight (right bottom), where the baseline struggles locate target objects or misinterprets spatial references. These emoos are less frequent with Cocos, which appears to enhance the policy's sensitivity to both visual and semantic nuances in the scene. Together, these case studies provide concrete support for our theoretical claims: Cocos prevents loss collapse not just in training metries, but how models behave under ambiguous or weakly conditioned scenarios.
(right
5.5 Ablation Studies
Cocos only requires the source distribution to be condition-dependent, i.c., (())c), and admits flexibility in its concrete form. For simplicity, we instantiate it as Gaussian with fixed stan- dard deviation and condition-dependen mean (Equation (8)). Our default setting uses 3 0.2, and we compare it against to investigate the impact of the standard devia- tion. Inspired by the similarity between our for- mulation and the variational autoencoder (VAE) posterior, we evaluate VAE-based version of Cocos, where the standard deviation is auto- matically learned from data. As shown in Ta- ble I, we observe that the default (B 0.2) and the wider variant (3 = 0.4), and the VAE ver- sion all achieve comparable performance across LIBERO task suites. However, the model with 0.1 performs substantially worse, suggest- ing that overly concentrated ((=cc) hinders train- ing. Overall, these results indicate that Cocos is robust to the choice of 33, as long as g(2[c) remains sufficiently spread.
0.1 and 8 0.4
X
Average feature cosine similarities (L)
view
w/ Cocos
w/o Cocos
scratch
We also compare the default two-stage training approach, where the source distribution is pre- trained and then fixed, with joint training strat- egy using exponential moving average (EMA) Figure 8: Case study of LIBERO plate movement. updates. In this version, a slowly updated target encoder Fø- is used during policy training to provide stable condition-dependen sampling. This online approach achieves performance on par with the two-stage method while simplifying the pipeline, making it practical choice for integrated training scenarios.
3r-person
0.514
0.843
0.971
p.lpeeson
0.206
0.324
0.820
Conclusion, Limitations, and Future Works
In this paper, identify and address a critical failure mode in diffusion policy training, loss collapse, where the model fails to distinguish between generative conditions and degenerates into modeling marginal action distributions. We propose Cocos, a simple yet effective modification that injects condition-awarenes. directly into the source distribution. We provide a theoretical analysis showing how Cocos prevents loss collapse and demonstrate its empirical benefits across extensive simulations and real-world tasks. Our results show that Cocos improves both training efficiency and policy performance, enabling even compact models match or exceed the performance of large-scale pretrained VLA systems. Despite its effectiveness, Cocos leaves open several important directions. First, while the method only requires the source distribution be condition-depenie have so far instantiated it using a fixed-variance Gaussian derived from autoencoding the condition. Exploring more expressive or learnable source distributions, such as flow-based or attention attention-guided priors, may further enhance its adaptability. Second, our current experiments are focused on multi-task imitation learning. Extending Cocos to large-scale pretraining and evaluating its ability to generalize across broad VLA domains remains promising area for future work.
(q(=lc)), we

---

Acknowledgments
This work is supported by the National Natural Science Foundation of China (Grant Nos. 62422605. 92370132), We would like to thank all the anonymous reviewers for their valuable comments and constructive suggestions, which have greatly improved the quality of this paper.
References
[1] Kevin Black, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Lachy Groom, Karol Hausman, Brian lchter, Szymon Jakubczak, Tim Jones, Liyiming Ke, Sergey Levine, Adrian Li-Bell, Mohith Mothukuri, Suraj Nair, Karl Pertsch, Lucy Xiaoyang Shi, James Tanner, Quan Vuong, Anna Walling, Haohuan Wang, and Ury Zhilinsky. Ty: A vision- language-action flow model for general robot control. In arXiv preprint 2Xiv:2410.241644 2024.
[2] Remi Cadene, Simon Alibert, Alexander Soare, Quentin Gallouedec, Adil Zouitine, and Thomas Wolf. Lerobot: State-of-the- machine learning for real-world robotics in pytorch. https: //github. httpll/goooogggneeeeeeeeeeeee
2024.
[3] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems, RSS, 2023.
[4] Eugenio Chisari, Nick Heppert, Max Argus, Tim Welschehold, Thomas Brox, and Abhinav Val- ada. Learning robotic manipulation policies from point clouds with conditional flow matching. In 8th Annual Conference on Robot Learning. CoRL, 2024.
[5] Zibin Dong, Yifu Yuan, Jianye HAO, Fei Ni, Yi Ma, Pengyi Li, and YAN ZHENG. Cleandiffuser: An easy-to-use modularized library for diffusion models in decision making. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, NIPS, 2024.
[6] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, NIPS, 2020.
[7] Wenlong Huang, Chen Wang, Yunzhu Li, Ruohan Zhang, and Li Fei-Fei. Rekep: Spatio- temporal reasoning of relational keypoint constraints for robotic manipulation. In 8th Annual Conference on Robot Learning, CoRL, 2024.
[8] Yuanchen Ju, Kaizhe Hu, Guowei Zhang, Gu Zhang, Mingrun Jiang, and Huazhe Xu. Robo- abc: Affordance generalization beyond categories via semantic correspondence for robot manipulation. In European Conference on Computer Vision, ECCV, 2025.
[9] Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan P Foster, Pannag R Sanketi, Quan Vuong, Thomas Kollar, Benjamin Burchfiel, Russ Tedrake, Dorsa Sadigh, Sergey Levine, Percy Liang, and Chelsea Finn. Open- VLA: An open-source vision-language-acttee model. In 8th Annual Conference on Robot Learning, CoRL, 2024.
[10] Moo Jin Kim, Chelsea Finn, and Percy Liang. Fine-tuning assion-language-acttoo models: Optimizing speed and success. In arXiv preprint (rXXi:25022.16455. 2025.
[11] Yinchuan Li, Xinyu Shao, Jianping Zhang, Haozhi Wang, Leo Maxime Brunswic, Kaiwen Zhou, Jiqian Dong, Kaiyang Guo, Xiu Li, Zhitang Chen, Jun Wang, and Jianye Hao. Generative models in decision making: A survey. In arXiv preprint 250.::25022.7000 2025.
[12) Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, ICLR, 2023.
[13] Bo Liu, Yifeng Zhu, Chongkai Gao, Yihao Feng, qiang liu, Yuke Zhu, and Peter Stone. LIBERO: Benchmarking knowledge transfer for lifelong robot learning. Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, NIPS, 2023,
10

---

[14] Jiaming Liu, Hao Chen, Pengju An. Zhuoyang Liu, Renrui Zhang, Chenyang Gu, Xiaoqi Li, Ziyu Guo, Sixiang Chen, Mengzhen Liu, Chengkai Hou, Mengdi Zhao, KC alex Zhou, Pheng- Ann Heng, and Shanghang Zhang. Hybridvla: Collaborative diffusion and autoregression in unified vision-language-actinnn acction model. In arXiv preprint arXiv:2503. 10631, 2025.
[15] Songming Liu, Lingxuan Wu, Bangguo Li, Hengkai Tan, Huayu Chen, Zhengyi Wang, Ke Xu, Hang Su, and Jun Zhu. RDT-1b: diffusion foundation model for bimanual manipulation. In The Thirteenth International Conference on Learning Representations, ICLR, 2025.
[16] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, ICLR, 2023.
[17] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps. In Advances in Neural Information Processing Systems, NIPS, 2022.
[18] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm- solver++: Fast solver for guided sampling of diffusion probabilistic models. In arXiv preprint arXiv:2211.01095, 2022.
[19] Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Charles Xu, Jianlan Luo, Tobias Kreiman, You Liang Tan, Pannag Sanketi, Quan Vuong, Ted Xiao, Dorsa Sadigh, Chelsea Finn, and Sergey Levine. Octo: An open-source generalist robot policy. In Proceedings of Robotics: Science and Systems, RSS, 2024.
[20] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V. Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel HAZIZA, Francisco Massa, Alaaeldin El-Nouby, Mido Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Po-Yao Huang, Shang-Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Herve Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. DINOv2: Learning robust visual features without supervision. Transactions on Machine Learning Research, TMLR, 2024.
[21] William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022.
[22] Karl Pertsch, Kyle Stachowicz, Brian Ichter, Danny Driess, Suraj Nair, Quan Vuong, Oier Mees, Chelsea Finn, and Sergey Levine. Fast: Efficient action tokenization for vision-language-action models. In arXiv preprini arXiv:2501.099 .09747, 2025.
[23] Zekun Qi, Wenyao Zhang, Yufei Ding, Runpei Dong, Xinqiang Yu, Jingwen Li, Lingyun Xu, Baoyu Li, Xialin He, Guofan Fan, Jiazhao Zhang, Jiawei He, Jiayuan Gu, Xin Jin, Kaisheng Ma, Zhizheng Zhang, He Wang, and Li Yi. Sofar: language-grounde orientation bridges spatial reasoning and object manipulation. In arXiv preprint arXiv:2503. 10631, 2025.
[24] Delin Qu, Haoming Song, Qizhi Chen, Yuanqi Yao, Xinyi Ye, Yan Ding, Zhigang Wang, Jia Yuan Gu, Bin Zhao, Dong Wang, and Xuelong Li. Spatialvla: Exploring spatial representations for 2025.
visual-language-attoon model. In arXiv preprint arXiv:2501.
[25] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lec, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, JMLR, 2020.
[26] Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, ICLR, 2021.
[27] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. arXiv preprint urXiv:1907.05660 2020.
[28] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, ICLR, 2021.

---

[29] Yang Tian, Sizhe Yang, Jia Zeng, Ping Wang, Dahua Lin, Hao Dong, and Jiangmiao Pang. Predictive inverse dynamics models are scalable learners for robotic manipulation. In The Thirteenth International Conference on Learning Representations, ICLR, 2025.
[30] Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based genera- tive models with minibatch optimal transport. Transactions Machine Learning Research, TMLR. 2024.
[31] Zihao Wang, Shaofei Cai, Zhancun Mu, Haowei Lin, Ceyao Zhang, Xuejie Liu, Qing Li, Anji Liu, Xiaojian Ma, and Yitao Liang. OmniJARVIS: Unified iisionlanguaggeacttinn tokenization enables open-world instruction following agents. In The Thirly-eighth Annual Conference on Neural Information Processing Systems, NIPS, 2024.
[32] Junjie Wen, Minjie Zhu, Yichen Zhu, Zhibin Tang, Jinming Li, Zhongyi Zhou, Chengmeng Li, Xiaoyu Liu, Yaxin Peng, Chaomin Shen, and Feifei Feng. Diffusion-vla Scaling robot foundation models via unified diffusion and autoregression. In arXiv preprint arXiv:2412.03295 2024.
[33] Junjie Wen, Yichen Zhu, Jinming Li, Minjie Zhu, Kun Wu, Zhiyuan Xu, Ning Liu, Ran Cheng, Chaomin Shen, Yaxin Peng, et al. Tinyvla: Towards fast, data-efficient vision- language- models for robotic manipulation. In IEEE Robotics and Automation Letters, RAL, 2025.
[34] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference Robot Learning, CoRL, 2019.
[35] Wentao Yuan, Jiafei Duan, Valts Blukis, Wilbert Pumacay, Ranjay Krishna, Adithyavairavan Murali, Arsalan Mousavian, and Dieter Fox. Robopoint: A vision-language model for spatial affordance prediction robotics. In Proceedings of The 8th Conference on Robot Learning. CoRL, 2025.
[36] Yanjie Ze, Gu Zhang, Kangning Zhang, Chenyuan Hu, Muhan Wang, and Huazhe Xu. 3d diffusion policy: Generalizable visuomotor policy learning via simple 3d representations. In Proceedings of Robotics: Science and Systems, RSS, 2024.
[37] Fan Zhang and Michael Gienger. Affordance-basee robot manipulation with flow matching. In arXiv preprint arXiv:2409.0 2025.
[38] Qinglun Zhang, Zhen Liu, Haoqiang Fan, Guanghui Liu, Bing Zeng, and Shuaicheng Liu. Flowpolicy: Enabling fast and robust 3d flow-based policy via consistency flow matching for robot manipulation. In arXiv preprint arXiv:2412.049882 2024.
12

---

A Proofs of Theorems
Proof of Lemma 1. As in Tong et al. [30] we assume that q. (2|22 are decreasing to zero at sufficient
speed as ||=|| and that Me, Ut, Ver, are bounded.
V pe 00
(x)c))) VoBpi(a1e) (Ioo(t,,,,))1 c)l2 - 2 (walt, 2, c), ux(z)))) + [
-ooBpittte) (1e0(t,2 e)ll - (vo (t) (11)
V (c (x VoBagsrnntalzo (IIvo(t, x, e) - (wo x, c), ue(a/2, e)) + [1((()2,,))
= VoBB (100 e)(l
(vo(t, (f (((k),))))
(12)
By bilinearity of the 2-norm and since is independent of 0. Next,
(pt() ||oo(,, xr, c)ll²= = / (x
(13)
// [Ivo(t,2,)) plal...)cc)
(14)
(15)
Finally,
Epp(zle) (vall, x,, c), ux(x)c)) = (wo(l, c), me(xlc)) pa (xlc)dr
(16)
vo 10(1, c), Fate) (1 pi(cle)dr
(17)
(a(e)e)
vo(t, c), / u || (va(L, pa..: (x)s)) (v0(t, x, c), (x
dædæ
(18)
(19)
(20)
Since the gradients of the two objectives are equal for any time f, the two objectives are equal.
L
Proof of Theorem I. Define
d(t, x, c) vo xz 20) f(z,c) Vouo(t,2, d(t,,,,),
(21)
where we write = (Io, 2), Introduce two measures:
Independent measure
(ddy) ((., e) g(xr) Pe (I)a), dxy dæy dx,
(22)
which does not depend on the particular choice of c, corresponding to the design choice where q(2) g( (z))
Conditional measure
(dy) = (( c) (Folc) Pe (a mо) dxy da dx,
(23)
which does depend on corresponding to the design choice where g(z) g( c)g(aee)...
13

---

We assume the following bounds hold for all and c E.C.
Voualt, 2, c) S M, 2 S SD,
(24)
and that the model's gradient is outpul-sensitive Lipschitz, namely there exists K > such that
[t, Vo I K [o,ott, xolt,,, - Ke.
(25)
ii) Independent measure. Under 14, the gradient
VeLcFMo(O,c)=2) f(2,c)
(26)
Hence for any Ct, c2;
VoLoFMe(O,c)) (Vocc(.l))))
(27)
= =21/ f(2,ai)((d)) / ((2,c))
(28)
-2 ==(/(((t.a))))))))))){ -(e, u(dy)
(29)
Expanding the integrand,
f(2,c1) f(2,02) vj [f 2(t,
(30)
= ((llt.)) d(t,
(31)
+ -
- Vovs((, d(l, 21).
(32)
Using the bounds [d(()) - < and (ou(c)) Vou(o)) S KE we obtain
/ (MMe-Ke))
(oL) - VoLcrMe(0, 02) 2 / (Mc+ KcD) u(dy)
(33)
=(
(34)
The above bound shows that when all C induce similar gradients, the network update over this region becomes nearly invariant to c. In the limit of gradient descent, the parameter thus evolves according to a shared direction, regardless of the exact value. As result, the learned model v0(1, x, e) tends to converge, for all e E C, to c-independent function v" (1, a)) that minimizes the averaged objective:
0 (f,z) arg min Beec [llu(t, 2) - (2)22
(35)
That is, the network effectively collapses the conditional dependency and behaves like shared estimator v" optimized for average performance across C.
(ii) Conditional measure. Under Po: the gradient becomes
(0,)) = / f(2,c) Poly).
(36)
For C2, write VoLCrMe(O,c)) VolcrMc(O,C2)
(37)
=2[/ 5(2,c1) (dz) - / f(2,02)
(38)
= 2 =2[[fstzcii ((E,(ti)) (dz) - fle,ci) + If f(z,ci) / f(z, 02) Mod (d2))
(39)
= T, + Tz.
(40)
14

---

The second term T2 coincides with the independent-measury analysis and thus is bounded by 2(M KD) e. The first term
Ti - ((,c)) == Pei (dz)
[(2 (oa
(41)
is the difference of the same function under two different measures, Without any further assumption on the relationship between Me, and Peg this difference can be made arbitrarily large. For instance, ((r))e1) = 6, and g(0o|c2) = 66. then l|Till = 2a c1) f(b, ci)l can diverge independently of Hence, under the conditional measure Pes no bound depending only on exists for VVL(0,, c1) V (8, c2)ll.
B Details of Benchmarks
Figure 9: LIBERO simulation benchmark. We conduct experiments on 40 tasks from four task suites in the LIBERO benchmark. We show task examples for each suitc here.
Figure 10: MetaWorke simulation benchmark. We conduct experiments on 30 tasks of three difficulty levels in the MetaWork World benchmark. We show all task examples here.
Camaro
Sponge
Hand- vlew
Charger
SeaiSenc LS15
Scibsor/PPae: Plan//ap
Real-worke world experimental setups. We conduct experiments on both SO100 and xArm platform. For each robot, we design a suite of 10 tabletop tasks involving diverse objects.
Figure 11:
LIBERO. The LIBERO simulation benchmark [13] features a Franka Emika Panda arm in simulation across four challenging task suites: Goal, Spatial, Object, and Long. Each suite comprises 10 tasks with 500 demonstrations and is designed to investigate controlled knowledge transfer related to goal variations, spatial configurations, object types, and long- horizon tasks. Unlike prior work [9, 10], we do not filter out unsuccessful demonstrations, aiming for more realistic evaluation setting, For policy training, the model predicts action chunks of length 16; after each chunk prediction, steps are executed before generating the next chunk. The observation space includes 2-view RGB images the current time step, without historical observations, During evaluation, following Liu et al. [13], each task is tested over 50 trials with different random seeds, and success rates reported. To
15

---

provide a clearer understanding of the task suites, we present agent-view observations in Figure 9 and detailed task descriptions in Table 3.
World. The MetaWorld simulation benchmark [34] includes 50 distinct tabletop manipulation tasks using Sawyer robot arm. We select 30 tasks from easy, medium, and very hard difficulty levels to evaluate VLA models. We use a scripted policy to collect 20 demonstrations for each task. For policy training, the model predicts action chunks of length 16; after each chunk prediction, 16 steps are executed before generating the next chunk. The observation space consists of a single RGB image at the current time step. without historical observations. During evaluation, each task is tested over 50 trials with 3 different random seeds, and success rates are reported. To better illustrate the task suites, we show agent- observations in Figure 10 and task descriptions in Table 4.
SO100 Robot Manipulation. The S0100 robot [2] is a low-cost, open-source 6-DoF manipulator, with both the leader and follower arms costing approximately $250. We assemble the hardware using a 3D-printed kit provided by the open-source community. The robot has two RGB cameras: one mounted on the wrist and the other positioned provide a third parrd-p view. Both cameras operate at a resolution of 640x480 and 25 FPS. The robot controller runs at 30Hz, and actions are defined as target absolute joint angles. Due its low-cost design, the platform has several hardware limitations, including significant arm jitter, low load capacity, and occasional camera lag, which present practical challenges for developing embodied AI systems. However, given the increasing adoption of such affordable open-source robots by the research community, we believe that evaluating models on these lower-performance systems offers valuable insights and broader applicability. We design four categories of tabletop manipulation tasks for the $0100 setup: (I) Tick&Placee involving 3 objects and placement zones (6 tasks). (2) MoveTo: navigating 2 objects to a single target zone (2 tasks), (3) Wipe: picking up cloth and wiping the table (1 task), and (4) Unfold: unfolding a cloth (1 task). In total, we evaluate performance on 10 distinct tasks. Language instructions for each task are listed in Table 5, and visual examples of the task environments are shown in Figure 11.
During data collection, we record 20 demonstrations per task. For policy training. the model predicts an action chunk of length 64; after each chunk prediction, 40 steps are executed before generating the next chunk. The observation space includes 2-view RGB images at the current time step, along with the absolute joint angles from the current and previous 10 steps. During evaluation, each task is tested over 10 trials, and success rates are reported.
xArm Robot Manipulation. xArm is a high-peeformance 7-DoF manipulator. The robot is equipped with a third-person view Intel RealSense L515 LiDAR camera, operating at 640x480 resolution and 30 FPS. We collect both RGB and depth images from the camera. The robot controller runs at 30Hz, and actions are defined as target absolute joint angles. We design four categories of tabletop manipulation tasks for the XArm setup: (1) Pick&Place: involving 3 objects and 2 placement zones (6 tasks), (2) Pot: taking off or putting on the pot lid (2 tasks). (3) Pour: pouring water from the kettle into the cup (1 task), and (4) Moka: placing the Moka pot on the cooker (1 task). In total, we evaluate performance on 10 distinct tasks. Language instructions for each task are listed in Table 5, and visual examples of task environments are shown in Figure 11.
During data collection, we record 20 demonstrations per task. For policy training, the model predicts an action chunk of length 64; after each chunk prediction, 40 steps are executed before generating the next chunk. The observation space includes a third-person view RGB image at the current time step, as well as the absolute joint angles from the current and previous 10 steps. During evaluation, each task tested over 10 trials, and success rates are reported.
C Details of Evaluation Metrics
Cosine similarity. To quantify the influence of condition inputs on our policy network, we measure the cosine similarity between hidden states before and after condition injection. Given an T--layer RDT policy network where conditions are integrated via cross-attention in each transformer block, we denote the hidden states before and after condition injection as h', ht' gSDD where is the sequence length and D is the transformer hidden dimension. Since earlier layers' transformations propagate through the network, affecting all subsequent computations, we calculate an exponentially weighted sum: E w/ mins siinss(Sim (Sim(h' (h', ]d)), where ming operates across the sequence dimension, Sim computes cosine similarity across the hidden dimension, and w 0.5 gives greater weight to
16

---

Table 3: Task description of each task in the LIBERO benchmark.
Task Suite
Task Description
open the middlle layer the put the howl the stove
put the wine bottle the top of the drawer
open the tup layer of the drawer and pat the bowl inside put the howl on the top of the drawer push the plate to the front of stove put cream cheese the bowl turn the
LIBERO-Goal
put the howl on the plate put the wine bottle on the rack
pick the akita black howl between the plate and the ramekin and place it on the plate pick the akita black bowl next to the ramekin and place it on the plate
pick akita black bowl from table center and place it on the plate black
pick
cnokies box and place on the plate
pick the akita black bowl in the top layer of the wooden cabinet and place it plate pick the akita black bowl the rumekin and place it on the plate
LIBERO-Spatial
pick the akita black bowl next to the cookies box and place on the plate
pick the akita black bowl on the stove and place it on the plate
pick akita black next to the plate and place it on the plate
pick akita black bowl
wooden cabinet and place the plate
pick alphabet soup place it in the basket pick cream cheese and place pick the salad dressing and place it in the basket pick the biq sauce and place it the basket
baskct
pick the ketchup and place the basket
LIBERO-Object
pick tomato sauce and place the basket pick butter and place pick the milk and place
basket
the basket
pick the chocolate pudding and place it in the basket pick the orange juice and place in the busket
put both the alphabet soup and the tomato sauce in the basket put both the cream cheese box and the butter in the basket turn on the stove and put the moka pot on it put black bowl in the bottom drawer of the cabinet and close put white mug on the left plate and put the yellow and white mug on the right plate
LIBERO-Long
pick up the book and place in the back compartment of the eaddy
put the white mug on the plate and put the chocolate pudding to the right of the plate put hoth the alphahet soup and the cream cheese box in the hasket put both moka pots on the stove put yellow and white mug microwave and
earlier layers. Lower similarity values indicate stronger condition influence on network representa- tions.
Norm scale change. Similarly, we quantify the magnitude of change induced by conditional inputs using exponentially weighted of relative norm changes: DL w' maxs R
|
where
maxş operates across the sequence dimension, Ill computes the norm across the hidden dimension. and w 0.5. Larger values indicate that conditional information substantially alters the magnitude of hidden representations, suggesting stronger condition integration within the policy network.
Cosine similarity in case study. Figure 8, we also calculate cosine similarity metric to quantify the degree to which the policy network differentiates between conditions. This cosine similarity computed between visual features before they are fused into the policy network (i.e., after the vision encoder's output has been processed by positional encoding and an MLP, yielding the features ultimately used for fusion with the policy network). We calculate this metric by randomly sampling image pairs from the dataset and then averaging the cosine similarity of their corresponding features.
17

---

Table 4: Task description of each task in the MetaWorld benchmark.
Task Name
Task Description
basketball
Dunk the basketball the basket.
bin-picking
Grasp the puck from one hin and place it into another hin.
butlon-press
Press a bullon.
button-press- patton-pres-topdownn
Press a button from the top.
button-press press-topdown-wat
Bypass a wall and press hutton from the top.
putlon-presswwll
Bypass a and press button.
coffee-button
Push button the coffee machine.
coffee-pull
Pull mug from coffee machine.
coffee-push
Push mug under a coffee machine.
dial-turn
Rotale dial 180 degrees.
disassemble
Pick nut out the peg.
door-lock
Lock the door by rotating the lock clockwise.
door-open
Open door with revolving juint.
door-unlock
Unlock the door by rotating the lock counter-clockwise. Push and close a drawer.
drawer-close
drawer-open
Open drawer. Rotate faucet clockwise.
faucet-close
faucet-open
Rotate the faucet counter-clockwise.
hammer
Hammer screw on the wall.
handle-press
Press a handle down.
handlc-press-sieee
Press a handle down sideways. a handle up.
handic-pull
handlc-pull-side
Pull handle up sidcways.
shelf-pha place
Pick and place puck onto shelf.
soccer
Kick soccer into the goal.
stick-push
Grasp a stick and push a box using the stick.
sweep
Sweep puck off table.
sweep-inlo
Sweep puck into hole.
window-close
Push and close window.
window-open
Push and open window.
Table 5: Task description cach task the SO100 and xArm benchmark, As each parameter combination introduces one task, each task suite has 10 tasks in total. For each task, we test the model for 10 trials.
Task Suite
Task Description
Parameter
pick place
side
[A]:
["scresdriver", "eponge", "charger"), ["lext", "righe")
move [A] to the center of the table pick
[A]: ['cup", "bos1"]
SO100
wipe table
Hone
unfold the cloth
Mone
pick [A] place
[A]: ["scissor", "plier", "tap"), [B] ["lete", "vight"]
the pot lid oz put the lid on the put pour the
[open, close]
xAm
cup
place the Moka pot on the cooker
Hone
D Policy Backbone Recipe
D.1 Condition Network
For vision encoding, utilize the pre-trained DINOv2-Base model, and for language encoding, employ the pre-trained T5-Base model. Both of these encoders are frozen during training. The outputs from these pre-trained modality encoders processed by linear projector and then augmented with positional encodings before being fused into the policy network. These positional encodings are trainable but are initialized with sin-cos positional encoding and scaled down by factor of 0.2.
D.2 Policy Network
We employ a compact Robot Diffusion Transformer (RDT) architecture, which fundamentally Diffusion Transformer incorporating cross-attention layers. Diffusion timestamps and robot kinematic information are integrated into the policy network using Adero [21]. The vision and language embeddings are used as the Keys and Values in the cross-attention layers to be integrated into the policy network alternately [15]. The Transformer architecture has hidden dimension of 384, with attention heads, and 12 layers.
18

---

D.3 Diffusion Model
We use a flow matching model defined by Equation (10). Diffusion timestamps are treated as continuous values within the range [0, 1]; we do not discretize them. Instead, they are represented using Fourier embedding with scale of 0.2 [5]. During training, diffusion timestamps are sampled from uniform distribution over the interval [0, 1|. For inference, we solve the corresponding ODE using the Euler method, dividing the interval [0, into equal-sized steps.
D.4 Computation Resources
All policy training and testing are conducted on a server equipped with 4 NVIDIA GeForce RTX 4090 GPUs and an Intel(R) (Ro CPU E5-2678 v3 @ 2.50GHz.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations of the work in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
20

---

Justification:
We provide the full set of assumptions and a complete proof for each theoretical
result.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
We provide every detail to reproduce the main experimental results.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [Yes]
Justification:
We will open-source the code in few days after submission.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.o cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ //nips. wttpccccpppccccceeececccccccccct httpa//ggiiieee/eeeuuuuiiiiiii
/CodeSubmie wwwinnoolll. onPolicy))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide every detail in training and testing.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Yes. We report error bars in the learning curve.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
22

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes. We list the GPU and CPU resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that NeurIPS Code of Ethics https: /ttp///eers/.////ee//ee//e cc httpppbliit/thicccuueeeeeee.eeel
didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification: Yes, do.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes, do.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
23

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Yes, they are.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
24

---

Answer: [Yes]
Justification: Yes.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do not include crowdsourcing experiments or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not include crowdsourcing experiments research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
25

---

Justification:
The core method development in this research does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to LLM policy (https:/ for what should or should not be described.
: /ttrr/ppp..c//ooofeeeeee///////2
26


Reviews:
[{"review_id": "BuugY7f6Qt", "reviewer": "Reviewer_rXXE", "summary": "This 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.", "questions": "Can 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?", "limitations": "Yes, the authors adequately describe limitations.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \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?)", "was_revised": true, "final_justification": "This is a very nice contribution that elucidates the effect of coupled priors on better action generation. I think this will be a useful contribution to the literature.", "rebuttal": "> Question 1: Improve Theorem quality\n\nWE 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.\n\n---\n---\n---\n\n> Question 2: Design choice about $\\mu(x_0|c)$\n\nOur 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.\n\nWe 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.\n\nHere are the results of this experiment:\n\n|                  | 5k          | 10k         | 20k         | 30k         |\n| :--------------- | :---------- | :---------- | :---------- | :---------- |\n| cocos (autoencoding) | 83.1 ± 0.9  | 81.1 ± 1.3  | 92.2 ± 1.3  | 93.8 ± 0.3  |\n| cocos (action prior) | 70.0 ± 1.8  | 81.0 ± 1.0  | 84.5 ± 0.6  | 87.2 ± 0.6  |\n| vanilla (w/o cocos)  | 66.5 ± 1.6  | 78.3 ± 0.5  | 83.7 ± 0.3  | 82.9 ± 0.4  |\n\nWe 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$.\n\n---\n---\n---\n\nThank 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.**"}, {"review_id": "1zcQIx0er8", "reviewer": "Reviewer_LRqf", "summary": "This 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.", "questions": "- 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)?", "limitations": "The paper briefly discusses limitations regarding the design of the condition-aware prior (e.g., fixed variance Gaussian) and the scope of current evaluations.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "Interesting analysis and neat solution.", "rebuttal": "> Question 1: Lack of Clarity in Analysis Metrics\n\nThanks 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.\n\n---\n---\n---\n\n> Question 2: How does the method prevent loss collapse, or does it only reduce the chance of collapse?\n\nWe 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.\n\n---\n---\n---\n\n> Question 3: Potential Limitations in Diversity and Generalization\n\nWe 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.\n\n---\n---\n---\n\n> Question 4: Questionable Semantic Prior Assumption\n\nAs 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:\n\nWe 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.\n\nHere are the results of this experiment:\n\n|                  | 5k          | 10k         | 20k         | 30k         |\n| :--------------- | :---------- | :---------- | :---------- | :---------- |\n| cocos (autoencoding) | 83.1 ± 0.9  | 81.1 ± 1.3  | 92.2 ± 1.3  | 93.8 ± 0.3  |\n| cocos (action prior) | 70.0 ± 1.8  | 81.0 ± 1.0  | 84.5 ± 0.6  | 87.2 ± 0.6  |\n| vanilla (w/o cocos)  | 66.5 ± 1.6  | 78.3 ± 0.5  | 83.7 ± 0.3  | 82.9 ± 0.4  |\n\n\nWe 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.\n\n> In some settings, similar instructions can require significantly different actions due to scene geometry or context.\n\nWe 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.\n\n---\n---\n---\n\nThank 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.**"}, {"review_id": "7Br1PIh7Xv", "reviewer": "Reviewer_JmtE", "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.", "questions": "1. Does the performance improvement seen with Cocos scale with model size?\n2. 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.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. The method is simple, yet effective. \n2. Cocos improves over not using it in many experiments.\n\nWeakness:\n1. The theoretical section is a bit hard to parse and understand. I would encourage the authors to simplify the writing a bit.\n2. 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.\n3. 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.\n4. 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?", "was_revised": true, "final_justification": "I thank the authors for all the additional experiments in the rebuttal. It addressed many of my concerns, and hence I upgrade my rating to borderline accept.", "rebuttal": "> Question 1: Simplify the Writing of Theorem\n\nThanks 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.\n\n---\n---\n---\n\n> Question 2: Adding More Experiments on Datasets like Robomimic\n\nWe 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:\n\n| | Square (ph) | Transport (ph) | ToolHang (ph) |\n| --- | --- | --- | --- |\n| w/ Cocos | 1.0 | 1.0 | 0.98 |\n| w/o Cocos | 1.0 | 0.98 | 0.98 |\n*Note: Since Robomimic tasks lack language instructions, we removed the T5 encoder, and the policy network only accepted vision embeddings. All other settings remained consistent with Libero.*\n\n---\n---\n---\n\n> Question 3: Cocos Scale with Model Size\n\nThat'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:\n\n| | Goal | Spatial | Object | Long |\n| --- | --- | --- | --- | --- |\n| w/ Cocos | 86.2 | 88.0 | 93.4 | 72.5 |\n| w/o Cocos | 80.4 | 82.7 | 91.6 | 61.2 |\n*Note: The \"w/o Cocos\" results show some discrepancy from the official OpenPI report. This is mainly because the limited time and computational resources during rebuttal forced us to adopt a lightweight fine-tuning strategy: **we froze PaliGemma, only trained the action expert, used bf16 precision**, a batch size of 32, and did not use PI0's pre-trained stats.*\n\nThe results show that Cocos can indeed bring benefits when fine-tuning large-sized VLA models.\n\n---\n---\n---\n\n> Question 4: The Method Still Lacks Behind Methods like Seer\n\nThe 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.\n\nOur 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.\n\n---\n---\n---\n\n> Question 5: Can this method be used for general tasks like text to image generation?\n\nTheoretically, 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.\n\nAIGC 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.\n\n---\n---\n---\n\nThank 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.**"}, {"review_id": "YboHtQEB1x", "reviewer": "Reviewer_2fVZ", "summary": "The 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.", "questions": "- 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", "limitations": "Limitations are appropriately discussed.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\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.", "was_revised": true, "final_justification": "Its a good paper with nice theoretical insights that uncovers a key problem in conditioning of diffusion policies. Results are also convincing and the authors could address my minor concerns in ther rebuttal", "rebuttal": "> Question 1: Other ways to learn the prior distribution\n\nThis 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$**.\n\nTherefore, 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.\n\n*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.\n\nHere are the results of this experiment:\n\n|                  | 5k          | 10k         | 20k         | 30k         |\n| :--------------- | :---------- | :---------- | :---------- | :---------- |\n| cocos (autoencoding) | 83.1 ± 0.9  | 81.1 ± 1.3  | 92.2 ± 1.3  | 93.8 ± 0.3  |\n| cocos (action prior) | 70.0 ± 1.8  | 81.0 ± 1.0  | 84.5 ± 0.6  | 87.2 ± 0.6  |\n| vanilla (w/o cocos)  | 66.5 ± 1.6  | 78.3 ± 0.5  | 83.7 ± 0.3  | 82.9 ± 0.4  |\n\nWe 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.**\n\n---\n---\n---\n\n> Question 2: Section 2 Clarity\n\nWe 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.\n\n---\n---\n---\n\nThank 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.**"}]

Author rebuttal:

Meta-review: This paper studies the training of CNF models for robot control, which approximate distributions over actions given state or goal information, such as images. A small modification is introduced into the modelling: instead of transporting a fixed source distribution to a conditional target by a learned conditional vector field, it is proposed to make the initial distribution dependent on the condition using a pretrained representation. This modification is shown to improve success rates on a variety of simulated and real robot control tasks.

Main strengths, all identified by multiple reviewers:
- This is a simple, indtuitive, widely applicable idea that delivers very strong results on a variety of tasks.
- Comprehensive experiments, including applications to real robots and well-chosen illustrations.
- Clear exposition.

Main weaknesses:
- The theoretical guarantees are not very meaningful and not stated in proper mathematical language. The authors did not respond to this point beyond promising improvement. The authors are strongly encouraged to improve this section.
- Questions about the manner of defining the conditional source distribution, addressed by a new experiment using a different pretrained representation that approximates the action distribution.
- Questions about how the method handles multimodality in the action distribution and about scaling, both answered satsifactorily in the rebuttal.

The reviewers unanimously recommend acceptance, and I agree with this recommendation.
Decision: Accept (poster)