Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning
Authors: Moritz Reuss
Jyothish Pari
Pulkit Agrawal
Rudolf Lioutikov
Venue: ICLR 2025 Poster (2025)
Keywords: Robotics
Imitation Learning
Primary area: applications to robotics, autonomy, planning
TL;DR: Mixture-of-Denoising Experts (MoDE), a novel Diffusion Policy that leverages a noise-conditioned routing strategy to achieve more efficient denoising and improved performance compared to prior approaches.

Abstract: Diffusion Policies have become widely used in Imitation Learning, offering several appealing properties, such as generating multimodal and discontinuous behavior.
As models are becoming larger to capture more complex capabilities, their computational demands increase, as shown by recent scaling laws. 
Therefore, continuing with the current architectures will present a computational roadblock. 
To address this gap, we propose Mixture-of-Denoising Experts (MoDE) as a novel policy for Imitation Learning.
MoDE surpasses current state-of-the-art Transformer-based Diffusion Policies while enabling parameter-efficient scaling through sparse experts and noise-conditioned routing, reducing both active parameters by 40\% and inference costs by 90\% via expert caching.
Our architecture combines this efficient scaling with noise-conditioned self-attention mechanism, enabling more effective denoising across different noise levels. 
MoDE achieves state-of-the-art performance on 134 tasks in four established imitation learning benchmarks (CALVIN and LIBERO). 
Notably, by pretraining MoDE on diverse robotics data, we achieve 4.01 on CALVIN ABC and 0.95 on LIBERO-90. 
It surpasses both CNN-based and Transformer Diffusion Policies by an average of $57\%$ across 4 benchmarks, while using 90\% fewer FLOPs and fewer active parameters compared to default Diffusion Transformer architectures. Furthermore, we conduct comprehensive ablations on MoDE's components, providing insights for designing efficient and scalable Transformer architectures for Diffusion Policies. 
Code and demonstrations are available at https://mbreuss.github.io/MoDE_Diffusion_Policy.

Published as a conference paper at ICLR 2025
EFFICIENT DIFFUSION TRANSFORMER POLICIES WITH MIXTURE OF EXPERT DENOISERS FOR MULTITASK LEARNING
Moritz Re Jyothish Pari Pulkit Agrawal', Rudolf Lioutikov 1 Intuitive Robots Lab (IRL), Karlsruhe Institute of Technology, Germany 2 *Departeent of Electrical Engineering and Computer Science (EECS), MIT, USA
ABSTRACT
Diffusion Policies have become widely used in Imitation Learning, offering several appealing properties, such as generating multimodal and discontinuous behavior. As models are becoming larger to capture more complex capabilities, their compu- tational demands increase, as shown by recent scaling laws. Therefore, continuing with the current architectures will present computational roadblock. To address Experts (MoDE) as a novel policy for
this gap, we propose cixture-of-DDeoiisiii Imitation Learning. MoDE surpasses current state-of-the-art Transformer-baase -based Diffusion Policies while enabling parameter-efficient scaling through sparse ex- perts and noise-conditioned routing, reducing both active parameters by 40% and inference costs by 90% via expert caching. Our architecture combines this efficient scaling with condicoondt self-attention mechanism, enabling more effective denoising across different noise levels. MoDE achieves toite--ffthe---e art perfor- mance on 134 tasks four established imitation learning benchmarks (CALVIN and LIBERO). Notably, by pretraining MoDE on diverse robotics data, we achieve 4.01 on CALVIN ABC and 0,95 on LIBERO-99 90. It surpasses both CNN-based and Transformer Diffusion Policies by average of 57% across 4 benchmarks, while using 90% fewer FLOPs and fewer active parameters compared default Diffusion Transformer architectures. Furthermore, we conduct comprehensive ablations MoDE's components, providing insights for designing efficient and scalable Trans- former architectures for Diffusion Policies. Code and demonstrations are available at whttts:: ftttps://mbbeusse github. 10/ D/MoDE Diffusionn
INTRODUCTION
Diffusion models learn to reverse an iterative process that adds Gaussian noise to data samples (Ho
al. 2020 Song et al.) 2020). After training, they can generate samples conditioned on goals
like instructions or images. Recently, diffusion models have gained widespread adoption as policies
Imitaaion Learning (IL) OOto Model Team al. 2023 Reuss al., 2023 Chi al. 20033. IL is
a powerful paradigm to train agents from expert demonstrations to learn versatile skills (Pomerlcau
1989.
Nair et al. 2017 Pari et al. 2021
Fu et al.
2024
Diffusion policies offer several appealing properties for IL they can generate diverse multimodal behavior (Jia al.) 2024), scale with more data (Octo Model Team et al., 2023). and handle discontinuities in the action space (Chi et al. 2023). However, major limitation is their high computational cost, resulting in slower training and inference speed as models become larger. Standard architectures contain hundreds of millions of parameters (Chi et al. 2023) and require many denoising steps to generate actions. Large encoder modules for images and text further increase the computational requirements for IL policies. This restricts their use in real-time robotics applications, particularly in scenarios with limited on- board computing resources, such as mobile robots.
To address these challenges, explore of (MoE) that can scale model capacity while using fewer FLOPs for training and inference. The core idea behind a sparse MoE is to utilize only subset of the total model parameters during each forward pass. This is achieved by having

---

Published as a conference paper at ICLR 2025
Roster Experts
((
Rouder
Router
Router Experta
Add Norts
Head Altertinn
Seloct
Seleut &
((Or))
Φ(0U)
Selected Experts
Denoising Process
Polee
Gaei
Figure I: The proposed MoDE architecture (left) uses transformer with causal masking, where each block includes noise-conditional self-attention and a noise-conditioned router that assigns tokens to expert models based on the noise level. This design enables efficient, scalable action generation. On the right, the router's s activation of subsets of simple MLP experts with Swish-GL. U activation during denoising is illustrated.
multiple expert subnetworks and a routing model, that sparsely activates experts and interpolates their outputs, based on the input.
We introduce Mixture-of- Denoising Experts Policy (MoDH) a scalable and efficient MoE Diffusion Policy. Our work is inspired by prior results showcasing the multitask nature of the denoising process (Hang al. 2024), where there is little transfer between the different phases in the denoising process. We present a novel noise-conditioned routing mechanism, that distributes tokens to our experts based on the current noise level. MoDE leverages noise-conditioned self-attention combined with a noise input token for enhanced noise-ineection Our proposed Policy surpasses previous Diffusion Policies with higher efficiency and demonstrates sota performance across 134 diverse tasks in challenging goal- col--ooditioned imitation learning benchmarks: CALVIN (Mees et al. (022b) and LIBERO (Liu et al. 2023). Through comprehensive ablation studies, we investigate the impact of various design decisions, including token routing strategies, noise-injection techniques, expert distribution and diverse pretraining on a large-scale robot dataset (Collaboration et al. 2023). We summarize contributions below:
We introduce MoDE novel Mixture-of- Diffusion Policy that achieves state-of- the-art performance while using 90% fewer FLOPs and less active parameters than dense transformer baselines thanks to moise-based expert caching and sparse Mo design.
We demonstrate MoDE's effectiveness across 134 tasks in 4 benchmarks, showing an average 57% performance increase over prior Diffusion Policies while maintaining improved computational efficiency.
We present detailed ablation studies that investigate the importance of routing strategies and noise-injection, visualizing expert utilization across denoising steps to identify key components of MoDE
2 RELATED WORK
Diffusion in Robotics. In recent years, Diffusion Models (ong & Ermon 2019 Ho et al. 2020
Karras et al.) 2022) have gained widespread adoption in the context of robotics. They are used as a
policy representation for Imitation Learning (Chi et al. 2023 Reuss et al.) 2024 Xian al. 2023
Ke et al. 2024 Li et al. 2023b Scheikl et al. 2023) and in Offline Reinforcement Learning AAjay
al.] Janner et al. [0022] Pari al. (022). Other applications of diffusion models in robotics
include robot design generation (Wang et al. 2023), video-generation (Du et al. [0233 Ko et al.
2023 Ajay et al. 2023b) and motion planning (Cavvalho et al.) 2023 Urain al. 2023). The most
common architecture for using diffusion models as policy in robotics is a Convolutional neural
network (CNN) with additional FiLM conditioning (Perez et al. 2018) guide the generation based
on context information. Recently, the transformer architecture has been adopted as strong alternative
"Correspondence to: moritz.reuss @ kit.edu Both authors contributed equally.
2

---

Published as a conference paper at ICLR 2025
backbone for diffusion policies, specifically in IL Examples include Octo (Octo Model Team et al. 2023), BESO (Reuss et al.] 2023) and 3D-Diffusion-Actor (Ke et al. (024). Sparse-Dr leverages MoE to learn specialist experts for different tasks (Wang et al. (024). However, prior work considers using a Mixture of Experts architecture for improving the computational efficiency and inference time of the denoising process for generalist policies.
Mixture-of-Experts, MoB are a class of models where information is selectively routed through the model. The modern version of MoE was introduced in (Shazeer al. 2017). where routing or gating network conditionally chooses a subset of experts to send an input to. After Transformers
Vassann et al. 2017) proved to be an effective model that scales well with data, they were modified to have expert feed-forward networks at each block of the model in (Fedus et al.] (022) which presented Switch Transformers. Switch Transformers laid the groundwork that still widely adopted in different Sarreeeenntaarreee.e
(L..M) (Jiang et al. 2024 Du et al.] 2022). This allowed for more total parameters while keeping the forward and backward FLOOss smaller than their dense counterpart, thus yielding significant performance gains. However, training both the router and experts parallel is non-trivial optimization problem, and it can yield suboptimal solutions such as expert collapse where experts learn similar functions instead of specializing (Chi et al., (022). In addition, router collapse occurs when the router selects a small subset the experts and doesn't t utilize all the experts. This is mitigated with load balaneing losses (Shazeer al. 2017: Fedus et al.. 2022) which encourage the router to distribute inputs more evenly across experts. Multiple works have explored different methods to perform routing, such as expert choice routing (Zhou et al. 2022) differential k-selection (Hazimeh et al. (021), frozen hashing functions (Roller et al., 202 and linear assignment (Lewis al.) 2021).
Multi Task Learning in Diffusion Models. It has been shown that the denoising process is multi-task Hang et al. (2024). Leveraging this idea, works have adopted architectures suited for multi-tat learning. Some works have explicitly scheduled which parameters are specialized to which stage in
the denoising process (Park al. 2023; Go et al. 2023). extension to this (Park et al. (024) uses the scheduling as guidance during training but also learns how to modulate representations based on the denoising stage. Finally, some works have employed different architectures based on the
denoising stage (Lee et al. 2024b).
METHOD
In this section, introduce MoDE our novel MoE Diffusion Policy. First, formulate the problem of learning a policy for IL Next, we summarize the framework used in Diffusion Policies and then introduce our MoDE architecture with our novel noise-conditioned routing and noise-conditioned self-attention that enable efficient policy design. Finally, we explain our expert caching mechanism for efficient inference and explain the pretraining MoDE
3.1 PROBLEM FORMULATION
We consider the problem of learning language-conditionel policy Ta(ā T 8, given a dataset of robot demonstrations T. The policy predicts sequence of future actions à (a..... a;+5-1) 1 of length j. conditioned on the history of state embeddings ø (8 8,) of length h and desired goal The dataset contains 7 E T trajectories, where trajectory consists of a sequence of triplets of state, actions, and goal (s,, ai, gi,g). g is a language instruction. Our policy is trained to maximize the log-likelihood of the action sequence given the context of state history and goal:
L1. = E
log (ālz,g)
(5,0.g)ET
3.2 DIFFUSION POLICY
MoDE uses the continuous- time diffusion model EDM (arras et al. 2022) as a policy repre- sentation. Diffusion models are type of generative model for generating data by initially adding noise through Gaussian perturbations and then reversing this process. MoDE applies the score-based diffusion model to represent the policy ma ( g) The perturbation and inverse process can be
3

---

Published as a conference paper at ICLR 2025
described using stochastic differential equation:
dā = (3.o; - log ( 1
+
where 3, controls the noise injection, duis refers to infinitesimal Gaussian noise, and Pt (.) is the score function of the diffusion process, that moves samples away from regions of high-data density in the forward process. To generate new samples from noise neural network is trained to approximate the score function Va log Pr s. g) via Score matching (SM) (Vincent 2011)
Ls = a. [a( + c, 3. 01) a|2]
where Dā c, ) is the trainable neural network. During training, we sample noise from training distribution and add an action sequence. The network predicts the denoised actions and computes the SM loss.
After training, we can generate new action sequence starting from random noise by approximating the reverse SDE or related ODE in discrete steps using numerical ODE integrators. Therefore, we sample noise from the prior at ~ (0, of and iteratively denoise it. MoDE uses the DDIM-solver, numerical ODE-solver designed for diffusion models (Song al. (021), that allows fast denoising of actions in a few steps. MoDE uses 10 denoising steps in all our experiments.
3.3 MIXTURF-OF-EXPREE DENOISING
We now introduce MModddueeMMMDH MMDE
a novel approach that employs ciisiccondiionee expert routing to enhance diffusion-basee policies. This novel routing mechanism enables us to precompute and fuse the required experts for more efficient inference. An overview of MoDE's architecture with the routing mechanism is shown in Figure
For language conditioning, MoDE leverages a frozen CLIP language encoder to generate latent goal vector, while image encoding utilizes FiLMMcondiiioned R The model processes a sequence of input tokens X E Riokens and noise level σ₁ A linear projection layer ((o,) encodes the noise level into token, which incorporated into X. The complete MoDE architecture, MoDE(X, )) consists transformer blocks, each specialized for different denoising phases.
We now define each block f as composition of self-attention (SA) layer and an MoE layer,
(' ((a))) MoE(SA(X) X, X.
key change in our approach the integration of noise-aware positional embeddings. By adding ((art) to all tokens in X before celf-attention:
X = = +X,
(5)
enable each token adapt its attention patterns based the current denoising phase. This design enhances denoising performance without introducing additional parameters or architectural complexity.
The self-attention mechanism follows the standard formulation Vaswani et al. 2017
1 XW
SA(X) softmax(
(6)
VD
Our MoE layer introduces a novel noise-conditioned routing strategy. Given experts Egi = 1 the layer output is:
MoB(X, ((oo))
R(d(..))E(())
where the routing function R(-): BR:::Riooeeese Iinkknss Iookensx determines expert selection:
R(o(ae)) =
(8)
Unlike traditional MoE approaches that route based on input content, MoDE's routing mechanism specifically considers the noise level. This enables specialized experts for different denoising phases,
4

---

Published as a conference paper at ICLR 2025
1
Expert,
Router
0+++
Caching
After training MoDE, the router is noise-conditioned, allowing pre-computation of the experts used at each noise level before inference. This enables removing the router and retaining only the selected experts, significantly improving network efficiency.
Figure 2:
allowing for both improved performance and computational efficiency through expert caching (de- tailed in /suusubsection 3.3.T). We utilize the same method from (Muennighofi et al. 2024) to initialize the router, which is from truncated normal distribution with std 0.02. We implement topk using multinomial sampling without replacement, selecting k elements according to softmax( (W) probabilities. To maintain gradient flow through the non-differentiable sampling process, scale expert outputs by routing probabilities and renormalize the selected probabilities. To prevent expert collapse, we incorporate load balancing loss (LB) (Fedus et al. 2022
LB N
1R()(o)))) 0)
1 B
Σ
= 11 (10ssftmax(((rr)Waa.)) (9)
with = 0.01 as the balancing coefficient.
3.3.1 ROUTER AND EXPERT CACHING
To make our method more efficient, we exploit the fact that our MoE is conise-conditioned, meaning where at each noise level, the routing path is deterministic and be precomputed. By doing so, can determine the chosen experts ahead of time for each noise level. Figure lillustrates this process. This allows to fuse the selected expert MLPs into single, composite MLP. effectively reducing the computation cost. Instead of looping through each expert individually, this fused expert MLP can run in parallel, which substantially decreases the overall latency of the network. By climinating the need for dynamically invoking each expert, we not only save time but also streamline memory access patterns, reducing the overhead typically associated with routing decisions a traditional |MoF setup. Our caching strategy reduces the FLOPs overhead by over 80% compared to standard MoElouuutts and makes it two times faster during inference.
3.4 GENERALIST POLICY PRE-TRAINIA
We pretrain MoDE on a diverse mix multi-robot datasets curated from the OXE dataset Collabora [ion et al. (2023). Our training data encompasses 196k trajectories selected from six diverse datasets, featuring various robot platforms and manipulation tasks. The pretraining process of MoDE runs for 300k steps on distributed Cluster of 6 GPUs over three days For finetuning we freeze the pretrained routers each layer and only finetune the other model components. comprehensive overview of our pretraining dataset composition and methodology is provided in the Appendix (subsection A.
In detailed evaluations using the real2sim benchmark SIMPLER et al. 2024b). MoDE demon- strates superior performance compared to state-of-the-art generalist policies. It achieves an average success rate of 26.30% across diverse manipulation tasks, outperforming both VLe (2024) (23.70%) and Octo Octo Model Team et al. (022) (17.75%). Full evaluation details provided in insuuuubsectton 1.2.1
VLA Kim et al.
4
EVALUATION
Our experiments aim to answer four key questions: (I) How does prior diffusion transformer architectures in terms of performance? (II) Does large-scale pre-training
to other policies and
5

---

Published as a conference paper at ICLR 2025
of diverse robotics data boost the performance of MoDE (III) What is MoDE efficiency and speed compared to dense transformer baselines? (IV) How does the model utilize different experts during the action-denoising process?
We compare MoDE against prior diffusion transformer architectures (Chi et al. 2023), ensuring fair comparisons by using similar number of active parameters. [MoDE]us uses 8 layers with 4 experts and
a latent dimension of 1024 in all experiments. Our pretrained variant slightly bigger with 12 layers and 4 experts with the same latent dimension of 1024.
We use an action chunking length of 10 and a history length of 1 for all experiments. MoDE does execute all 10 generated actions without early re-planning temporal aggregation. Detailed hyperparameters are provided in the Appendix (Table 3).
4.1 LONG-HORIZON MULTI-TAST EXPERIMENTS
0.10
Avg. Succoss Hate
0.8
0.6
0.4
10
90
IDDP-T
DEDP-CN
I QueST
MoDE
PrPr-MMoDE
(a) LIBERO-90 Tasks
(b) Results for LIBERO-10 and LIBERO-90
Figure 3: Visualization and Results for LIBERO environment. (a) Few example environments and tasks of the LIBERO-90 task suite. (b) Average results for both LIBERO challenges averaged over 3 seeds with 20 rollouts for each task.
We first evaluate MoDE on the LONG-challenge and LIBERO-90 challenge of the LIBERO bench- mark (Liu et al., 2023). The LONG challenge requires model to learn 10 tasks in different settings. demands long-horizon behavior generation with several hundreds of steps for completion. The 90 variant tests policies on 90 diverse short-horizon tasks in different environments. Figure 3a visualizes a few examples of these tasks. All environments feature two cameras: static one and wrist-mounted camera, used to encode the current observation using FiLM-ResNets- 18. We test each policy 20 times on each task and report the average results over 3 seeds. MoDE] and all other diffusion architectures use conddiiondt ResNets- 18 with CLIP sentence embedding encode the goal and the images.
Baselines. We compare MoDE against three state-of-the-art baselines: 1) The Diffusion Transformer (DP-T) architecture (Chi et af. 2023), which conditions on noise and observations using cross- attention module. The standard Diffusion Policy CNN-based architecture (DP-CNN) 3) QueST (Mete et al. 2024), a transformer-bas based policy that learns discrete action representations using vector- quantized embeddings of action sequences. We tested all baselines ourselves, except for QueST, whose results are taken directly from their paper.
The performance of all models on the benchmark summarized in Figure 3b Overall, MoDE achieves the highest average performance in both benchmarks, while the QueST baseline the second best in the LIBERO-99 setting and the CNN-architecture is second best the long horizon setting. These results demonstrate MoDE]s MoDDE ability to learn long-horizon tasks with high accuracy. The performance gap is more pronounced the challenging LIBERO-10 experiment, where MoDE is the first policy to achieve over 90% success rate. Furthermore, MoDE surpasses prior best Diffusion baselines by an average of 16% in both settings, all while maintaining its computational advantage. The pretrained|MoDE variant achieves even higher performance in both settings, showing the potential of diverse pretraining. This showcases MoDEEss ability to achieve state-of-the-an performance with efficient use of computational resources.
Results.

---

Published as a conference paper at ICLR 2025
No. Instructions in a Row (1000 chains)
Trasn--Tost
Method
Active Params in Million
PrT
2
3
4
Avg. Len.
DP-CNN
321
86.3%
72.7%
60.1%
51.2%
41.7%
3.16=0.06
DP-T
194
78.3%
53.9%
33.8%
20.4%
11.3%
1.98+0.09
RoboFlamingo
1000
96.4%
89.6%
82.4%
74.0%
66.0%
4.09
ABCD-D
GR-1
130
94.9%
89.6%
84.4%
78.9%
73.1%
4.21=#0.00
MoDE
277
X
96.6%
90.6%
86.6%
80.9%
75.5%
4.30=00.22
MoDE
436
97.1
92.5%
87.9%
83.5%
77.9%
4.39+
DP-CNN
321
X
63.5%
35.3%
19.4%
10.7%
6.4%
1.35
DP-T
194
62.2%
30.9%
13.2%
5.0%
1.6%
1.13:0.02
RoboFlamingo
1000
82.4%
61.9%
46.6%
33.1%
23.5%
2.40=0.00
ABC-D
SuSIE
860+
87.0%
69.0%
49.0%
38.0%
26.0%
2.69=0.00
GR-1
130
85.4%
71.2%
59.6%
49.7%
40.1%
3.00
MoDE
307
x
91.5%
79.2%
67.3%
55.8%
45.3%
3,39
MoDE
436
96.2%
88.9%
81.1% %
71.8%
63.5%
4.01=00.44
Table Performance comparison on the two CALVIN challenges. The table reports average success rates for individual tasks within instruction chains and the average rollout length (Avg. Len.) to complete consecutive instructions, based on 1000 chains. Zero standard deviation indicates methods without reported average performance. "Prt" denotes models requiring policy pretraining. Parameter counts exclude language encoders.
4.2 SCALING MULTI-TASK EXPERIMENTS
'iitt the pink block sliding cabinet"
"store the grasped block drawer"
push the red open drawer"
"pull the handle to
"gresp and lift the block"
(a) Environments
(b) Example CALVIN-Rollout
Figure 4: Overview of the CALVIN environment. (a) CALVIN contains four different settings (A,B,C,D) with different configurations of slides, drawers and textures. (b) Example rollout consisting of 5 tasks in sequence. The next goal is only given to the policy, if it manages to complete the prior.
We evaluate MoDE on the CALVIN Language-Skinn Benchmark (Mees al. 2022b), com- prehensive image-basee benchmark for In the ABCD +D challenge, models learn from 22, 966 interaction sequences across four environments (A, B. C. D), each comprising 64 timesteps and 34 diverse tasks requiring complex skill chaining (Figure 4a). Following the official protocol, we test all policies on 1000 instruction chains of 5 sequential tasks in environment D (Figure 4b). Models receive 1 point per completed task and must complete each task to progress. We report average sequence length over 3 seeds with 1000 instruction chains each.
Baselines. We test MODE against several methods specialized for learning language-conditionel behavior and against other baseline diffusion policy architectures. We also compare MoDE against RoboFlamingo and GR-1. RoboFlamingo is fine- tine-tuned Vision-Language Action model, that contains around 3 billion parameters and has been pre-trained on diverse internet data. GR-1 is generative decoder-only Transformer pretrained on large-scale video generation and then co-finetuned on CALVIN (Wu al. 2024). If available, we report the average performance of all prior work directly from their paper, given the standard evaluation protocol in CALVIN (Mees et al. 2022a).
Results. Our findings (Table show that MoDE achieves the highest average success rate across all policies, even without pretraining. This surpasses extensively pretrained baselines like RoboFlamingo and GR-1. While GR-1 uses fewer parameters (130M vs 436M), MoDE requires significantly fewer FLOPs during inference (1.53 vs 27.5 GFLOPS) while maintaining comparable speed (12.2 12.6 ms). The combination of state-of-th -ftte--rt performance and computational efficiency makes MoDE particularly suitable for multitask applications.

---

Published as a conference paper at ICLR 2025
4.3 ZERO- SHOT GENERALIZATION EXPERIMENTS
Finally, then extend our investigation to the ABC-D challenge in the second phase, testing MoDEIs zero-shot generalization abilities. In this experiment, models are only trained on data from the first three CALVIN environments A,B,C and tested the unseen setting of environment D, which has different positions of relevant objects and texture of the table. This requires policies, that are able generalize their learned behavior new environment configurations and different textures, which is especially challenging. We evaluate MoDE trained from scratch and MoDE pretrained on a sub-set of Open-X-Emenn E-mo-mmmmmdm Data. This enables us study the zero-shot performance and pretraining effectiveness of MoDE
Baselines. For this experiment, we compare MODE against the previous CALVIN baselines, with the addition of SuSIE (Black et al., 2023). A hierarchical policy utilizing finetuned mmage-generation model, Instruct2Pix (Brooks al. 2023), to generate goal images, which guide a low-level diffusion policy. The high-level goal generation model requires large-scale pretraining. SuSIE's results are based on 100 rollouts only, without standard deviation, due to the computational cost of generating subgoal images.
The results of this experiment are summarized in the lower part of Table MoDE outper- forms all tested baselines and surpasses all other Diffusion Policy architectures by wide margin. Further, by pretraining MoDE on diverse robotics data, we achieve a new sota performance of 4.01. Therefore, in response to Question (I), affirmatively conclude that Mixture-of- of--Exerts models not only enhance scaling performance but also demonstrate strong zero-shot generalization capabili- ties. In addition, challenging zero-shot settings.
Results.
answer Question (II) by concluding that pretraining boost performance in
4.4 COMPUTATIONAL EFFICIENCY OF MODE
Avx. Inferace Time (mz)
120
7,000
6,000
5,772
64
60
3,000
2.71
20
16
25
2,000
1,000
0.7
189
1,445
361
128
512
128
512
Dasch Size
Datch Size
MoDE w/ Cache
MoDE w/o Cache
Dense-T
Figure 5: Computational efficiency comparison between MoDE and a Dense Transformer model with the same number of parameters. Left: Average inference speed over 100 forward passes for various batch sizes. Right: FLOP count for MoDE with router cache and without compared against a dense baseline. MoDE demonstrates superior efficiency with lower FLOP count and faster inference thanks to its router caching and sparse expert design.
We compare MoDE
against dense transformer baseline with similar parameters by measuring average inference time and FLOPs across batch sizes. As shown in Figure Figure 3) MoDE with caching significantly improves computational efficiency at batch size 1, inference is 20% faster (1 12ms vs 15ms), while at batch size 512, MoDE requires 16x fewer FLOPs (361 vs 5,772 GFLOPS) and achieves nearly 40% faster inference (64ms vs 104ms). These results demonstrate that MoDE delivers both superior task performance and substantial computational efficiency through its architecture and caching mechanism. detailed comparison of inference speed and FLOPS against all other baselines on CALVIN is summarized in issubsection A.5)
4.5 ABLATION STUDIES
To thoroughly evaluate MoDE's design choices, conducted series of ablation studies. These experiments address our research questions: the computational efficiency MoDE (uestiion III), the impact of routing strategies (Question IV), and the token distribution (Question V).
8

---

Published as a conference paper at ICLR 2025
4.5.1 WHAT DESIGN DECISIONS AFFECT MoDE's PERFORMANCE?
First, we assess the impact of various design decisions on MoDE performance. We ablate the choice of noise-conditioning and various MoE strategies on the LIBERO-10 benchmark. The results are summarized in Table 2
Lojse-Injection Ablations. Our experiments reveal the importance of proper noise conditioning in MoDE The full MoDE which uses both imput noise tokens and coise-conditioned self- attention, achieves the best performance with an average success rate of 0.92. Removing the input noise token slightly decreases performance to 0.90, highlighting the complementary nature of both conditioning methods. Using only the noise token for conditioning, without noise-conditioned self- attention, further reduces performance to 0.85. Interestingly, using FiLM conditioning (Perez et al., 2018), common approach in image-diffusion (Peebles & Xie 2023 yields the lowest performance in this group at 0.81. These results underscore the effectiveness of our proposed noise conditioning strategy in MoDE demonstrating a clear performance advantage of 0.08 over the FiLM approach.
MoE Ablations. Next, ablate several design decisions re- garding Mixture-of-Experts. First, we test the topk number of used experts. Setting topk only marginally lowers the av-
Arg. Success.
0.92
Impul Noior Token
0.90
Notse cond Attention
0.85
erage performance from 0.92 to 0.91. MoDE maintains robust
FILM Noise Conditioning
0.81
performance even with single expert. We also examine the
TopK=1 Shorod Expert 7= 0.1
0.91
0.90
effect of using a shared expert, where the model consistently employs the same expert in all cases. This approach achieves a comparable average success rate of 0.90. Different choices for the token-distnn -okenndistribution loss are also tested. While|MoDE] usee 0.01 a default value, we experiment with 7 values of 0.1 and 0.001, which result in average success rates of 0.90 and 0.86, respectively. These results indicate that y value of 0.01 strikes the best performance.
0.90
y 0.001
0.86
256 Emnhed Dim 512 Embedd Dims
0.86
0.87
Table 2:
Ablation Studies for MoDE on LIBERO-10. All results are averaged over 3 seeds with 20 rollouts each.
Latent Dimension. We investigate the impact of varying the latent dimension in MoDE testing dimensions of 256, 512, and 1024 (our default). The results show that increasing the latent dimension from 256 512 yields modest performance improvement from 0.86 to 0.87, while further increasing to 1024 provides a more substantial boost 0.92. This suggests that larger latent dimension allows MoDE to capture more complex representations, leading to improved performance.
4.5.2 OPTIMAL. ROUTING STRATEGY FOR DIFFUSION TRANSFORMERS
Next, we answer Question (III) by testing different routing strategies for our diffission-rrasffmmee policy across several environments. We test two different token routing strategies: 1) Token-only conditioned Routing and 2) Noise-only Token Routing. (1) is commonly used in I.IMs, where the routing decided based on the tokens only. We test these strategies in five experiments and report the average performance over 3 seeds: Noise-only Routing achieves an average normalized performance of 0.851, slightly outperforming Token-only Routing, which achieves 0.845. Detailed results are summarized in Table in the Appendix. The results demonstrate the effectiveness of our proposed routing strategy. While the performance difference is small, Noise-only Routing offers an additional advantage: the ability to predict all used experts based on noise levels once before roll-outs, enabling faster inference, as described in subsubsection 3.3.1 This is particularly beneficial for robotics applications.
4.5.3 HOW DOES THE MODEL DISTRIBUTE THE TOKENS TO DIFFERENT EXPERTS?
To address Question IV, we analyzed how MoDE distributes tokens to different experts using pre-trained model. Figure 6 visualizes v the average usage of each expert in each model layer during inference across various noise levels, using 10 denoising steps for clarity. Our analysis reveals that MoDE learns to utilize different experts for various noise levels, suggesting that the router has specialized for different noise regimes. A transition in expert utilization occurs around dg. In the first layer, the model learns an expert specialized for low-n levels, primarily used in the last denoising step min-- We conduct more ablations studies with our pretrained model and various other versions

---

Published as a conference paper at ICLR 2025
Denoising Process
6
6
83
4
65
10
Layezs
Experts
Figure 6: Visualized Expert Utilization. The average usage of all experts in MoDE across all layers shown. Purple color corresponds to low usage and yellow color high one, and each image is separately normalized. The average activation shows that MoDE learns to utilize different experts for different noise levels.
of MoDE subsubsection A.7.1 of the Appendix. These findings affirmatively answer Question IV, demonstrating that MoDE effectively distributes tokens across experts based on noise levels.
4.5.4 How DOES THE MODEL SCALE WITH MORE EXPERTS?
Finally, analyze the effect of increasing the number of experts in MoDE The results presented in(Figure 8jin the Appendix, where we evaluate MoDEjon the CALVIN ABCD and CALVIN ABC benchmarks using 2, 4, 6, and experts. For comparison, Dense-small and Dense-large. Dense-small shares the same latent dimensionality as MoDE while Dense-large is scaled up 2024 dimensions, matching MoDE overall parameter count. Our analysis focuses on how scaling affects both general performance (C-ABCD) and zero-shot generalization (C- ABC). In the ABCD environment, MoDE with 4 experts achieves the best performance. Interestingly, increasing beyond 4 experts degrades performance, possibly due to overfitting or increased routing complexity. the zero-shot generalization (ABC), MoDE with 4 experts still performs best. Notably, (ABC),MMoDE the Dense-small variant consistently underperforms across both tasks, underscoring the efficiency of the MoE architecture utilizing parameters more effectively. We hypothesize, that 4 experts has an ideal trade-off for the context of noise-only routing of Diffusion Policies. The different expert specialization patterns observed in Figure and Figure 12 from the Appendix show that expert specialization based noise regions. MoDE with more than 4 experts does not have a performance benefit. Overall, MoDE demonstrates that it can achieve comparable or superior performance dense transformer models while requiring fewer computational resources.
include two dense MoDE baselines:
4.6 LIMITATIONS
MoDE still has certain limitations. In our experiments, find that exhibits slightly higher standard deviation compared to the baselines. We hypothesize that the router's initialization has significant impact on overall optimization, requiring future work on stabilizing routing models. In addition, when visualizing expert utilization, in some of our experiments we noticed that only a subset of the total experts were being utilized phenomenon known as expert collapse (Chi et 2022).
5 CONCLUSION
In this work, we introduced fixture-of-Denoistte Experts (MoDE), novel Diffusion Policy that leverages a mixture of experts Transformer to enhance the performance and efficiency diffusion policies. We also proposed noise-conditions routing strategy for learning specialized experts within our model. In our extensive experiments and ablation studies across diverse benchmarks, we demonstrated the advantages of MoDE to outperform prior Diffusion Policies with a lower number of parameters and 90% less FLOPS during inference. future work, we want to experiment with more routing strategies, such expert-choice routing (Zhou et al. 2022
10

---

Published as a conference paper at ICLR 2025
6 ACKNOWLEDGMENTS
We would like to thank Adam Wei, Anurag Ajay, Hao-Shu Fang, Anthony Simeonov, Yilun Du for their insightful discussions and feedback. The work was funded by the German Research Foundation (DFG) 448648559, The authors also acknowledge support by the state of Maden-Württembergg through HoreKa supercomputer funded by the Ministry of Science, Research and the Arts Baden- Württemberg and by the German Federal Ministry of Education and Research. The research was also sponsored by the Army Research Office and was accomplished under ARO MURI Grant Number W911NF-23-1-00277.
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In International Confer- on Learning Representations, 2023a. URL https::/ net ferttfooou?idd sP1fo2K9DFG
Anurag Ajay, Seungwook Han, Yilun Du, Shaung Li, Abhi Gupta, Tommi Jaakkola, Josh Tenenbaum, Leslie Kaclbling. Akash Srivastava, and Pulkit Agrawal. Compositional foundation models for hierarchical planning. arXiv preprint arXiv:2309 N.XV::779.08887, 2023b.
Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: visual language model for few-shot learning. Advances in Neural Information Processing Systems, 35:23716- 23736, 2022.
Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint http:2220915577.. 2022.
Homanga Bharadhwaj. Jay Vakil, Mohit Sharma, Abhinav Gupta, Shubham Tulsiani, and Vikash Ku- mar. Roboagent: Generalization and efficiency in robot manipulation via semantic augmentations and action chunking, 2023.
Kevin Black, Mitsuhiko Nakamoto, Pranav Atreya, Homer Walke, Chelsea Finn, Aviral Kumar, and Sergey Levine. Zero-shot robotic manipulation with pretrained image-editing diffusion models. arXiv preprint arXiv:2310. 10639, 2023.
Denis Blessing, Onur Celik, Xiaogang Jia, Moritz Reuss, Maximilian Xiling Li, Rudolf Lioutikov, and Gerhard Neumann. Information maximizing curriculum: A curriculum- approach for learning versatile skills. Thirty-seventhh Conference on Neural Information Processing Systems, 2023. URL BRRLhtps:: httttp:://peezeviee..o / wttp://eeeetoodddeeeeeeeeeemaeee 7eW6N22E4
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Tomas Jackson, Sally Jesmonth, Nikhil Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Isabel Leal, Kuang-Huei Lee, Sergey Levine, Yao Lu, Utsav Malla, Deeksha Manjunath, Igor Mordatch, Ofir Nachum, Carolina Parada, Jodilyn Peralta, Emily Perez, Karl Pertsch, Jornell Quiambao, Kanishka Rao, Michael Ryoo, Grecia Salazar, Pannag Sanketi, Kevin Sayed, Jaspiar Singh, Sumedh Sontakke, Austin Stone, Clayton Tan, Huong Tran, Vincent Van- houcke, Steve Vega, Quan Vuong, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Tianhe Yu, and Brianna Zitkovich. Rt-1: Roboties transformer for real-world control at scale. In arXiv preprint arXiv:2112.06817, 2022.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choromanski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, et al. Rt-2: Vision Vision-language-accio models transfer web knowledge to robotic control. arXiv preprint arXiv:2307. 15818, 2023.
Tim Brooks, Aleksander Holynski, and Alexei A. Efros. Instructpix2pix: Learning to follow image editing instructions. In CVPR, 2023.
11

---

Published as a conference paper at ICLR 2025
Joao Carvalho, An T Le, Mark Baierl, Dorothea Koert, and Jan Peters. Motion planning diffusion: Learning and planning of robot motions with diffusion models. In 2023 JEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pp. 1916-1923. IEEE, 2023.
Onur Celik, Dongzhuoran Zhou, Ge Li, Philipp Becker, and Gerhard Neumann. Specializing versatile skill libraries using local mixture of experts. In Aleksandra Faust, David Hsu, and Gerhard Neumann (eds.), Proceedings of the Sth Conference on Robot Learning, volume 164 of Proceedings of Machine Leurning Research, pp. 1433-1433. PMLR, 08-11 Nov 2022. URL m
htttps:/pproced/ www.gngww...........ppp.p..tt.ee prwss/v 64/ ce wit.iikk2aa....... html
Onur Celik, Aleksandar Taranovic, and Gerhard Neumann. Acquiring diverse skills using curriculum reinforcement learning with mixture of experts. arXiv preprini arXiv:2403. 06966, 2024.
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 Robolics: Science and Systems (RSS), 2023.
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, Heyan Huang, and Furu Wei. On the representation collapse of sparse mixture of experts. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https //openreview.oo net nwt.:orunntetetteee/eaaeaaaooaat worumm id- www.dwwwwacccc22
Open X-Embodiment Collaboration, Abhishek Padalkar, Acorn Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, Antonin Raffin, Ayzaan Wahid, Ben Burgess-Limerick, Beomjoon Kim, Bernhard Schölkopf, Brian Ichter, Cewu Lu, Charles Xu, Chelsea Finn, Chenfeng Xu, Cheng Chi, Chenguang Huang. Christine Chan, Chuer Pan, Chuyuan Fu, Coline Devin, Danny Driess, Deepak Pathak, Dhruv Shah, Dieter Büchler, Dmitry Kalashnikov, Dorsa Sadigh, Edward Johns, Federico Ceola, Fei Xia, Freek Stulp, Gaoyue Zhou, Gaurav S. Sukhatme, Gautam Salhotra, Ge Yan, Giulio Schiavi, Hao Su, Hao-Shu Fang. Haochen Shi, Heni Ben Amor, Henrik I Christensen, Hiroki Furuta, Homer Walke, Hongjie Fang, Igor Mordatch, Ilija Radosavovic, Isabel Leal, Jacky Liang, Jaehyung Kim, Jan Schneider, Jasmine Hsu, Jeannette Bohg, Jeffrey Bingham, Jiajun Wu, Jialin Wu, Jianlan Luo, Jiayuan Gu, Jie Tan, Jihoon Oh, Jitendra Malik, Jonathan Tompson, Jonathan Yang. Joseph J. Lim, João Silvério, Junhyek Han, Kanishka Rao, Karl Pertsch, Karol Hausman, Keegan Go, Keerthana Gopalakrishnan, Ken Goldberg, Kendra Byrne, Kenneth Oslund, Kento Kawaharazuka, Kevin Zhang, Keyvan Majd, Krishan Rana, Krishnan Srinivasan, Lawrence Yunliang Chen, Lerrel Pinto, Liam Tan, Lionel Ott, Lisa Lee, Masayoshi Tomizuka, Maximilian Du Michael Ahn, Mingtong Zhang, Mingyu Ding, Mohan Kumar Srirama, Mohit Sharma, Moo Jin Kim, Naoaki Kanazawa, Nicklas Hansen, Nicolas Heess, Nikhil J Joshi, Niko Suenderhauf, Norman Di Palo, Nur Muhammad Mahi Shafiullah, Oier Mees, Oliver Kroemer, Pannag Sanketi, Paul Wohlhart, Peng Xu, Pierre Sermanet, Priya Sundaresan, Quan Vuong, Rafael Rafailov, Ran Tian, Ria Doshi, Roberto Martin-Martin, Russell Mendonca, Rutav Shah, Ryan Hoque, Ryan Julian, Samuel Bustamante, Sean Kirmani, Sergey Levine. Sherry Moore, Shikhar Bahl, Shivin Dass, Shuran Song, Sichun Xu, Siddhant Haldar, Simeon Adebola, Simon Guist, Soroush Nasiriany, Stefan Schaal, Stefan Welker, Stephen Tian, Sudeep Dasari, Suneel Belkhale, Takayuki Osa, Tatsuya Harada, Tatsuya Matsushima, Ted Xiao, Tianhe Yu, Tianli Ding, Todor Davchev, Tony Z. Zhao Travis Armstrong. Trevor Darrell, Vidhi Jain, Vincent Vanhoucke, Wei Zhan, Wenxuan Zhou, Wolfram Burgard, Xi Chen, Xiaolong Wang, Xinghao Zhu, Xuanlin Li, Yao Lu, Yevgen Chebotar, Yifan Zhou, Yifeng Zhu, Ying Xu, Yixuan Wang, Yonatan Bisk, Yoonyoung Cho. Youngwoon Lee, Yuchen Cui, Yueh hua Wu, Yujin Tang, Yuke Zhu, Yunzhu Li, Yusuke Iwasawa, Yutaka Matsuo, Zhuo Xu, and Zichen Jeff Cui. Open X- Embodiment Robotic learning datasets and RT-X models. /ttp/v/xxvvv////////v///e////.vv httgg//bs/233
08864 2023.
Zichen Jeff Cui, Yibin Wang, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. From play to policy: Conditional behavior generation from uncurated robot data. In International Confer- ence on Learning Representations, 2023. URL htttps:///ppeereeeeetee ew. net. wtt.moorrrmmmmmmmmmm
t 7 1M1777101N
Nan Du Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma,
12

---

Published as a conference paper at ICLR 2025
Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Hellserr,,,, Toju Duke, Lucas Dixon, Kun Zhang, Quoc V Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. Glam: Efficient scaling of language models with
2022.
Yilun Du, Mengjiao Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Joshua B Tenenbaum, Dale Schu- urmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. arXiv preprint arXiv:2302.001 11, 2023.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(1 23(120):1-39
2022.
Zipeng Fu, Tony Z Zhao, and Chelsea Finn. Mobile aloha: Learning bimanual mobile manipulation with low-cost whole-body -ody teleoperation. arXiv preprint arXiv:2401.02117, 2024.
Hyojun Go, Yunsung Lee, YounYoung Kim, Seunghyun Lee, Myeongho Jeong, Hyun Seung Lee, and Seungtaek Choi. Towards practical plug-and-play diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1962-1971, 2023.
Jiayuan Gu, Sean Kirmani, Paul Wohlhart, Yao Lu, Montserrat Gonzalez Aremas, Kanishka Rao, Wenhao Yu, Chuyuan Fu, Keerthana Gopalakrishnan, Zhuo Xu, Priya Sundaresan, Peng Xu, Hao Su, Karol Hausman, Chelsea Finn, Quan Vuong, and Ted Xiao. Rt-trajectory: Robotic task generalization via hindsight trajectory sketches. International Conference on Learning Representations, 2024.
Tiankai Hang. Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via min-snr weighting strategy, 2024.
Kay Hansel, Julen Urain, Jan Peters, and Georgia Chalvatzaki. Hierarchical policy blending as inference for reactive robot control. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 10181-10188. IEEE, 2023.
Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdhery, Maheswaran Sathiamorrthy, Yihua Chen, Rahul Mazumder, Lichan Hong, and Ed H. Chi. Dselect-k: Differentiable selection in the mixture of experts with applications multi-task learning. CoRR, abs/2106.003600 2021. URL
https::/axxiv.co wtg.
wwwp/ss..//ooooooooooooooooooa abs 2205 03760
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, pp. 9902-9915. PMLR, 2022.
Xiaogang Jia, Denis Blessing, Xinkai Jiang, Moritz Reuss, Atalay Donat, Rudolf Lioutikov, and Gerhard Neumann. Towards diverse behaviors: A benchmark for imitation lcarning with human demonstrations. In The Twelfth International Conference Learning Representations, 2024. URL hhttps://openrewe
net. forum? L
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.00 040888, 2024.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Tsung- Wei Ke, Nikolaos Gkanatsios, and Katerina Fragkiadaki. 3d diffuser actor: Policy diffusion with 3d scene representations. arXiv preprint arXiv:2402. 10885, 2024.
Tsung-Wei
Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan Foster, Grace Lam, Pannag Sanketi, et al. Openvla: An open-source vision-language-action model. arXiv preprint 46924.2406.09226, 2024.
13

---

Published as a conference paper at ICLR 2025
Po-Chen Ko, Jiayuan Mao, Yilun Du, Shao-Hua Sun, and Joshua B Tenenbaum. Learning Act from Actionless Video through Dense Correspondences. arXiv:2310.08576, 2023.
An Thai Le, Kay Hansel, Jan Peters, and Georgia Chalvatzaki. Hierarchical policy blending as optimal transport. In Learning for Dynamics and Control Conference, pp. 797-812. PMLR, 2023.
Seungjae Lee, Yibin Wang, Haritheja Etukuru, H. Jin Kim, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. Behavior generation with latent actions. arXiv preprint arXiv:2403.03111 031181 2024a.
Yunsung Lee, Jin Young Kim, Hyojun Go, Myeongho Jeong, Shinhyeok Oh, and Seungtaek Choi. Multi-architecture multi-expert diffusion models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 13427-13436, 2024b.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models, 2021.
Maximilian Xiling Li, Onur Celik, Philipp Becker, Denis Blessing, Rudolf Lioutikov, and Gerhard Neumann. Curriculum-bast based imitation of versatile skills. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 2951-2957, 2023a. doi: 10.1009/CCRA48811.200
10160543.
Xiang Li, Varun Belagali, Jinghuan Shang, and Michael S Ryoo. Crossway diffusion: Improving diffusion-based visuomotor policy via self-supervised learning, arXiv preprint arXiv:2307.01849) 2023b.
Xinghang Li, Minghuan Liu, Hanbo Zhang, Cunjun Yu, Jie Xu, Hongtao Wu, Chilam Cheang, Ya Jing, Weinan Zhang, Huaping Liu, et al. Vision-language foundation models as effective robot imitators. In International Conference on Learning Representations, 2024a.
Xuanlin Li, Kyle Hsu, Jiayuan Gu, Karl Pertsch, Oier Mees, Homer Rich Walke, Chuyuan Fu, Ishikaa Lunawat, Isabel Sich, Sean Kirmani, Sergey Levine, Jiajun Wu, Chelsea Finn, Hao Su, Quan Vuong, and Ted Xiao. Evaluating real-world robot manipulation policies in simulation. arXiv preprint ar 2xXx:2445.0094// 2024b.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.027447 2022.
Bo Liu, Yifeng Zhu, Chongkai Gao, Yihao Feng, Qiang Liu, Yuke Zhu, and Peter Stone. Libero: Benchmarking knowledge transfer for lifelong robot learning. arXiv preprint arXiv:2306.03310, 2023.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play, 2019.
Oier Mees, Andreas Eitel, and Wolfram Burgard. Choosing smartly: Adaptive multimodal fusion 2016 JEEE/RŠ. International Conference on
for object detection in changing environments. Intelligent Robots and Systems (IROS), pp. 151-156. IEFE, 2016.
Oier Mees, Lukas Hermann, and Wolfram Burgard. What matters in language conditioned robotic imitation learning over unstructured data. IEEE Robotics and Automation Letters (RA-I.), 74)) 11205-11212, 2022a.
Oier Mees, Lukas Hermann, Erick Rosete-Beas, and Wolfram Burgard. Calvin: A benchmark for language-conditiong policy learning for long-horizon robot manipulation tasks. IEEE Robotics and Automation Letters. 2022b.
Atharva Mete, Haotian Xue, Albert Wilcox, Yongxin Chen, and Animesh Garg. Quest: Self- supervised skill abstractions for learning continuous control. arXiv preprint a0X.v.240.118840 2024.
Niklas Muennighoff, Luca Soldaini. Dirk Groeneveld, Kyle Lo. Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, et al. Olmoe: Open mixture-of-expert. language models. arXiv preprint 20f.:2409002064 2024.
14

---

Published as a conference paper at ICLR 2025
Ashvin Nair, Dian Chen, Pulkit Agrawal, Phillip Isola, Pieter Abbeel, Jitendra Malik, and Sergey Levine. Combining self-supervised learning and imitation for vision-based rope manipulation. In 2017 IEEE international conference robotics and automation (ICRA), pp. 2146-2153. IEEE, 2017.
Johan Obando-Ceron, Ghada Sokar, Timon Willi, Clare Lyle, Jesse Farebrother, Jakob Foerster, Gintare Karolina Dziugaite, Doina Precup, and Pablo Samuel Castro. Mixtures of experts unlock parameter scaling for deep rl. arXiv preprint urXiv2242.088609 2024.
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, Dorsa Sadigh, Chelsea Finn, and Sergey Levine. Octo: An open-source generalist robot policy. https
Pfocto-modelle github. io 2023.
Jyothish Pari, Nur Muhammad Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The surprising effectiveness of representation learning for visual imitation, 2021.
Jyothish Pari, Nur Muhammad (Mahi) Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The Surprising Effectiveness of Representation Learning for Visual Imitation. In Proceedings of Robotics: Science and Systems, New York City, NY, USA, June 2022. doi: 10.15607/RSS,2000 2002.. XVIII.1O.
Byeongjun Park, Sangmin Woo, Hyojun Go, Jin- Young Kim, and Changick Kim. Denoising task routing for diffusion models. arXiv preprint arXiv:2310.07138, 2023.
Byeongjun Park, Hyojun Go, Jin- Young Kim, Sangmin Woo, Seokil Ham, and Changick Kim. Switch diffusion transformer: Synergizing denoising tasks with sparse mixture-of-expprts, arXiv preprint arXiv:2403.09176, 2024.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pp. $195-4205 2023.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Dean Pomerleau. Alvinn: An autonomous land vehicle in a neural network. In D.S. Touretzky (ed.), Proceedings of (NeurIPS) Neural Information Processing Systems, pp. 305 313. Morgan Kaufmann, December 1989.
Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal conditioned imitation learning using score-based diffusion policies. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Moritz Reuss, Ömer Erdinç Yagmurlu, Fabian Wenzel, and Rudolf Lioutikov. Multimodal diffusion transformer: Learning versatile behavior from multimodal goals. In Robotics: Science and Systems, 2024.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Su- sano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts, 2021.
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. Hash layers for large sparse models. CoRR, abs/2106.044e 04426, 2021. URL httpt:ttpss.o httpp//:///xxxvv.rr.///
04426
Paul Maria Scheikl, Nicolas Schreiber, Christoph Haas, Niklas Freymuth, Gerhard Neumann, Rudolf Lioutikoy, and Franziska Mathis-Ullrich. Movement primitive diffusion: Learning gentle robotic manipulation of deformable objects. arXiv preprint arXiv:2312. 10008, 2023.
Nur Muhammad Mahi Shafiullah, Zichen Jeff Cui, Ariuntuya Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning modes with one stone. In Thirry- Conference on Neural Information Processing Systems, 2022. URL URL http::p:://peeeeeeeeeeeeeeeeeeee http://orrummidaa-----v--00
15

---

Published as a conference paper at ICLR 2025
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gatee -ated mixture-of-expprts layer, 2017.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In ICLR, 2021.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
Hugo Touvron, Thibaut Lavril, Gautier Izacard. Xavier Martinet, Maric-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1 13971, 2023.
Julen Urain, Niklas Funk. Jan Peters, and Georgia Chalvatzaki. Se (3) difiifffusiossssdd:: Learning smooth cost functions for joint grasp and motion optimization through diffusion. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 5923-5930. IEEE, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computa- tion, 23(7):1661-1674 2011. 1. doi: 10. 10.1162/NECO a_ 4412222.
Homer Rich Walke, Kevin Black, Tony Z Zhao, Quan Vuong, Chongyi Zheng, Philippe Hansen- Estruch, Andre Wang He, Vivek Myers, Moo Jin Kim, Max Du, et al. Bridgedata v2: A dataset for robot learning at scale. In Conference Robot Learning, pp. 1723-1736 PMLR, 2023.
Tsun-Hsuan Wang. Juntian Zheng, Pingchuan Ma, Yilun Du Byungchul Kim, Andrew Everett Spielberg, Joshua B. Tenenbaum, Chuang Gan, and Daniela Rus. Diffusebot: Breeding soft robots with physics-augmented generative diffusion models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL htttts:///ppeneeeeeeepeetp ew. .w..et/.eeeeeeeeeee..e www.neetfoooro forum? 1d=1204 LQUEs
Yixiao Wang, Yifei Zhang, Mingxiao Huo, Ran Tian, Xiang Zhang, Yichen Xie, Chenfeng Xu, Pengliang Ji, Wei Zhan, Mingyu Ding, al. Sparse diffusion policy: A sparse, reusable, and flexible policy for robot learning. arXiv preprint arXiv:2407.01531, 2024.
Hongtao Wu, Ya Jing, Chilam Cheang, Guangzeng Chen, Jiafeng Xu, Xinghang Li, Minghuan Liu, Hang Li, and Tao Kong. Unleashing large-scale video generative pre-training for visual robot manipulation. In International Conference on Learning Representationss 2024.
Zhou Xian, Nikolaos Gkanatsios, Theophile Gervet, Tsung-Wee Ke, and Katerina Fragkiadaki. Chaineddiffuser: Unifying trajectory diffusion and keypose prediction for robotic manipulation. In 7th Annual Conference on Robot Learning. 2023. URL Forum? Mit.::ddd-0299222B2112
net/
Tony Z Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:200.11375,, 2023.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai, Zhifeng Chen, Quoc Le, and James Laudon. Mixture-of-experts with expert choice routing, 2022.
Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Ted Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Ste- fan Welker, Ayzaan Wahid, Quan Vuong, Vincent Vanhoucke, Huong Tran, Radu Soricut, Anikait Singh, Jaspiar Singh, Pierre Sermanet, Pannag R Sanketi, Grecia Salazar, Michael S Ryoo, Krista Reymann, Kanishka Rao, Karl Pertsch, Igor Mordatch, Henryk Michalewski. Yao Lu, Sergey Levine, Lisa Lee, Tsang- Wei Edward Lee, Isabel Leal, Yuheng Kuang, Dmitry Kalashnikov, Ryan Julian, Nikhil Joshi, Alex Irpan, brian ichter, Jasmine Hsu, Alexander Herzog, Karol Hausman,
16

---

Published as a conference paper at ICLR 2025
Keerthana Gopalakrishnan, Chuyuan Fu, Pete Florence, Chelsea Finn, Kumar Avinava Dubey, Danny Driess, Tianli Ding, Krzysztof Marcin Choromanski, Xi Chen, Yevgen Chebotar, Justice Car- bajal, Noah Brown, Anthony Brohan, Montserrat Gonzalez Arenas, and Kehang Han. RT-2: Vision- language-action models transfer web knowledge to robotic control. In 7th Annual Conference on Robot Learning, 2023. URL htt.thttt///pmmette
ew wwt//etttuuuee J=XMM
17

---

Published as a conference paper at ICLR 2025
A APPENDIX / SUPPLEMENTAL MATERIAL
Hyperparameter
CALVIN ABCD
CALVIN ABC
LIBERO-10
LIBERO-90
Pret-MoDE
Number of Transformer Layers
12
Number Experts
4
Heads
Action Chunk Size
10
10
10
10
History Leagth
1
Embedding Dimension
1024
1024
1024
1024
1024
Image Encoder
FILM-RoaNott8
FILM-ResNct50
FiLM-ResNeti8 CLIP VII-B/32
FiLM-RosNctt8 CLIP VIT-B/32
FiLM-ResNcc50
Lang Encoder
VIT-B/32
VIT-B/32
VIT-B/32
Attention Dropont
0.3
0.3
0.3
0.3
0.3
Residual Dropout
MLP Dropout
0.1
0.1
0,1
0.1
0,1
Optimizer
Adamm
Adam
Adam
Adam
Adam
Belus
[0.9, 0.95]
0.955
[0.9, 0.95)
[0.9. 0.95]
[0.9. 0.95)
Learning Rate
lc-4
le-4
lc-4
1o-4
lo-4
Transformer Weight Decay
0.1
Other weight decay
0.05
0.05
0.05
0.05
0,1
Batch Size Train Steps in Thousands
512
512
512
30
25
15
30
300
80
80
80
0
0.001
0.001
0.001
0.001
0.5
0.5
0.5
0.5
0.5
True
True Exponential DDIM
True
Time steps
Exponential DDIM
Exponential DDIM 460
Exponential DDIM
Exponential
Sampler
DDIM
Parameter Count (Millions)
460
460
460
685
Table 3: Summary of all the Hyperparameters for the MoDE policy used in experiments.
A.1 PRETRAINING DETAILS
Dataset
Weight
BC-Z
0,258768
LIBERO-10
0.043649
BRIDGE
0.188043
CMU Play-Fusion
0.101486
Google Fractal
0.162878
DOBB-E
0.245176
Total
1.000000
Table 4: Dataset sampling weights used for training MoDE on small subset of trajectories. The total dataset consists of 196k trajectories.
We pretrain a large variant of MoDE on a subset of available datasets from the Open-X-Embodimen Collaboration al. (2023) study the generalization ability of MoDE An Overview of the used dataset is summarized in Table 4 Our pretraining dataset comprises 196K trajectories from different sources, with weighted sampling across BC-Z (0.259), LIBERO-10 (0.044), BRIDGE (0.188), CMU Play-Fusion (0. 101), Google Fractal (0.163), and DOBB-E (0.245). This dataset includes demonstrations from diverse robot platforms including Google robots, Franka Pandas, and Hello-Stretch robots, covering wide range of manipulation tasks. The pretraining was conducted on NVIDIA A6000 GPUs with 40GB VRAM each over days, completing 300K training steps. We used batch size of 1024 with learning rate of le-4 and weight decay of 0.1. To ensure balanced dataset mixing during training, we implemented a large shuffle buffer of 400K samples. Each dataset was individually normalized to account for different scales and ranges across the various robot platforms. This diverse pretraining significantly improved MoDE's zero-shot generalization, particularly challenging benchmarks like CALVIN ABC-D where achieved the-art performance of 4.01 average rollout length. For reproducibility, we will release our pretrained model weights and preprocessing code to the community.
state-of-
18

---

Published as a conference paper at ICLR 2025
Benchmark
MoDE
DP-T
DP-CNN
Avg. Baseline
Improvement
CALVIN ABC-DD (norm.)
0.678
0.226
0.270
0.248
+151.1%
CALVIN ABCD+D (norm.) LIBERO-90 LIBERO-10
0.860
0.396
0.632
0.514
+36.1%
0.910
0.690
0.780
0.735
+16.7%
0.920
0.510
0.730
0.620
+26.0%
Average Improvement Over Second- -econd-Best:
57.5%
Table 5: Detailed Performance Improvement Analysis. CALVIN scores normalized by dividing by 5 to align with LIBERO scale. Improvement calculated as: (MoDE Avg. Baseline) / Avg. Baseline 100. Final average is the mean of improvements across all four benchmarks compared to the second best Diffusion Policy variant on each one.
Google Fracta
Berkeley Bridge V2
Figure 7: Example Scenes of the SIMPLER Li et al. (022b) benchmark used to test generalist
policies on various tasks from the Bridge and Google Fractal dataset.
For finetuning we freeze the routers of the model and remove the load- ballncilnni loss and train the local domain for 10k steps on LIBERO and 15k steps for CALVIN with batch size of 64 per GPU and 4 GPUs.
A.2 EXPERIMENTS DETAILS
Average Performance Increase. To quantify MoDE's advantages existing Diffusion Poli- cies, we compared its performance against the second-bes best method across all benchmarks. MoDE demonstrated substantial improvements, particularly in challenging transfer scenarios like CALVIN ABC-D where it outperformed the next best method by 15.11%. Even on the more standardized LIBERO benchmarks, MoDE maintained significant advantages of 16.7% 26.0%. Averaging across all tasks, MoDE achieved a 57.5% improvement over the second-best performing method while maintaining its computational efficiency with 90% fewer FLOPPs compared to dense tranformers with similar number of parameters.
A.2.1 MoDE EVALUATION ON SIMPLER
We evaluate MoDE's capabilities generalist policy by comparing it against two state-of-the-art models trained on substantially larger datasets from Open -- Embodiment: Octo (800K trajectories) OOtoo Model Team et al.) 2023) and Open (Kim et al.) (IM trajectories). We conduct this comparison using the SIMPLER benchmark, which provides real2sim variants of the BridgeV2 (Walee et al.) 2023) and Google Fractal datasets used to train RT-1 (Brohan et al. 2023). The
benchmark encompasses diverse manipulation tasks across multiple environments, as illustrated in
Figure
19

---

Published as a conference paper at ICLR 2025
OpenVLA
Octo Base
MoDe (ours)
Metric
Score
Rank
Score
Rank
Score
Rank
Drawer Open
16% %
1
0%
4.23%
Drawer Close
20%
2
2%
34.92%
Pick Can Horizontal
71 %
0%
33.78%
Pick Can Vertical
27%
2
0%
29.78%
Pick Can Standing
65 %
I
0%
36.44%
2
Move Near
48 %
I
3%
30%
Drawer Open
19%
2
1%
21.30%
Drawer Close
52%
2
44%
76.85
Pick Can Horizontal
27 %
21%
22%
Pick Can Vertical
3%
3
21%
40%
Pick Can Standing
19%
2
9%
35%
Move Ncar
46 %
4%
45.42%
Partial Put Spoon on Tablecloth
4%
3
35%
29.17%
Spoon Tablecloth
12%
12.5%
Partial Put Carrot on Plate
33%
2
53%
29.17%
Put Carrot on Plate
0%
3
8%
8.33%
Partial Stack Green Block on Yellow Block
12%
2
32%
8.33%
3
Stack Green Block Yellow Block
0%
2
0%
0%
Partial Put Eggplant Basket
67%
37.5%
Eggplant Basket
3
43%
8.33%
Average
23.70%
1.95
17.75%
2.1
26.30%
1.65
Table Detailed comparison of MoDE against two sota Generalist Policies OpenVLA (Kim et al.
2024 and Octo (Octo Model Team al. 2023) tested all SIMPLER tasks with 2952 evals.
The results of the evaluations are summarized in Table 6 On average MoDE achieves the highest average success rate of 26.30% and the best average ranking of 1.65 across all tasks, surpassing both Octo (17.75% success, 2.1 rank) and the 7.7B parameter Open model (23.70% success, 1.95 rank). MoDE shows particularly strong performance on challenging manipulation tasks like drawer manipulation (34.92% on drawer close) and precise object interactions (40% on vertical can picking). While specialized tasks like stacking blocks remain challenging for all models, MoDE's consistent performance across diverse tasks demonstrates effectiveness as scalable architecture for generalist policies.
A.2.2 CALVIN BENCHMARK
The CALVIN benchmark (ees et al. 2022b) is an established IL benchmark for learning language- conditioned behavior from human play data. In contrast other benchmarks the data does not contain structured demonstrations, where the robot completes one task, but instead, the dataset was collected by humans, that randomly interact with the environment. From these long-horizon trajectories across the 4 different settings, the authors randomly cut out short sequences with 64 frames and labeled them with the task label. While the dataset offers models to train on the unlabeled part too, restricted MoDD to only train on the labeled parts. The Franke Emika Panda robot is controlled using Bnnd-fffector Space with discrete gripper. We use two cameras encode the current scene: static camera and a wrist one and predict the next 10 actions, before receiving the next observations and generating another set 10 actions.
CALVIN ABC. We train MoDE and our dense transformer baseline for 25k training steps with
batch size of 512 on a GPU Cluster Node with 4 A6000 NVIDIA GPUs for 6.5 hours with all 1000 rollouts at the end of training. We report the mean results averaged over 3 seeds done in all relevant prior work. All baselines are reported from the original paper given the standardized evaluation protocol of CALVIN (Mees et al. 2022b).
CALVIN ABCD. We train MoDE and our dense transformer baseline for 30k training steps with batch size of 512 4 GPU Cluster Node with 4 A6000 NVIDIA GPUs for 7.5 hours with all 1000 rollouts at the end of training. We report the mean results averaged over 3 seeds as done in all relevant prior work.
20

---

Published as a conference paper at ICLR 2025
Model
Block Push
Relay Kitchen
CAL
CAL ABCD
L-10
Average
Dense T
0.96++0.02
3.73+0.12
2.83=
4.1340.11
0.839 144 0.91
Token- Rokeer
0.97.10.01
3.85.10.03
2.6710..04
4.29:10.08
0.90
0.845
Oc-Router
0.97-0.01
3.79:10.04
2.7910.16
4.3011.02
0.92
0.851 0.151
Table 7: Overview of the performance of all different token routing strategies used for MoDD across 5 benchmarks. We mark the best result for each environment in bold and the second best in cursive. We use CAL to represent CALVIN. To average the results, we normalize all scores and compute the average over all environments.
A.2.3 LIBERO BENCHMARK
LIBERO-10. The LIBERO-10 10 benchmark consists of 50 demonstrations for 10 different tasks that are all labeled with text instruction. The Franka Emika Panda robot is controlled using an end-effector controller. Similar to CALVIN all models have access to two camera inputs: a static one and wrist camera. We train MoDE]and and our dense transformer baseline for 50 epochs with a batch size of 512 on a GPU Cluster Node with 4 A6000 NVIDIA GPUs for 2 hours with all 200 rollouts at the end of training. The benchmark does require to test models on 10 different long- long-horizon tasks. We test each task 20 times for each model and report the final average performance overall 10 tasks.
LIBERO-90. The LIBERO-10 benchmark consists of 50 demonstrations for 90 different tasks that are all labeled with a text instruction. The Franka Emika Panda robot is controlled using an end-effector controller. We train MoDE and our dense transformer baseline for 50k steps with batch size of 512 on a 4 GPU Cluster Node with A6000 NVIDIA GPUs for 12 hours with all 1800 rollouts the end of training. The benchmark does require test models 90 different tasks in many different environments. We test each task 20 times for each model and report the final average performance overall 90 tasks.
A.3 BASELINES
Below explain several baselines used in the experiments in detail:
Diffusion Policy-CNNTT Inspired by (Chi al. 2023), we evaluate extension of the DDPM based Diffusson Policy framework for goal-conditioned Multi-task learning. We evaluate two versions: the CNN-based variant and the Diffusion-Transformee variant, that is conditioned on context and noise using cross-attention. For experiments also EDM-based Diffusion framework for fair comparison against[MoDE] MMDE We optimized the ideal number of layers and latent dimension for the Transformer baseline and our final version uses layers with latent dimension of 1024. Larger or smaller variants resulted in lower average performance.
RoboFlamingo. RoboFlamingo [Li et al. 2024a is Vision-LanggugeeMMddee (VLM) finetuned for behavior generation. The authors use a 3 billion parameter Flamingo model (Alayrac et al. 2022) and fine-tune on CALVIN by freezing the forward blocks and only fine-tuning new Perceiver Resampler module to extract features from frozen vision- transformerr image encoder and the cross- attention layers to process the image features. Finally, new action head is learned to generate actions. Overall, the finetuning requires training approx. 1 billion of the parameters. We report the reported results from the paper since they use the standard CALVIN evaluation suite.
SuSIE. This model first finetunes Instruct2Pix, an image-generation diffusion model, that generates images conditioned on another image and a text description (Brooks et al. 2023) on the local CALVIN robotics domain and uses it as high-level goal generator. The lewelevel policy is a based Diffusion Policy, that predicts the next 4 actions given the current state embedding and desired sub-goal from the high-level policy (Black et al. 2023).
GR-1 A causal GPT-Transformer model (Wu et al., 2024), that has been pretrained on large-scale generative video prediction of human videos. Later, the model is finetuned using co-training of action prediction and video prediction on CALVIN. We report the results directly from their paper for the CALVIN benchmark.
21

---

Published as a conference paper at ICLR 2025
A.4 RECTIFIED FLOW FINETUNING FOR MODE
We investigate the compatibility of MoDE with alternative denoising objectives by combining our
architecture with Rectified Flow (RF) (Albergo & Vanden-Eijnden, 2022; Lipman et al. 2022). Using our PxE-rretraiee MoDH model finetuned with an RF objective, we evaluate performance with reduced inference steps on the ABC-D setting.
No. Instructions in a Row (1000 chains)
Method
Steps
2
3
Avg. Len.
MoDE
10 96.2%
88.9%
81.1%
71.8%
63.5%
4.01=:0.04 3.87=00006
RFDE
2
93.5%
86.2%
77.5%
68.8%
60.5%
Table 8: Performance comparison of MoDE with standard vs RF objective on CALVIN ABC-D. Ro achieves comparable performance with just denoising steps.
As shown in Table 8
MoDE RF maintains strong performance while requiring only 2 denoising
steps compared to the original 10 steps. The modest performance drop (4.01 3.87 average rollout length) suggests that our noise-specialized experts successfully adapt to different denoising objectives. Future work could explore combining MoDE with Rectified Flow to further reduce inference steps while maintaining high performance.
A.S AVERAGE FLOPS COMPUTATION AND INFERENCE SPEED
We provide an in-depth comparison of the total parameters and FLOPs used for every method in Table 9) Additionally, we provide the computational efficiency metric per GFLOPS (10" FLOPS) compare the various methods and measure the average prediction time for a single action. In the following, we detail the average GFLOPS computation for all relevant baselines on the CALVIN ABC benchmark. Specifically, we compare the average GFLOPS required to predict single action.
To guarantee a fair comparison we evaluated all methods on the same NVIDIA A6000 GPU with 40 GB VRAM. To compute the average inference speed, we tested each method 100 times and removed large outliers to compute an average time.
MoDE. We benchmark the large pretrained variant with 12 layers, 4 experts and hidden dim of 1024. The average GFL.OPS for forward pass are 0.7 GFLOPS.. Without router caching, the model would require 5 GFLOPS, indicating that the router caching reduces the overall computational cost by over 90%. The architecture processes 14 tokens in total (1 noise + 1 goal 2 images 10 noisy action actions). MoDE predicts a sequence of 10 actions with 10 denoising passes. For the variant with ResNet-50, the image encoder requires 8.27 GFLOPS. On average for a single action, MoDE requires 10 forward passes with the transformer and a single pass with the ResNet-50. Consequently. the pretrained variant of MoDD needs 1.53 GFLOPS average to predict single action. The inference time for that model depends on the hardware. We measure an average inference time per action of 12.2.
DP-CNN/T. DP-CNN utilizes 0.8 GFLOPS for an average forward pass. The ResNet- require 3.62 GFLOPS. The model predicts 10 actions with 10 denoising steps and executes 10 actions without replanning. This results in the CNN version requiring 1.28 GFLOPS to predict a single action. For the Transformer version, the architecture predicts 10 actions using 10 denoising steps and processes 14 tokens in total (1 noise + 1 goal + 2 images + 10 noisy action actions) similar to MoDE It achieves an average GFLOPS usage of 1.8 GFLOPS per forward pass through the transtormer. The DP-T baseline requires 2.16 GFLOPS on average to predict a single action. The CNN version requires 11.7 ms to predict a single action on average and the transformer variant with its cross-attention conditioning needs 16.2 ms.
RoboFlamingo. For computational analysis, the model requires 34 GFLOPS to encode a single image with a ViT. For the policy backbone, we evaluated the "mpt-Ib-ree hegt--bbreppjjaa-00b--ee as used in the paper. This architecture requires 656 GFLOPS for an average sequence of 32 tokens per forward pass. While multiple variants of RoboFlaminge exist, this provides rough estimate of the required GFLOPS. In total, we estimate an average of 690 GFLOPS to predict an action in CALVIN. To predict single action, the model requires 65 on average.
variant
22

---

Published as a conference paper at ICLR 2025
Method
Active Params (M)
Total Params (M)
GFLOPS
PrT
Avg. Length
SF-Ratio
Inf. Time [ms]
Difi-P-CNN
321
321
1.28
1.35
1.05
11.7
Difl-P-T
194
194
2.16
1.13
0.53
16.2
RoboFlamingo
1000
1000
690
2.47
0.004
65
SuSIR
860+
860+
60
2.69
0.045
199
GR-1
130
130
27.5
3.06
0.11
12.6
MoDE (ours)
436
740
1.53
4.01
2.6
12.2
Table 9: Comparison of total and active number of parameters of methods used in the CALVIN benchmark. Additional overview of average FLOPS required by the different methods together with their average performance on the ABC benchmark. SF-Ratio compares average rollout length with GFLOPS.
SuSIE, In our computational analysis, we tested Instruct2Pix with 50 denoising steps as implemented by SuSIE. The resulting 1026 GFLOPS are divided by 20 the model only generates new subgoals every 20 timesteps. The low-level policy uses ResNet-50 image encoder with 8.27 GFLOPS. Contrast to other policies, SuSIe only computes one image per state and predicts actions every timestep. These are then averaged using exponential averaging. Thus, we omit the small diffusion head get an estimate of 60 GFLOPS per action. For the average inference speed we meassure the time to generate single goal image and divide it by 20, next we add the average time to encode two images with the ResNet-50 and forward passes through small MI.P. Every 20 timesteps, when SuSIE generates new image it requires 3777.62 ms for single action generation. Otherwise its lot faster with 10.7 ms. On average SuSIE requires 199 ms to generate a single action, which make the slowest policy overall.
The pretrained MAE Vision Transformer requires approx. 17.5 GFLOPS to encode a single image. The transformer backbone processes 150 tokens with a history length of 10 and 15 tokens per timestep (10 image tokens + goal token + proprioceptive token + 2 video readout tokens + action token per timestep). Consequently, the average GFLOPS for single action prediction utilizing the decoder with latent dimension 384 and 12 layers are 10 GFLOPS. In total, this results in an average computational cost of 27.5 for predicting single action in CALVIN. For the average single action prediction the model requires 12.6 ms.
GR-1
Analysis. Overall, MoDE offers the best performance GFLOPS ratio across all tested baselines. Although MoDE is significantly larger in total size compared to the other Diffusion Policy architec- tures requires similar inference speed and low FLOP count. Additionally, has demonstrated superior efficiency in terms of computational resources while maintaining high performance on the CALVIN benchmark tasks. For inference speed MoDE is the second fastest although has a high total parameter count.
A.6 DETAILED EXPERIMENTAL RESULTS
We summarize the ablations regarding the choice of routing in Table 7 Therefore, test two 2 different routing strategies across 5 benchmarks.
A.7
STATE-BASED EXPERIMENTS
We conduct additional experiments with MoDE on two established multi-task state- state-based environ- ments:
Relay Kitchen. We utilize the Franka Kitchen environment from (Lynch et al. 2019) to evaluate models. This virtual kitchen environment allows human participants to manipulate seven objects using VR interface: kettle, microwave, a sliding door, hinged door, light switch, and two burners. The resulting dataset consists 566 demonstrations collected by the original researchers, where each participant performed four predetermined manipulation tasks per episode. The Franka Emika Panda robot is controlled via a 9-dimensional action space representing the robot's 'ss joint and end-effector positions. The 30-dimensiona observation space contains information about the current state the relevant objects in the environment. As desired goal state, we randomly sample future states desired goal to reach.
23

---

Published as a conference paper at ICLR 2025
For this experiment, train all models for 40k training steps with a batch size of 1024 and evaluate
them 100 times as done in prior work (Shafiullah et al.] [2022] Cui et al.] [023] Reuss et al.] (2023) to guarantee a fair evaluation. All reported results averaged over 4 seeds. We train our models
on a local PC RTX with an RTX 3070 GPU for approx. 2 hours for each run with the additional
experimental rollouts.
Block Push. The PyBullet environment features an XArm robot tasked with pushing two blocks into two square targets within a plane. The desired order of pushing the blocks and the specific block-target combinations are sampled from the set of 1000 demonstrations as a desired goal state. The demonstrations used for training our models were collected using a hard-coded controller that selects block to push first and independently chooses a target for that block. After pushing the first block to a target, the controller pushes the second block to the remaining target. This approach results in four possible modes of behavior, with additional stochasticity arising from the various ways of pushing block into a target. The models only get a credit, if the blocks have been pushed in the correct target position and order. We consider block successfully pushed if its center is within 0.05 units of a target square.
All models were trained on a dataset of 1000 controller-generatee demonstrations under these randomized conditions. All models have been trained for 60k steps with a batch size of 1024. To evaluate them we follow prior work (Shafiullah et al. 2022 Cui et al., 2023 Reuss et al. 2023) and test them on 100 different instructions and report the average result over 4 seeds. We train models local PC RTX 3070 GPU for approx. 3 hours for each run with final evaluation. Demonstrations are sourced from a scripted oracle, which first pushes randomly chosen block to a selected square, followed by the other block different square. The policies are conditioned to push the blocks in the desired configuration using a goal state-vector. We chose an action sequence length of given a history length of for these experiments, which are inspired by our dense diffusion transformer baseline BESO (Reuss et al. 2023)
Baselines. In this setting, we compare MoDE against several SOTA goal-conditioned policies. We test two transformer architectures, C-BeT (Cui al. 2023) and VQ-BeT (Lee et al. 2024a), that predict discretized actions with an offset. C-BeT uses k-means clustering together with an offset vector while VQ-BeT leverages residual Vector Quantization to embed actions into a hierarchical latent space. Further, we test against a dense diffu- sion policy transformer model BESO (Reuss al.) 2023). BESO uses the same continuous- time dif- fusion policy combined with a dense transformer to predict a single action given sequence of prior states. To enable fair comparison, we chose the same hyperparameters for BESO and MoDE in both settings. We test all models averaged over 4 seeds and report the mean values directly from prior work (Lee et al.
Block Push
Relay Kitchen
C-BeT
0.87=)00007)
(.01
VQ-BeT BESO MoDE
0.80=(0.022)
(..04)
(.902+(0.02)
3.73+) 3.79-1(0.02))
0.97++00.01)
Table 10: Comparison of the performance of different policies on the state-based goal- conditioned relay-kitchen and block-push envi- ronment averaged over 4 seeds. MoDE outper- forms the dense transformer variant BESO and other policy representations on all baselines.
2024a
The results of both experiments are summarized in Table 10 MoDE achieves a new SOTA performance on both benchmarks and outperforms the dense transformer variant of BESO in both settings. Further, MoDE achieves higher performance compared to other policy representation methods such as VQ-BeT and C-BeT.
Results.
A.7.1 MIXTURE-OF- ABLATIONS
Q: How does the Load Balancing Loss influence the Expert Distribution?
We analyze how the load balancing loss affects expert distribution across noise levels by training MoDElon LIBERO-10 10 with varying load balancing weights TLB E [0.1, 0.01, 0.001, 0.0001 Figure visualizes the resulting expert distributions.
With a high load balancing loss of 0.1, experts are used almost uniformly across all layers with minor variations in two out of eight layers (Figure 9a). However, this enforced uniformity comes at cost
24

---

Published as a conference paper at ICLR 2025
CALVIN ABC
CALVIN ABCD
3
4.5
Performance
2.5
2
3.5
2
8
2
4
N- N-Experts
N-Experts
MoDE
Dense-large
Dense-small
Figure 8: Scaling performance of MoDE and Dense-MoDE on CALVIN ABC and ABCD envi- ronments, showing average performance for to 8 experts using best-performing variants for each environment.
Expert Usage across Layers
Expert Usage across Layers
25%
25%
35%
15%
25%
25%
20%
30%
Number
Layer Numbrer
30%
20%
140
20%
25%
25%
25%
25%
20%
25%
25%
25%
20%
10%
25%
20%
32%
25%
1.0%
45%%
25%
25%
25%
25%
25%
25%
35%
15%
Index
Index
(a) "LD
0.1
(b)
0.01
Expert Usage across Layers
Expert Usage across Layers
20%
15%
45%
30%
4.3%
5%
5%
30%
15%
5%
35%
Number
Layer Number
50%
20%
5%
3
Export
Expert
(c) QLB
0.001
(d) Tt.8 0.0001
Figure 9: Average Expert Utilization for different Load Balancing Weights across all denoising levels.
the average performance drops to 0.9. This result suggests that enforcing equal expert usage across noise levels may constrain the model's learning capacity.
25

---

Published as a conference paper at ICLR 2025
Expert Usage
Expert Usage Layers
30%
20%
50%
231
40%
10%
25%
30%
20%
50%
2.5%
25%
25%
25%
40%
30%
30%
2004
25%
25%
30%
Number
Layer Number
30%
30%
20%
23%
2504
25%
25%
4.
40%
20%
40%
23%
30%
20%
25%
5-
90%
20%
30%
20%
35%
30%
15%
6-
30%
20%
50%
15%
30%
35%
20%
7-
40%
20%
40%
25%
10%
35%
30%
Expert Index
Expert Index
(a) Experts
(b) Experts
Export Usage across Layers
Expert Usage across Layors
30%
20%
10%
20%
0%
20%
20%
30%
20%
10%
20%
33%
0%
30%
20%
10%
20%
0%
20%
0%
Layer Number
Layer Number
1.0%
10%
2.%
10%
10%
-40
20%
0%
30%
10%
20%
0%
20%
1.0%
10%
23%
0%
1.0%
20%
10%
0%
20%
20%
0%
20%
10%
20%
30%
20%
10%
20%
20%
30%
20%
0%
0%
10%
30%
12%
20%
30%
10%
10%
30%
20%
10%
0%
20%
20%
30%
20%
0%
30%
0%
0%
30%
10%
20%
20%
20%
2
7
Index
Expert Index
(c) 6 Experts
(d) Experts
Figure 10: Average Expert Utilization for different number of used experts across all denoising levels.
At 0.01, we observe a more flexible distribution while maintaining good overall expert utilization Figuure 9b). Within individual layers, expert usage varies from 10% to 40% in various layers. This enables for specialization while preventing any expert from becoming dominant. This configuration achieves the best average performance and provides evidence that moderate specialization benefits the policy most.
As reduce 7LB further to range of [0.001, 0.0001], the expert distribution gets worse and the model there is an expert collapse is happening for lb of 0.0001 in several layers (Figure 9c and Figure 9d). These values show, that the load balancing weight is crucial guarantee good usage of all experts for the model. The lower performance of these models (0,86 for TLB = 0.001 and 0.83 for TLB = 0.0001) suggests that balanced expert participation important for optimal Diffusion Policy performance.
Q: What happens with increased number of Experts? Next, we study the load balancing distribu- tion for models that use more experts. Our experiments with 6 and 8 expert variants on LIBERO-LE 10 reveal challenges in scaling beyond 4 experts. As shown in Figure 10, configurations with more ex- perts consistently underutilize their capacity, with most layers actively using only 4 experts even with increased load balancing regularization (7 = 0.1). These larger configurations also show decreased performance, indicating that experts provides optimal balance between representation capacity and learning efficiency. We hypothesize that additional experts impair the model's s ability to learn effective representations from image and language inputs while maintaining noise-level specialization. This phenomenon persists even when increasing the load balancing loss factor, further supporting our finding that 4 experts offers the best trade-off between model capacity and performance.
26

---

Published as a conference paper at ICLR 2025
Expert Usage
Expert Usage
0%
075
60%
4.04
095
40%
60%
60%
07
40%
0%
6075
40%
0%
0/
2015
095
0.1%
0%
20%
80%
Number
Layer Number
100%
0%
0%
094
20%
50%
0%
0%
40%
20%
4.14
0%
50%
20%
24
0%
60%
4.54
50%
40%
0%
0%
60%
40%
20%
60%
20%
40%
0%
40%
20%
80%
20%
0%
Expert Index
Expert Index
(a) topk=1
(b) topk=1 with shared Expert
Figure 11: Average Expert Utilization for two variants that topk=1: left MoDE with 4 experts, right MoDE with 4 experts, where I is used in all settings as a shared one. Overall, MoDE struggles to equally distribute tokens across experts in this setting.
Q: How is the expert distribution with topk=1? We further study the impact of using top; 1 for MoDE We analyze MoDE's behavior with topk=1 routing, both with and without a shared expert, using 4 experts total. As visualized in Figure 1T topk=1 configurations struggle maintain balanced expert utilization, typically collapsing to using only two of the four available experts. While this configuration shows only modest performance degradation on LIBERO-10, its impact is more visible on the challenging CALVIN ABC benchmark, where the average rollout length drops from 3.34 (topk=2) 2.72 (topk=1). These results highlight the importance of expert interpolation across noise levels, with topk=2 providing the necessary flexibility to blend expert specializations for optimal performance.
Q: What expert distribution patterns emerge for different noise levels?
Finally, we analyze the different experts loads for various noise levels of our biggest model after pretraining on diverse robotics data. Detailed results are shown in Figure 12
Our analysis of the expert distribution reveals several key insights about how the model learns to organize its mixture-of-experts architecture:
Noise-Level Specialization
Experts specialize between high-noise (01-07) and low-noise (ag-010) regimes, partic- ularly visible in L4 where E3 and EO show distinct high activation patterns different noise regimes
CCler transition point around dg where expert utilization patterns shift, most evident in L44 and L5 with stark changes in expert activation
Smooth handoff between denoising phases, exemplified in L. where expert activations gradually transition across noise levels
-
Layer-wise Organization
Early layers (LO-L4) show distinct expert specialization, particularly visible in L4's strong alternating expert patterns
Middle layers (L.5-L77) demonstrate more distributed expert utilization, shown by more balanced activation patterns
Later layers (L8-L11) return clearer specialization, evident in Lil's distinct expert preferences
First layer (LO) shows particularly strong specialization for low-noise scenarios, with clear expert preferences in the final denoising steps
Expert Role Distribution
Different experts develop specialized roles, most clearly visible where experts show strong preferences for specific noise ranges
27

---

Published as a conference paper at ICLR 2025
Some experts consistently handle high-noise denoising (evident in L3 and L4) while others focus on low-noise refinement (for example L7-L8)
"Transitional experts" emerge in middle noise ranges, particularly visible in L5's balanced activation patterns
Model employs different expert combinations across layers, well demonstrated in the contrasting patterns between L4 and L8
Load Balancing Properties
Effective load balancing achieved across experts, particularly visible in the distributed patterns of 1.5-1.7
No single expert dominates across all noise levels, as shown by the varied activation patterns in each layer
Smooth transitions between noise regimes, most evident in the gradual activation changes in L4 and L5
These findings provide strong evidence that MoDH effectively partitions the denoising process across its experts, with each expert naturally specializing in different phases. The emergence of this organization without explicit supervision supports the effectiveness of our architectural choices and routing strategy.
100
This image shows the expert usage distribution across 12 layers in a Mixture of Experts (MoE) model. Each subplot represents different layer (from layer 0 to layer 11), with experts labeled EO, E1, E2, and E3 on x-axis. The y-axis indicates the log-scaled token count, displaying the frequency of tokens routed each expert within that layer. The color gradient indicates the proportion of tokens assigned to each expert, where darker colors represent higher usage.
Figure 12:
28

---

Published as a conference paper at ICLR 2025
A.8 EXTENDED RELATED WORK
Mixtre-of-Expertt in Robotics. In the context of robotics, MoB models are used in many settings without being combined with transformer architecture. Several works mixture of small MLP policies, that focus on different skills in Reinforcement Learning (Ohando-Ceron et al. 2024 Celik
al., 2022; 2024) for robot motion generation(Hansel et al. 2023; Le al., 2023). another
body of work utilizes combinations of small CNNs robot perception (Riguelme et al. [ Mees
al. 2016). Further applications include learning multimodal behavior using a mixture of Gaussian policies (Blessing et al. 2023 Li et al.
2023a). Despite the extensive usage of MoElin in many domains, no prior work has tried to utilize MoE logether with Diffusion Policies for scalable and more efficient
Diffusion Policies.
Transformers for Robot Learning. Transformer models have become standard network archi- tecture for many end-to-end robot learning policies in the last few years. They have been combined with different policy representations the context of IL. One area of research focuses on generating sequences of actions with Variational Autoencoder (VAE) models (Bharadhwaj al. 2023 Zhao
jet al. 2023). These action-chunking transformer models typically use an encoder-decoder transforme
policy architecture. Several Diffusion Policies, such as Octo (Octo Model Team al. (023).
BESO (Reuss et al.) 2023). ChainedDiffuser (Xian et al. 2023) and 3D-Diffusion-Acto leverage transformer model as a policy backbone. Another direction of research treats behavior generation
as discrete next-token predictions similar to auto-regressive language generation (Touvron et al.
2023). C-Bet, RT-1, and RT-2 use discretized action binning to divide seen actions into k-classes (Cui
et al.
2023 Shafiullah et al. 2022 [Brohan et al. 2022 Zitkovich et al., (023)) while VQ-BeTT (Lee al.
2024a) learns latent actions with residual Vector Quantization. Several works have shown the
advantages of using pre-trained LLM or VLM
policy backbone, which are then finetuned for
action generation (Brohan et al. 2023 Gu et al.
2024, Collaboration et al., 2023 Li et al. 2024a). None of the recent work considers using any Mixture-of-Expee architecture for policy learning.
MoDE is the first architecture to leverage MoE
architecture combined with diffusion for behavior
generation.
29


Reviews:
[{"review_id": "gcnrg0htNI", "reviewer": "Reviewer_DSi8", "summary": "This paper presents a novel policy architecture for imitation learning based on sparse Mixture of Experts and EDM diffusion policies. The architecture is shown to improve performance in most cases across four domains (two from LIBERO, and two from CALVIN benchmarks), while jointly improving efficiency by reducing the number of active parameters required to perform inference and sample from the policy. The main architectural contributions are (1) a noise-conditioned routing mechanism that selects experts based on the timestep and noise level of the diffusion process, and (2) a noise-conditioned attention mechanism.\n\nThe proposed method, MoDE, is the first policy to achieve a `>90%` success rate in the LIBERO-10 benchmark, and performs competitively with models pre-trained on significantly more data (GR-1, RoboFlamingo), despite training only on data available from target benchmarks. Thorough ablations showcase the importance of each of the two main architectural contributions, and explore the space of key hyper-parameters, including the number of experts, and the extent of router load balancing regularization.\n\nWhile the idea of sparse routing layers for transformers is not new, and has been shown to lead to comparable gains in scaling + inference efficiency in pure language modeling, this paper takes the important step of applying these techniques to Diffusion-based policies in decision-making tasks.", "questions": "Questions have been woven into the strengths section, see above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "## Originality:\n\nThe idea of sparse routing layers is not new, as mentioned in my summary. This technique is generally expected to produce gains in inference efficiency by reducing the number of active parameters, and has been used in LLM with Mixtral [1], and in Diffusion transformers in [2], showing similar efficiency gains as this paper shows for imitation learning tasks. The main originality in this paper comes from the architectural contributions, including the noise-conditioned routing mechanism, and noise-conditioned attention mechanism. In previous work in [2], the sparse MoE router is conditioned on previous layer hidden states, whereas in this paper the sparse MoE router is conditioned on just the noise token. This design choice has the potential to simplify the router’s learning task, acting as an implicit regularization, and causing experts to specialize to different denoising steps, rather than different tokens.\n\n[1] Mixtral of Experts, Albert Q. Jiang et al. 2024\n\n[2] Scaling Diffusion Transformers to 16 Billion Parameters, Zhengcong Fei et al. 2024\n\nTo my knowledge, these architecture contributions are original, and their great performance across four benchmark domains highlights their value for efficient learning.\n\n---\n\n## Quality:\n\nThe paper is organized well, concepts are introduced and explained thoroughly, and experiments are chosen to highlight the improvements in policy performance, reduced inference cost, and the role that each component of the proposed method has to overall performance. Experiments are thorough, and appropriate baselines are chosen, including the dense version of the proposed method, previous diffusion-based policies, and previous models pretrained at scale on internet data (which would be expected to have an advantage to the larger source of data). Its impressive that on CALVIN ABC→D tasks, MoDE performs comparable to GR-1, despite lacking a large-scale pre-training step.\n\nA minor quality issue with Table 1: performance is reported as a percent without a confidence interval. The experimental quality is generally good in this paper, but could be further improved by reporting confidence intervals for the results, so that readers can judge how statistically significant the improvements are.\n\nAnother manner by which to improve the paper would be to include more examples in “HOW DOES THE MODEL DISTRIBUTE THE TOKENS TO DIFFERENT EXPERTS?” In particular, for different load balancing weights, how does this impact the observed usage of the experts? In addition, results in Figure “(b) Scaling performance of MoDE” (the figure number seems to be missing here) show that performance saturates at 4 experts. Its unclear from the results if this is just an artifact of the limited training data available in LIBERO and CALVIN domains, and it would strengthen results to explain this behavior, and to include more examples of what happens to the experts' usage as the number of experts increases.\n\nA similar analysis can be done to study the impact of the top k routing parameter.\n\n---\n\n## Clarity:\n\nWriting is clear, organized well, and easy to follow. Design choices for the main contributions are explained carefully, and clearly. Questions that remained after reading the main text, such as hyperparameters, benchmark settings, are presented in the supplementary materials.\n\n---\n\n## Significance:\n\nThe main contributions, while performant and efficient, are not particularly surprising. Previous works such as [1] and [2] have shown the viability of sparse Mixture of Experts for improving performance and inference efficiency in other domains, and [2] in particular has shown this for Diffusion transformers. What’s significant about this paper is the empirical study of what type of sparse MoE mechanism works best for Imitation Learning in robotics, where data is often much smaller than for pure text and image generation. Getting MoE to perform well in this data-limited regime is challenging, and findings in this paper, such as the noise-conditioned routing that pre-computes routing paths upfront, improve the viability of MoE in this setting.", "weaknesses": "Weaknesses have been woven into the strengths section, see above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NDre8TyZTK", "reviewer": "Reviewer_jByy", "summary": "The paper proposes to implement the architecture for a deep imitation learning policy model by a Mixture of Denoising Experts (MoDE). It enables parameter efficient scaling and reduces the inference cost. The paper introduces a routing scheme that conditions the expert selection on the current noise level. The experiments show that the proposed MoDE model performs better than CNN and Transformer-based baselines.", "questions": "- Does the model generate all of the sequential actions' decisions at the same time? For instance, in the ABCD->D experiment, it involves interactions each consisting of 64 timesteps and 34 diverse tasks. Does the model determine the 64 timesteps' actions at the same time after denoising in Fig. 1?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The idea is straightforward and the writing is easy to follow.\n- The experiments show that the proposed idea performs better than the baselines.\n- The experiments are thorough, exploring different aspects of performance, speed, routing strategy, and expert utilization.", "weaknesses": "- I think the novelty of the paper is a bit limited. I am not expert in the imitation learning literature, but the paper seems to be mostly applying mixture of experts idea (that is already been in the LLMs literature) to the diffusion-based imitation learning setup. For instance, using the top-k router and load balancing regularizations are not new contributions.\n- The improvements on the LIBERO-90 and CALVIN datasets are marginal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yB9WZKPWh6", "reviewer": "Reviewer_nWR1", "summary": "To address the issue of models becoming increasingly large and difficult to capture more complex functions, along with their rising computational demands, this paper proposes a Mixture of Denoising Experts (MoDE) as a new strategy for imitation learning. MoDE is based on a diffusion strategy using Transformers and achieves parameter-efficient scaling, significantly reducing inference costs. To achieve this, MoDE employs sparse experts combined with a novel routing strategy that selects experts based on the current noise level of the denoising process. This is further enhanced by a noise-regulated self-attention mechanism. MoDE achieves state-of-the-art performance across 134 tasks in four established imitation learning benchmarks (Calvin and LIBERO).", "questions": "1、Can you elaborate on the specific innovations introduced by your Mixture-of-Experts Diffusion Policy that differentiate it from existing methods in the literature? What novel insights does your approach provide beyond the combination of MOE and diffusion methods?\n2、In Figure 1, how does your proposed router strategy select different experts based on the current noise level? Could you provide additional details or examples to clarify this process and better illustrate its implementation?\n3、The methods section appears to be brief and heavily based on existing research. Can you elaborate on the unique contributions of your work and how they extend beyond the existing literature? What specific methodologies have you developed or modified?\n4、Regarding the experimental results for ABC→D in Table 1, which are lower than those of the baseline methods, could you provide further justification or insight into why this occurred? What factors might have influenced these results, and how should they be interpreted in the context of your overall findings?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe structure of the paper is logical, and the text is clearly articulated and easy to read.\n\n2.\tThe layout of the tables in the article is reasonable, and the overall formatting of the paper is quite aesthetically pleasing.\n\n3.The paper provides rich experimental results with a sufficient number of figures and tables.", "weaknesses": "1、The innovation of this paper lies solely in the introduction of a Mixture-of-Experts Diffusion Policy, merely combining MOE with diffusion methods without sufficiently demonstrating the novelty of their approach. Additionally, the overall workload appears to be insufficient.\n\n2、In Figure 1, the authors propose a new router strategy, but the figure does not clearly illustrate how this strategy selects different experts, failing to convey the core idea of the paper. Furthermore, the framework diagram is overly simplified and lacks sufficient depth.\n\n3、The methods section of the paper is too brief, and much of the content is drawn from existing research, with the authors not highlighting their contributions effectively.\n\n4、In the experimental section, the authors do not provide a comparison of FLOPs and inference times with other baseline methods, which may lead to unfair comparisons. Moreover, the experimental results for ABC→D in Table 1 are worse than the baseline methods, and the explanations in the paper are not sufficiently convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MNWRwLhI0p", "reviewer": "Reviewer_3nco", "summary": "This paper introduces a MOE diffusion model for embodied AI, which applies a noise-based router for policy diffusion. This MOE structure reduces the active parameters and FLOPs, and also achieves better performance in CALVIN challenges.", "questions": "Please refer to the weakness. My major concern is w1-2 &5. I would like to increase my rating  if authors address my concern, especially w5.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Authors claim that the proposed method increases the performance with fewer parameters and FLOPs.\n2. Good writing. This paper is easy to follow.", "weaknesses": "1. This paper mainly measures the computational and storage costs of diffusion models with FLOPs and parameters. However, latency (response time for one input instance) and memory footprint are more important metrics. I highly recommend authors report them since the reduction in FLOPs and parameters does not always bring real benefits. Besides, I'm really curious about the computation costs of diffusion models in the overall pipeline of embodied AI, since it seems that the diffusion model only contains 8 layers, which are much smaller than the perception model. This will make compression and acceleration for the diffusion model not important for embodied AI.\n2. MOE may bring additional parameters. How do the overall parameters (instead of active parameters) in Table 1.\n3. Some typos. 285, wrong usage for quotes. You should use `` instead.\n4. Table 1 shows that MoDe works better than other methods without pertaining, which is good. However, it can be better to report MoDe with pretraining for a fail comparison with GR-1.\n5. Section 3.3 seems to be a direct adaptation of the classific MOE method. Could you provide a concrete comparison between the MOE utilized in this work and previous works?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)