Sparse Diffusion Policy:  A Sparse, Reusable, and Flexible Policy for Robot Learning

Authors: yixiao  wang
yifei  zhang
mingxiao  huo
ran  tian
xiang  zhang
yichen  xie
chenfeng  xu
pengliang  ji
wei  zhan
mingyu  ding
masayoshi  tomizuka
u c berkeley
moe module moe module
Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-07 07:17:42

Abstract: The increasing complexity of tasks in robotics demands efficient strategies for multitask and continual learning. Traditional models typically rely on a universal policy for all tasks, facing challenges such as high computational costs and catastrophic forgetting when learning new tasks. To address these issues, we introduce a sparse, reusable, and flexible policy, Sparse Diffusion Policy (SDP). By adopting Mixture of Experts (MoE) within a transformer-based diffusion policy, SDP selectively activates experts and skills, enabling task-specific learning without retraining the entire model. It not only reduces the burden of active parameters but also facilitates the seamless integration and reuse of experts across various tasks. Extensive experiments on diverse tasks in both simulators and the real world show that SDP 1) excels in multitask scenarios with negligible increases in active parameters, 2) prevents forgetting in continual learning new tasks, and 3) enables efficient task transfer, offering a promising solution for advanced robotic applications. More demos and codes can be found on our https://anonymous.4open.science/w/sparse_diffusion_policy-24E7/.

Main paper: 
 introduction: 
 Generalist robots are gaining substantial attention in both academia and industry, capable of performing a wide range of tasks and continually learning new ones without losing previously acquired skills [1, 2, 3, 4, 5, 6] . Traditional approaches often rely on a universal and monolithic policy [1, 2] for all tasks, activating all the parameters in the large network for even simple tasks like pushing. Besides, given the diverse nature and lifelong requirements of robot learning tasks [7, 8] , when encountering a new task, these approaches typically require costly fine-tuning [9] , which carries the risk of catastrophic forgetting of previously acquired skills. Task-specific adapters, such as LoRA [10] , entail expanding active parameters during inference. An alternative approach is to train separate policies for different tasks, though this requires independent and from-scratch training for each task and prevents knowledge transfer across tasks.Recent works on skill discovery [5, 11, 12] and chain of skills [13, 14] show promise in addressing the above challenges. These methods necessitate meticulous design with knowledge guidance such as visual features [5, 15, 16, 17, 18] and language prompts [14, 19] , to learn different skills for different tasks, with the goal of reusing these skills in unseen scenarios. However, their skill abstraction modules are typically not scalable and the network structure is not designed to be sparse for efficient computing. As a result, the influence of network structure has not yet been thoroughly explored. Recently, Mixture of Experts (MoE) [20] has proven successful in large-scale applications across NLP, vision, and multimodal domains [21, 22, 23] . It selectively activates only a subset of expert Figure 1 : Overview of Sparse Diffusion Policy (SDP). 1) Multitask Learning: SDP can simultaneously acquire experts from different human demonstration datasets. Due to its sparsity, SDP can activate different experts for different tasks. Additionally, with its reusability, SDP can activate the same expert to share knowledge among tasks. 2) Continual Learning: With its flexibility, SDP can transfer to new tasks by adding only a few new experts to learn the new tasks. This approach mitigates catastrophic forgetting by retaining the old experts and routers. 3) Task Transfer: Leveraging its reusability, SDP can transfer to new tasks by tuning the old experts and routers for expert selection. This allows SDP to acquire new skills based on the previously learned knowledge.networks selected by a router, allows experts to be utilized across various tasks and over time, and facilitates the integration of additional networks while preserving the functionality of existing ones. This observation raises a natural question: Can the mere employment of a sparse, reusable, and flexible MoE structure overcome the challenge without the extensive integration of human-derived knowledge?Motivated by the above observation, we introduce Sparse Diffusion Policy (SDP), as depicted in Figure 1 , a framework for multitask and continual learning by exploring the potential of integrating MoE architecture within a transformer-based diffusion policy [24] . SDP offers several advantages: 1) Sparsity. Only a select set of skills is activated at one time, significantly enhancing computational efficiency during inference. 2) Reusability. Skills are systematically reused across different tasks, for example, "pick and place" is a common skill frequently utilized in robotic tasks. 3) Flexibility. Skills for new tasks can be merged or added to the existing skill pool, enabling their flexible use in future tasks. We conceptualize the experts in MoE as specialized skills and the router as a skill planner (as illustrated in Figure 1 ). Furthermore, we explore specific training and application strategies of MoE for robotic learning.Our extensive experiments in both simulations and real-world settings demonstrate the effectiveness of SDP in multitask, continual, and transfer learning for robotic tasks. It achieves superior multitask performance, with only a 1% increase in active parameters compared to the single-task model. For continual learning, SDP maintains a higher success rate on new tasks without forgetting previously learned tasks, whereas the baseline model [10] requires activating over 62% of its parameters. Furthermore, we investigate the potential for task transfer to complex, long-horizon tasks using a very small pretrained model, initially trained on only two half-length tasks. By training a highly lightweight router (less than 0.4% of the total parameters), SDP outperforms models trained from scratch. Through experiments, we observe that the SDP is capable of extracting a broad range of skills through the combination of experts, and the router functions effectively as a skill planner. 
multitask and continual learning in robotics: 
 In the field of robot learning, significant advancements have been made in multitask [25, 26, 27, 28, 29, 30, 13, 31, 32, 2, 33, 34, 35, 36, 37] and continual learning [38, 39, 40, 41, 5, 42] , allowing robots to efficiently acquire and retain multiple skills over time. Multitask learning approaches, such as policy distillation [43, 44, 45] and hierarchical reinforcement learning [46, 47, 48, 49] , enable robots to learn and perform multiple tasks simultaneously by leveraging shared representations and decomposing complex tasks into manageable subtasks. However, these methods cannot induce sparsity during policy learning, which can enhance the efficiency of the policy network in multitask learning. Continual learning techniques, including regularization-based methods like Elastic Weight Consolidation (EWC) [50] , memory-based strategies like experience replay, and architectural innovations such as Progressive Neural Networks (PNNs) [51] , are developed to mitigate catastrophic forgetting, allowing robots to incrementally acquire new skills while retaining previously learned ones. Also, there are works in meta-learning [52, 53, 54] and few-shot learning [55, 56, 57, 58] , that provide robots with the ability to quickly adapt to new tasks with minimal data. However, the MoE structure can naturally support continual learning without forgetting old tasks due to its unique architecture. This approach requires fewer additional techniques and can seamlessly integrate with multitask learning to create a dynamic task pool.The Mixture of Experts (MoE) approach has seen significant advancements in both computer vision and large language models, offering a promising strategy to enhance model performance by leveraging specialized sub-models as "experts". In computer vision, MoE frameworks have been employed for multitask learning and transfer learning [59, 22, 60, 61, 62] , demonstrating their efficiency in handling diverse and complex datasets such as segmentation, image classification. Moreover, many works [63, 64, 65] integrated MoE into the Transformer architecture, showing substantial gains in natural language processing tasks. These advancements underscore the potential of MoE systems to address the growing demands for computational efficiency [66, 67] and model accuracy [68, 69, 70] in both computer vision and language processing domains. This work focuses on leveraging the sparsity of MoE to conduct multitask and continue learning for the robot learning area. We also make full use of the MoE module to explore the efficient finetuning for task transfer. 
method: 
 Our approach integrates Mixture of Experts (MoE) layers into a transformer-based diffusion policy network [24] , combined with specifically designed training and application strategies for multitask and continual learning in robotics. Owing to the network's structural sparsity, we refer to our method as Sparse Diffusion Policy (SDP). In the subsequent sections, we first outline the problem formulation for multitask and continual imitation learning. We then discuss the integration of the Mixture of Experts (MoE) structure and explore how its sparsity, flexibility, and reusability can be specifically utilized for robot learning. Finally, we present the training strategies we have developed to further unleash the potential in the domain of robot learning.We consider a set of robot tasks C = {T j } J j=1 . For task j, there are N expert demonstrations {τ j,i } N i=1 . Each demonstration τ j,i is a sequence of state-action pairs. We formulate robot imitation learning as an action sequence prediction problem [24, 36] , training a model to minimize the error in future actions conditioned on historical states. Specifically, for task j, imitation learning minimize the behavior clone loss L j bc formulated asL j bc = Es t-o:t+h ,a t-o:t+h ∼T j T t=0L (π(a t:t+h |st-o+1:t, Tj; θ), a t:t+h ) .(EQUATIONwhere a is action, s is state, h is the prediction horizon, o is the number of historical steps, L is a supervised action prediction loss such as mean squared error or negative log-likelihood, T is the length of demonstrations and θ represents the learnable parameters of the network. In a multitask setting (to learn {T j } J-1 j=1 ), the behavior cloning loss is given by the case of continual learning, when only task J is to be learned, we have access exclusively to the corresponding demonstrations T J in this learning cycle, and the behavior cloning loss is L bc = L J bc .L bc = J-1 j=1 L j bc . In Action Head 𝑠 !"# 𝑠 ! 𝑎 ! 𝑎 !$# 𝑎 !$% 𝑎 !$& 𝑎 !$' 𝑎 !$( 𝑎 !$) 𝑠 !$' 𝑠 !$(We utilize the transformer-based diffusion policy [24] and replace the Feed-Forward Networks (FFN) with MoE layers. For the n-th MoE layer [20] , n = {1, 2, ..., N }, there are L experts {E n l } L l=1 and one router R n . Each expert network E n l is composed of multilayer perceptrons (MLP). Router R n compares the input x ∈ R 1×M with the expert embeddings W n ∈ R M ×L and get their scores. Only the Top-K expert networks are activated for inference. Specifically, the MoE layer output y is derived asEQUATION)where Top-K(v, l) is the l-th element of vector v if it is largest K elements otherwise 0.Multitask Learning. Since different tasks require distinct experts for completion, we assign a taskspecific router R n j to enable different tasks to select different experts based on the same historical states. As depicted in Figure 2 , the same experts can be reused at different times within the same task and across various tasks, facilitated by the task-specific router and time-varied state. On the other hand, state-specific and task-specific experts can also be utilized and learned. More importantly, the activation of a limited number of experts demonstrates computational efficiency.Continual Learning. MoE layers facilitate straightforward model expansion and support continual learning [61] . Specifically, for each new task, we freeze the previously learned experts and routers, integrate new trainable experts into each MoE layer and train the corresponding task-specific routers (See Figure 2 ). Catastrophic forgetting is mitigated by fixing previously learned parameters, thus enabling lifelong learning. Upon mastering the new task, the experts related to it are abstracted and become reusable in subsequent learning processes. Moreover, the computational cost remains constant despite the continuous integration of new tasks.Intuitive Interpretation and Task Transfer. Intuitively, each unique combination of the TopK experts within every layer of a Mixture of Experts (MoE) architecture (refer to Figure 2 ) represents a distinct "skill". The routing mechanism acts as a planner for skill chains, selecting specific experts to assemble a skill. In this structure, the number of potential combinations of experts across N layers is given by ( L! (L-K)!K! ) N . This formula suggests the capacity to cover a broad spectrum of diverse skills using a finite set of experts and layers. It also benefits continual learning. For instance, with parameters L = 4, N = 2, K = 1, adding one expert per layer (two in total) generates 9 new combinations. More promisingly, when confronted with an unseen long-horizon task, the inherent 3) and behave well in a novel and even long-horizon task (denoted as T J for convenience, the same in continual learning), enabling fast task transfer (See Figure 2 ).Directly minimizing the behavior clone loss L bc often leads to favoring certain experts and reinforcing their selection through a cycle of increased training and preference [71, 20] . To avoid overload of the expert, Previous studies [20, 63, 72, 73] have incorporated an auxiliary load balancing loss to prevent the overloading of any specific expert. However, in robot learning, certain skills are consistently utilized across various tasks, suggesting that some experts should be frequently engaged. For instance, pick-and-place skills are commonly required in numerous robotic tasks, resulting in excessive activation of the combination of experts related to the pick-and-place skill. Consequently, expert overload is typical in multitask robot learning, whereas task overload is atypical due to the distinct nature of each task. Each task includes specific components that necessitate unique skills for successful completion. Inspired by this observation, we propose encouraging experts to specialize in specific tasks [22, 61] to prevent task overload. Specifically, we would like to maximize the mutual information between the task T and the expert E n for each MoE layer:EQUATIONwhere we assume that each task is equally important, i.e., p(T j ) = 1 J . Implementation details for I(T , E n ) are presented in Appendix A. Thus, the total training objective isEQUATIONwhere L bc = J-1 j=1 L j bc in multitask learning, and L bc = L J bc in continual learning and task transfer.4 Experimental Results 
multitask learning: 
 In this section, we discuss sparsity and reusability of our SDP, highlighting its ability to simultaneously acquire multiple experts efficiently while sharing experts across tasks. As shown in Figure 3 . We evaluate the performance of SDP on 8 tasks in Mimicgen [74] . Mimicgen includes 1K-10K human demonstrations per task with broad initial state distributions, effectively showing the generalization for multitask evaluation. To our knowledge, this is the first work to explore multitask training on the Mimicgen benchmark. We choose task-conditioned diffusion (TCD) [76, 19] , fine-tuning the action head (TH), fine-tuning last three transformer layers (TT w/ 3Layer), Octo [77] as our baselines. TCD and Octo require the activation of all network parameters, exhibiting a dense structure rather than a sparse one (Ours). We train our standard SDP model as well as a smaller version, referred to as Light SDP.Each model is trained for 300 epochs using one A6000 GPU for 130-150 hours and evaluated every 50 epochs. We report the average success rate of the best three checkpoints in Table 1 . Thanks to the sparsity, our SDP outperforms all baselines with the same level of active parameters in the policy network. Additionally, our Light SDP also surpasses all baselines with significantly fewer active parameters, while maintaining the same level of total parameters. Another observation is that sparse models (SDP, TH, TT) outperform dense models (TCD and Octo). This suggests that dense structures may hinder the learning of distinct policies across tasks, whereas sparse structures, which allocate separate parameters for different tasks, facilitate the learning of more diverse actions. Task-expert frequency map in Figure 4 shows that each task activates only a subset of the experts, contributing to computational efficiency. More importantly, the map reveals that different tasks can activate the same expert, demonstrating the reusability of experts across various tasks.3D Simulation Results. We evaluate our approach on 3 tasks in DexArt [78] and 3 tasks in Adroit [79] . Specifically, we follow [75] and integrate the Mixture of Experts into the Feed-Forward Network blocks of the 3D diffusion policy network [75] . The number of active parameters is set to be equivalent to the original. We train for 6000 epochs and evaluate every 200 epochs. We report the average success rate of the best three checkpoints. As shown in Table 2 , our method outperforms the baseline, showing the effectiveness of our SDP in 3D perception settings.Figure 5 : Multitask evaluation on real-world tasks.Real Robot Experiment Results. We conduct real robot experiments on FANUC LRMate 200iD/7L robotic arm outfitted with an SMC gripper. We choose three diverse and universal: pulling a circle, picking and placing a cup, and hanging a cup. For the first two tasks, we collect 20 demonstrations each, and for the last task, we collect 40 demonstrations with two distinct hanging sticks. More details can be found in the Appendix B.2. The robot is manipulated using admittance control [80] , which can achieve compliant robot motion to ensure safety during manipulation. The number of training epoch is 2000. As shown in Figure 5 , our method greatly outperforms TCD [76, 19] . Additionally, we found that TCD cannot clearly distinguish the multimodal action distributions across different tasks. We hypothesize that this may be caused by insufficient sparsity. More details and visualizations can be found in the Appendix B.2. In this section, we evaluate the performance of the Sparse Diffusion Policy (SDP) in continual learning, with a focus on the framework's flexibility. When integrating a new task, we add a small number of experts and a new router to the Mixture of Experts (MoE) structure, while freezing the existing experts, routers, and other non-MoE modules. We then train only the newly added router and experts, enhancing the efficiency of learning new tasks and preventing catastrophic forgetting [81] of previously learned tasks. Assume we continuously learn three tasks, a → b → c. We first train the model on task a (Stage 1), then on task b (Stage 2), and finally on task c (Stage 3), evaluating the performance on both the current and previous tasks.Comparison with baselines. We choose three tasks from robimimic [82] for evaluation: Can → Lift → Square. First, we compare our approach with LoRA [10] and a fully fine-tuned (FFT) version of our method. Each method is trained for 500 epochs and evaluated every 50 epochs. We report the average success rate of the best three checkpoints. The results, presented in Table 3 , demonstrate that our method generally exhibited superior performance in both new tasks and previous tasks. Conversely, FFT exhibited significant forgetting of previously learned tasks when learning new ones, whereas LoRA struggled with a substantial increase in the number of active parameters. These findings underscore the flexibility of our SDP.Figure 6 : Ablation Study on MI loss. We report the success rate for the new tasks, as the performance on previous tasks remains the same thanks to our SDP structure.Ablation study on continual learning. In this section, we report ablation study on mutual information (MI) loss in Section 3.3. We selected three tasks from MimicGen [74] : Stack → Hammer Cleanup → Coffee, where task complexity progressively increases (e.g., the Coffee task requires not only picking and placing but closing the coffee cap.). Since our SDP avoids catastrophic forgetting [81] , we report only the performance of the newly added tasks in Figure 6 . The results emphasize the significance of MI loss. We argue that the router initially favors previously trained experts, as newly added experts, being untrained, produce random actions, reinforcing reliance on frozen experts. However, MI loss encourages the router to select task-specific experts, promoting the use of seldom-selected experts. Additional ablation studies are presented in Appendix C. 3 , where experiments further demonstrate the critical role of comprehensive MoE structures and MI loss in achieving optimal performance in complex continual learning settings.In this section, we evaluate the reusability of our SDP, highlighting its ability to efficiently transfer to new tasks by learning how to leverage previously acquired experts (router tuning). Additionally, we visualize the evolution of expert scores to assess the role of prior skills in enabling the transfer to new tasks. When the robot interacts with the coffee machine (not encountered in Mug Cleanup), experts trained on the Coffee achieve higher scores and are activated, offering interpretable insights into the SDP. Also, our model only requires less than 0.1M (0.5%) trainable parameters for efficient task transfer.task transfer experiment. We first train our SDP on base tasks, then freeze all experts and fine-tune only the router for the new task. Specifically, we use Coffee and Mug Cleanup as the base tasks and Coffee Preparation as the new task, all from MimicGen [74] . Coffee Preparation is a long-horizon, partially unseen task requiring unique skills, such as moving the mug to the drip tray. After pretraining on the base tasks, we fine-tune only the MoE router and the vision encoder for the new task. For comparison, we also train a model with the same structure from scratch. Both models are trained for 100 epochs. As shown in Table 4 , our method uses less than 0.5% of the parameters compared to the train-from-scratch approach while achieving better performance, demonstrating that the experts learned from the base tasks effectively cover diverse skills and offer strong representation for distinct tasks. Skills visualisation. As shown in Figure 7 , we visualize the evolution of expert scores on the Coffee, Mug Cleanup, and Coffee Preparation tasks. We observe that when knowledge from Coffee is required but not present in Mug Cleanup, the experts trained on Coffee are activated. Four typical examples are illustrated on the right of Figure 7 . By composing these experts, SDP can acquire new skills, such as moving the mug to the coffee machine's drip tray. Additionally, we find that some experts are consistently activated in Coffee Preparation, while others are seldom used. This insight suggests the possibility of policy distillation, which we leave for future work. More experiment details and explanations are presented in Appendix D. 
discussion and limitation: 
 In this paper, we have introduced the Sparse Diffusion Policy (SDP) framework, which integrates Mixture of Experts (MoE) layers into the diffusion policy. Our approach is designed around three key principles: sparsity, flexibility, and reusability. By activating only relevant portions of the network for specific tasks, SDP can induce the sparsity for efficient multitask learning. The flexibility of our algorithms is capable of acquiring new tasks without forgetting existing skills, while the reusability of existing knowledge allows for efficient multitask and task transfer learning. Our model has great potential for future large-scale robot learning.Limitation. Our SDP may fail if the shared knowledge in the network is too limited but same experts are activated by different routers. Additionally, the router in SDP is task-specific, which hinders the ability for universal task completion. By combining it with large language models, the SDP conditioned on language with a broader policy can be applied to a wider range of robot learning scenarios in the future. In this section, we conduct additional continual learning experiments using a fully fine-tuned visual encoder. As shown in Table 5 , after transferring to the final task (Square), our model outperforms the LoRA method in terms of transferability with the fully fine-tuned visual encoder. However, fully fine-tuning the visual encoder results in complete forgetting of previous tasks. Therefore, utilizing a comprehensive set of MoEs is essential for effective continual learning.In this section, we present ablation study results on three easy tasks: Can → Lift → Square. Can involves picking and placing a can, Lift focuses on lifting a cube, and Square requires inserting a square peg into the correct post. We report the performance on the Square task at Stage 3 in Figure 10 . These results highlight the importance of employing a comprehensive set of MoEs and the Mutual Information loss to achieve optimal performance in complex continual learning scenarios. Compared to the results in Figure 6 , we observe that in more challenging continual learning settings (Stack → Hammer Cleanup → Coffee), Mutual Information loss plays a more significant role in learning new tasks.Recall that Coffee and Mug Cleanup are the base tasks, Coffee Preparation is the new task. In Coffee, the robot is required to place the pod into the holder and close it. For Mug Cleanup, the robot must open a drawer, place the mug inside, and close the drawer. Coffee Preparation involves placing the mug into the drip tray, opening the holder, and the drawer, placing the pod into the holder, and then closing it. This task can be seen as a composite of Coffee and Mug Cleanup, but it includes unique actions, such as moving the mug to the drip tray, which are not present in the initial training tasks.Our task is to leverage the frozen experts developed from the previous two tasks and learn taskspecific routers that combine these experts to acquire new skills (e.g., moving the mug to the drip tray). To illustrate the router's functionality, we aim to activate different experts for each base task. Thus, we set eight experts per layer and activate only the top two. We visualize Expert Score versus Timesteps in the final transformer layer at the last fifth diffusion timestep in Figure 7 . We observe that Coffee primarily activates experts 0127, while Mug Cleanup activates experts 3456, indicating no shared experts between these tasks. This allows us to identify which base task's knowledge is being used for the new task, Coffee Preparation.We observe that experts trained on the Coffee task are more likely to be selected when the task requires information specific to the coffee machine, which does not appear in the Mug Cleanup task. For example, Coffee-related experts are activated when the action involves placing the pod. This observation highlights that the router functions as a skill planner, with experts serving as individual skills. The router effectively composes these skills across tasks to tackle complex and previously unseen tasks.Interestingly, we find that new skills can emerge through the combination of experts from different tasks. For example, when the robot moves the mug to the coffee machine's drip tray-a new skill-it must determine the coffee machine's location, triggering the activation of Coffee-related experts. In most cases, however, experts from Mug Cleanup are activated, indicating that their combination suffices for the majority of Coffee Preparation, except for actions that specifically require Coffee-related information. This demonstrates the broad skill coverage achieved through expert composition, underscoring the expressive power of our SDP. Additionally, we observe that some experts are seldom activated. Given the frequency of activation across tasks, we can prune, merge and enlarge the experts, suggesting a potential connection between our SDP and policy distillation, which we leave for future work.Based on the discussion above, aligning with the intuitive explanation in Section 3.2, the experts could be viewed as a pool of skills, with the router functioning as a chain planner, enabling our SDP to efficiently transfer knowledge to new tasks. 
 
 Appendix:  


Decision: The authors have provided thorough responses with additional experiments  to the questions and concerns of the reviewers. The final scores are 1 weak accept, 1 accept and 1 strong accept
------
The paper "Sparse Diffusion Policy: A Sparse, Reusable, and Flexible Policy for Robot Learning" presents an innovative approach to multitask and continual learning in robotics using a mixture of experts within a diffusion policy. The strengths of this work lie in its novel idea and comprehensive experimentation demonstrating its feasibility and potential. However, the paper has several weaknesses, including concerns about memory consumption, the limited variety of tasks in experiments, and insufficient clarity in certain methodological details, making the experimental results less convincing.

Reviews: review: title: title: A promising idea for multi-task and continual robot learning with difussion policy
paper_summary: paper_summary: The authors propose to adopt a Mixture-of-Expert layer in the diffusion policy for high computation efficiency and continual learning. The introduced model is called sparse diffusion policy, featured by its sparsity and reusaibility for multi-task learning, continual learning and efficient task transfer. Experiments in both 2D/3D simulation and real robot verify the claimed properties.
main_review: main_review: ### Strengths
- The problem tackled by this work is of great importance in the imitation learning/robot learning community;
- The proposed idea of equipping the Mixture-of-Experts (MoE) layer to the policy network is highly sensible and feasible;
- The presentation of this work is clear and easy to follow;
- The experiments are comprehensive, covering all the claimed properties of the model.

### Weaknesses/ Suggestions
 - The first question is about the claimed **sparsity**. In experiments such as Tables 1 and 2, the proposed method has similar active parameters with the baselines, while active parameters denote only a subset of the experts in the proposed model and the whole one for the baseline models. In this sense, the parameters of the whole model (all experts) to be stored would be much more than the baselines. Does this cause the problem of high memory consumption? Should a proper comparison be made with the same amount of the **overall** model for both the proposed method and the baselines? 
 -  The tasks evaluated in the experiments are insufficient in terms of amount and task diversity [1, 2]. There are only 7 tasks in 2D simulation, 3 tasks in 3D simulation, 3 tasks on real robots for multi-task learning, and 3 similar tasks for continual learning. With the marginal improvement in the success rate, the experimental results are less convincing.   
 - To better support the claim in the paper, additional baselines of open-sourced large models for robot learning such Octo [1] can be evaluated with the proposed idea. The reason behind this is that the diffusion policy is known to work better for single tasks than multi-tasking. To enable diffusion policy to be a multi-task learner, Octo with diffusion policy as its action head is one of the popular models for such purpose;
 - In the continual learning experiment, does the size of the newly added expert and router ( number of active parameters) influence the performance?  
 - Figure 5 is hard to read. The minor improvement of having mutual information loss renders it less necessary. Moreover, the three tasks here, e.g., "Lift" and "Square," share similar skills with "Can," such as picking and lifting, etc. That means the claimed effect of mutual loss to enforce specific tasks to specific experts is less sensible.
 - The last question is for the task transfer. As task transfer fine-tuning is similar to continual learning except for the frozen experts, it would be less confusing for readers if such design differences could be better motivated and explained.    

[1] Team, Octo Model, et al. "Octo: An open-source generalist robot policy." arXiv preprint arXiv:2405.12213 (2024).

[2] Liu, Shiqi, et al. "Continual vision-based reinforcement learning with group symmetries." Conference on Robot Learning. PMLR, 2023.
questions: questions: See the list of weaknesses above.
limitations: limitations: 2
review_summary: review_summary: Though this is an interesting idea, I vote for weak reject due to the less convincing experimental results.
title: title: Appealing idea but results not convincing
paper_summary: paper_summary: This paper presents a technique named Sparse Diffusion Policy (SDP) for multitask and continual learning of manipulation skills. SDP utilizes a transformer-based diffusion policy where the final layers of the model are composed of Mixture-of-experts layers, where each expert is an MLP. Task-specific trained routers activate different combinations of experts for different tasks. The reuse of experts across tasks leads to the reuse of learned skills without unbounded growth in parameters. Experimental results include multitask and continual learning scenarios in simulation and the real world.
main_review: main_review: I understand the appeal of the proposed technique and agree that using a mixture of experts can lead to the reuse of skills for multitask/continual learning. Although MoEs with transformer-based architectures have been used in other domains, it is not trivial to transfer this to robotics, and the proposed SDP framework can be potentially useful for the robot learning community. However, there are several open questions and I think the clarity of the paper needs to be improved, especially in the part where SDP is described. The biggest cause of concern for me is related to the significance of some of the results, where the difference between the proposed method and other baselines is not very clear (see rebuttal questions below).

*Strengths*:
1. Appealing idea. MoE is potentially a good way to acquire new manipulation skills without catastrophic forgetting and unbounded parameter growth.
2. Addresses both multitask and continual learning scenarios

*Weaknesses*:
1. The paper lacks clarity in several places
2. Results do not clearly show that the proposed method is better than other baselines
3. Many open questions
questions: questions: 1. Line 76: Please clarify what is meant by "additional techniques"

2. Line 115: Please explain how the Router $\mathcal{R}^n$ compares the input $x$ with the expert embeddings $W^n$. How are the embeddings trained?

3. Fig. 2(a) is confusing. Some questions related to this:
(i) What do the colors (blue, green, red) denote?
(ii) Does each rectangle denote a single MLP expert?
(iii) I assume each row of boxes (green border) is a single MoE layer. If an MoE layer produces a single output, what is the significance of the outputs between blue and blue and between green and green boxes?

4. Please state in the caption what the flame symbol signifies in Figs 2(b) and (c)

5. Eq. (2) does not match Eqs. 2, 3 from [20]. Please explain the difference.

6. Line 127 states "[we] integrate new trainable experts into each MoE layer and train the corresponding task-specific routers". Do you have multiple routers in each MoE layer or only a single router per layer (as shown in Fig. 2)?

7. Line 131: In what sense does the computational cost remain constant (storage, FLOPS, or something else)? As you add some experts and a router for each task, does the cost remain **nearly** constant or absolutely constant?

8. In equation (3), how is $p(\mathcal{T}_j, \mathcal{E}^n_l)$ computed? 

9. Please consider commenting on the suitability of the baselines used in the experiments (i.e. why did you choose the baselines used in the paper?). 

10. In Table 1, the average of the best three checkpoints is presented. Is each checkpoint an independent trial? Why do you use the best three and not the average over all checkpoints?

11. In the results reported in Tables 1, 2, and 4, the difference between the proposed method and other baselines is quite small. These results are the average over the three best trials (as far as I understood) and the authors do not report the standard deviation. Given the small number of checkpoints, the difference between the methods may not be significant. It is difficult to be convinced that the proposed method is superior to the baselines without considering more checkpoints and looking at the average and SD, or conducting explicit statistical tests.

12. In Table 1, does the phrase "active parameters" refer to the parameters belonging to the experts selected by the router (plus, other shared parameters)? Please also report the total parameter size.

13. Line 180: A quantitative comparison of computational efficiency would help the reader. If I understand correctly, the active parameters of the sparse method (SDP) are nearly the same as the other non-sparse baselines. Why then is SDP more efficient computationally? 

14. Please consider commenting on the reason for the very long training times needed by your models. Do all the baselines need to be trained for similar durations?

15. For the ablation study (line 223), why do you only report the results for the last task? Since this is a CL ablation, the results of all tasks should be reported. If space is a constraint, consider reporting the full results in the supplementary material/appendix.

Some minor comments: 

1. Many citations point to the Arxiv versions of published papers. Please change the citations to the published versions. Some examples (this is not an exhaustive list):
- Ajay et al. Is conditional generative modeling all you need for decision-making - Published in ICLR 2023
- Liang et al., Skilldiffuser: Interpretable hierarchical planning via skill abstractions in diffusion-based task execution - Published in CVPR 2024
- Chi et al., Diffusion policy: Visuomotor policy learning via action diffusion - Published in RSS 2023

2. Tables 1,2,3 use "SDP (ours)" and Table 4 uses "MoE (ours)". Please be consistent.

3. In the caption of Table 1, "activate" should be "active"

4. Line 88: "continue" should be "continual"
limitations: limitations: 3
review_summary: review_summary: The idea of using MoE layers to avoid catastrophic forgetting while continually learning manipulation skills is appealing. However, the results do not fully support the authors' claims about the superiority of their approach. Moreover, several questions remain unanswered by the paper in its current form.
title: title: A very clear contribution overall, requires minor improvements
paper_summary: paper_summary: The paper proposes a sparse (computationally efficient) diffusion policy, which is flexible in its knowledge addition and adaptation, and reusable by simply switching the skill selection model.
main_review: main_review: The paper proposes a sparse (computationally efficient) diffusion policy, which is flexible in its knowledge addition and adaptation, and reusable by simply switching the skill selection model. The work would be stronger if it presented comparisons with more standard baselines. 

Strengths: 

- The authors have released the code and demonstrations used, significantly strengthening the reproducibility of the method’s results, as well as making the work more convincing.  

- Strong results supporting all claims made by the paper.  

- The writing is clear and concise. 

Weaknesses. 

- Brief related works section does not cover many relevant works. 

- The future works section is lacking. 

- Some method details are unclear, e.g. how are the expert embedding obtained. While the code is provided, based on the paper alone, there might not be enough information provided to reproduce the work. 

- Justification of the diffusion-based model is insufficient. 

- While the chosen baselines are good, it would be better to include further baselines from multi-task learning and continual learning.
questions: questions: Other comments: 

Very good strong motivation and introduction. 

Fig 2. It’s unclear what the gray blocks represent exactly. 

The comparison with LoRA does not show significant difference between the proposed method and LoRA, while L218 claims that it “struggled”. The results presented do not seem to support that statement.  

Fig. 6. It is unclear what the heatmap axes are. 

While diffusion policies have gained traction in recent years, it is unclear why 3D diffusion policy was chosen as the particular – could any model be used as the backbone? Could there be a mix of differently structured experts? Is “diffusion” core to the contribution? If yes, it asks for more justification and related works. If not, it is unclear why the contribution is focused around it. 

The reasoning behind some of the references is unclear – Table2 / Fig 3. refers to previous works, but do the depicted tasks originate from those works?  

One strength of the method which is not disccussed is interpretability. The skill-based router model allows for more interpretability of the model’s behaviour – given specialized experts it would be clear what the model “intends”.
limitations: limitations: 3
review_summary: review_summary: I recommend the work is accepted due to clear strengths. However, during the rebuttal phase I could be swayed based on author responses and reviewer comments.