Variational Distillation of Diffusion Policies into Mixture of Experts

Authors: hongyi  zhou
denis  blessing
ge  li
‡  onur celik
xiaogang  jia
gerhard  neumann
rudolf  lioutikov
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 12:27:57

Abstract: This work introduces Variational Diffusion Distillation (VDD), a novel method that distills denoising diffusion policies into Mixtures of Experts (MoE) through variational inference. Diffusion Models are the current state-of-the-art in generative modeling due to their exceptional ability to accurately learn and represent complex, multi-modal distributions. This ability allows Diffusion Models to replicate the inherent diversity in human behavior, making them the preferred models in behavior learning such as Learning from Human Demonstrations (LfD).
However, diffusion models come with some drawbacks, including the intractability of likelihoods and long inference times due to their iterative sampling process. The inference times, in particular, pose a significant challenge to real-time applications such as robot control.
In contrast, MoEs effectively address the aforementioned issues while retaining the ability to represent complex distributions but are notoriously difficult to train.
VDD is the first method that distills pre-trained diffusion models into MoE models, and hence, combines the expressiveness of Diffusion Models with the benefits of Mixture Models.
Specifically, VDD leverages a decompositional upper bound of the variational objective that allows the training of each expert separately, resulting in a robust optimization scheme for MoEs.
VDD demonstrates across nine complex behavior learning tasks, that it is able to: i) accurately distill complex distributions learned by the diffusion model, ii) outperform existing state-of-the-art distillation methods, and iii) surpass conventional methods for training MoE. The code and videos are available at https://intuitive-robots.github.io/vdd-website.

Main paper: 
 introduction: 
 Diffusion models [1] [2] [3] [4] have gained increasing attention with their great success in various domains such as realistic image generation [5] [6] [7] [8] . More recently, diffusion models have shown promise in Learning from Human Demonstrations (LfDs) [9] [10] [11] [12] [13] . A particularly challenging aspect of LfD is the high variance and multi-modal data distribution resulting from the inherent diversity in human behavior [14] . Due to the ability to generalize and represent complex, multi-modal distributions, diffusion models are a particularly suitable class of policy representations for LfD. However, diffusion models suffer from several drawbacks such as long inference time and intractable likelihood calculation. Many diffusion steps are required for high-quality samples leading to a long inference time, limiting the use in real-time applications such as robot control, where decisions are needed at a high frequency. Moreover, important statistical properties such as exact likelihoods are not easily obtained for diffusion models, which poses a significant challenge for conducting post hoc Figure 1 : VDD distills a diffusion policy into an MoE. LfD is challenging due to the multimodality of human behaviour. For example, tele-operated demonstrations of an avoiding task often contain multiple solutions [13] . Lower: A diffusion policy can predict high quality actions but relies on an iterative sampling process from noise to data, shown as the red arrows. Upper: VDD uses the score function to distill a diffusion policy into an MoE, unifying the advantages of both approaches.optimization such as fine-tuning through well-established reinforcement learning (RL) approaches like policy gradients or maximum entropy RL objectives.A well-studied approach that effectively addresses these issues are Mixture of Experts (MoE).During inference, the MoE first selects an expert that is subsequently queried for a forward pass. This hierarchical structure provides a fast and simple sampling procedure, tractable likelihood computation, and the ability to represent multimodal distributions. These properties make them a well-suited policy representation for complex, multimodal behavior. However, training Mixture of Experts (MoEs) is often difficult and unstable [15] . The commonly used maximum likelihood objective can lead to undesired behavior due to mode-averaging, where the model fails to accurately represent certain modes. Yet, this limitation has been alleviated by recent methods that use alternative objectives, such as reverse KL-divergence, which do not exhibit mode-averaging behavior [14, 16] .To obtain the benefits of both models, i.e., learning highly accurate generative models with diffusion and obtaining simple, tractable models using a mixture of experts, this work introduces Variational Diffusion Distillation (VDD), a novel method that distills diffusion models to MoEs. Starting from the variational inference objective [17, 18] , we derive a lower bound that decomposes the objective into separate per-expert objectives, resulting in a robust optimization scheme. Each per-expert objective elegantly leverages the gradient of the pre-trained score function such that the MoE benefits from the diffusion model's properties. The resulting MoE policy performs on par with the diffusion model and covers the same modes, while being interpretable, faster during inference, and has a tractable likelihood. This final policy is readily available to the user for post hoc analysis or fast fine-tuning for more specific situations. A high-level architecture of the VDD model and its relation to the diffusion policy is shown in Fig. 1 . VDD is thoroughly evaluated on nine complex behavior-learning tasks that demonstrate the aforementioned properties. As an additional insight, this paper observed that one-step continuous diffusion models already perform well, a finding not discussed in prior work.In summary, this work presents VDD, a novel method to distill diffusion models to MoEs, by proposing a variational objective leading to individual and robust expert updates, effectively leveraging a pre-trained diffusion model. The thorough experimental evaluation on nine sophisticated behavior learning tasks show that VDD i) accurately distills complex distributions, ii) outperforms existing SOTA distillation methods and iii) surpasses conventional MoE training methods. 
related works: 
 Diffusion Models for Behavior Learning. Diffusion models have been used in acquiring complex behaviors for solving sophisticated tasks in various learning frameworks. Most of these works train diffusion policies using offline reinforcement learning [19] [20] [21] [22] [23] [24] , or imitation learning [9, 12, 11, 10, 25] . In contrast, VDD distills diffusion models into an MoE policy to overcome diffusion-based policy drawbacks such as long inference times or intractable likelihoods instead of optimizing policies directly from the data.Mixture of Experts (MoE) for Behavior Learning. MoE models are well-studied, provide tractable likelihoods, and can represent multi-modality which makes them a popular choice in many domains such as in imitation learning [14, 26-31, 16, 13] , reinforcement learning [32] [33] [34] [35] [36] [37] [38] and motion generation [39] to obtain complex behaviors. Although VDD also uses an MoE model, the behaviors are distilled from a pre-trained model using a variational objective and are not trained from scratch.The empirical evaluation demonstrates that VDD's stable training procedure results in improved performance compared to common MoE learning techniques.Knowledge Distillation from Diffusion Models Knowledge distillation from diffusion models has been researched in various research areas. For instance, in text-to-3D modeling, training a NeRFbased text-to-3D model without any 3D data by mapping the 3D scene to a 2D image and leveraging a text-to-2D diffusion is proposed [7] . The work proposes minimizing the Score Distillation Sampling (SDS) loss that is inspired by probability density distillation [40] and incentivizes the 3D-model to be updated into higher density regions as indicated by the score function of the diffusion model.To overcome drawbacks such as over-smoothing and low-diversity problems when using the SDS loss, variational score distillation (VSD) treats the 3D scene as a random variable and optimizes a distribution over these scenes such that the projected 2D image aligns with the 2D diffusion model [8] . In a similar context, the work in [41] proposes distilling a trained diffusion model into another diffusion model while progressively reducing the number of steps. However, even though the number of diffusion steps is drastically reduced, a complete distillation, i.e. one-step inference as for VDD is not provided. Additionally, the resulting model suffers from the same drawbacks of diffusion models such as intractable likelihoods. In contrast, in consistency distillation (CD), diffusion models are distilled to consistency models (CM) [42] [43] [44] [45] such that data generation is possible in one step from noise to data. However, one-step data generation typically results in lower sample quality, requiring a trade-off between iterative and single-step generation based on the desired outcome. As CMs, VDD performs one-step data generation but distills the pre-trained diffusion model to an MoE which has a tractable likelihood and is efficient in inference time. The experimental evaluations show the advantages of VDD over CMs. Diff-Instruct [46] proposes a two-step framework for distilling diffusion models into implicit generative models, whereas VDD considers an explicit generative model where the model's density can be directly evaluated. In addition, Diff-Instruct requires training an auxiliary diffusion model, while VDD only optimizes a single model. Score Regularized Policy Optimization (SRPO) [47] also leverages a diffusion behavior policy to regularize the offline RL-based objective. However, in contrast to SRPO, VDD learns an MoE policy instead of a uni-modal Gaussian policy and explicitly distills a diffusion model instead of using it as guidance during optimization. Furthermore, VDD trains MoEs policies in imitation learning instead of reward-labeled data as in offline RL. A concurrent work, EM-Distillation (EMD) [48] , introduces an EM-style distillation objective derived from the mode-covering forward KL divergence. In contrast, VDD proposes an EM-style objective based on the mode-seeking reverse KL but encourages mode-covering behavior by having multiple experts. 
preliminaries: 
 Here, we introduce the notation and foundation for Denoising Diffusion and Mixture of Experts policies. Throughout this work, we assume access to samples from a behavior policy π * and the corresponding state distribution µ, that is a ∼ π * (•|s) and s ∼ µ(•), respectively.Denoising Diffusion Policies. Denoising diffusion policies employ a diffusion process to smoothly convert data into noise. For a given state s ′ , a diffusion process is modeled as stochastic differential equation (SDE) [3] EQUATION) with drift f , diffusion coefficient g(t) and Wiener process w t ∈ R d . The solution of the SDE is a diffusion process (a t ) t∈[0,T ] with marginal distributions π * t such that π T ≈ N (0, I) and π 0 = π * . [49, 50] showed that the time-reversal of Eq. 1 is again an SDE given by dat = f (a t , t) -g 2 (t)∇ at log π * t (a t |s ′ ) dt + g(t)d wt , a T ∼ N (•|0, I).(2) Simulating the SDE generates samples from π * (•|s ′ ) starting from pure noise. For most distributions π * , however, we do not have access to the scores (∇ at log π * t (a t |s ′ )) t∈[0,T ] . The goal of diffusion- based modeling is therefore to approximate the intractable scores using a parameterized score function, i.e., f θ (a, s, t) ≈ ∇ log π * t (a|s). To that end, several techniques have been proposed [51, 52] , allowing for sample generation by approximately simulating Eq. 2. The most frequently employed SDEs in behavior learning are variance preserving (VP) [2, 10] and variance exploding (VE) [9] . For further details on diffusion-based generative modeling, we refer the reader to [3, 4] . While we only consider VE and VP in this work, VDD can be applied to any score-based method.Gaussian Mixtures of Expert Policies. Mixtures of expert policies are conditional discrete latent variable models. Denoting the latent variable as z, the marginal likelihood can be decomposed as q ϕ (a|s) = z q ξ (z|s)q νz (a|s, z),where q ξ (z|s) and q νz (a|s, z) are referred to as gating and experts respectively. ξ and ν z denote the gating and expert parameters and thus ϕ = ξ ∪ {ν z } z . The gating is responsible for soft-partitioning the state space into sub-regions where the corresponding experts approximate the target density.To sample actions, that is, a ′ ∼ q ϕ (•|s ′ ) for some state s ′ , we first sample a component index from the gating, i.e., z ′ ∼ q ξ (•|s ′ ). The component index selects the respective expert to obtain a ′ ∼ q νz (•|s ′ , z ′ ). For Gaussian MoE the experts are chosen as q νz (a|s, z) = N (a|µ νz (s), Σ νz (s)), where µ νz , Σ νz could be neural networks parameterized by ν z . From the properties of Gaussian distributions, it directly follows that this model class admits tractable likelihoods and fast sampling routines. Furthermore, given enough components, Gaussian MoEs are universal approximators of densities [53] , which makes them a good representation for distillation of diffusion policies. 
variational distillation of denoising diffusion policies: 
 In this section, we outline the mathematical formulations of the VDD model. Detailed descriptions of the model architecture and algorithms can be found in Appendix A. We aim to distill a given diffusion policy π(a|s) by using a different policy representation q ϕ with parameters ϕ. This is useful, e.g., if q ϕ (a|s) has favorable properties such as likelihood tractability or admits fast inference schemes.Assuming that we can evaluate π point-wise, a common approach is to leverage variational inference (VI) to frame this task as an optimization problem by minimizing the reverse Kullback-Leibler (KL) [54] divergence between q ϕ and π, that is,EQUATIONfor a specific state s ′ . To obtain a scalable optimization scheme, we combine amortized and stochastic VI [55] . The former allows for learning a conditional model q ϕ (a|s) instead of learning a separate q ϕ for each state, while the latter allows for leveraging mini-batch computations, that is,EQUATIONwith batch size M ≤ N . Thus, J(ϕ) can be minimized using gradient-based optimization techniques with a gradient estimator such as reinforce [56, 57] or the reparameterization trick [58] . Note that, while the states are sampled from the given data set of the behavior policy, the actions needed to evaluate the KL are generated using our estimated model q ϕ (a|s i ). Yet, there are two difficulties to directly apply this scheme in distilling a diffusion model into an MoE: i) we are not able to evaluate the likelihood π(a|s) of a diffusion model, ii) training of MoE models is notoriously difficult [15] . We will address these two issues in Section 4.1 and Section 4.2, respectively.Although we cannot directly evaluate the likelihood of the diffusion policy π(a|s), we have access to its score functions ∇ at log π t (a t |s) = f θ (a t , s, t), where t ∈ [0, T ] is the diffusion time step. In practice, we would like to evaluate the score in the limit of t → 0 as ∇ a log π * (a|s) ≈ lim t→0 f θ (a t , s, t).Yet, this might lead to an unstable optimization [59] as this score is often not estimated well throughout the action space, and, hence other diffusion time-step selection processes are needed [47] . For now, we will omit the diffusion time-step for the sake of simplicity and refer to Section 4.3 for a detailed discussion about time-step selection. Moreover, we will, for now, assume that the parametrization of q ϕ is amendable to the reparameterization trick [58] . In this case, we can express a ∼ q ϕ (•|s) using a transformation h ϕ (ϵ, s) with an auxiliary variable ϵ ∼ p(•) such that a = h ϕ (ϵ, s). We then express the gradient of J w.r.t. ϕ asEQUATIONUsing the chain rule for derivatives, it is straightforward to see thatEQUATIONAs ∇ a log π(a|s i ) can be replaced by the given score of the pre-trained diffusion policy, we can directly use of VI for optimizing J without evaluating the likelihoods of π.To distill multimodal distributions learned by diffusion models, we require a more complex family of distributions than conditional diagonal Gaussian distributions, which are commonly used in amortized VI. We will therefore use Gaussian mixture of experts. To that end, we construct an upper bound of J which is decomposable into single objectives per expert, allowing for reparameterizing each expert individually and therefore avoiding the need for techniques that perform reparameterization for the entire MoE [60, 61] . The upper bound U (ϕ, q) can be obtained by making use of the chain rule for KL divergences [62, 63, 15] , i.e.,EQUATIONwhere q is an arbitrary auxiliary distribution and upper bound U (ϕ, q) = E µ(s) [E q ξ (z|s) [E q νz (a|s,z) [log q νz (a|s, z) -log π(a|s) -log q(z|a, s)]EQUATIONwith U s z (ν z , q) being the objective function for a single expert z and state s. For further details see Appendix B. Since the expected KL term on the right side of Eq. 8 is always positive, it directly follows that U is an upper bound on J for any q. This gives rise to an optimization scheme similar to the expectation-maximization algorithm [64] , where we alternate between minimization (M-Step) and tightening of the upper bound U (E-Step), that is,min ϕ U (ϕ, q)and min q E µ(s) E q ϕ (a|s) D KL q ϕ (z|a, s)∥q(z|a, s) ,respectively. Please note that q is fixed during the M-Step and ϕ during the E-Step. In what follows, we identify the M-Step as a hierarchical VI problem and elaborate on the E-Step.Step for Updating the Experts. The decomposition in Eq. 8 allows for optimizing each expert separately. The optimization objective for a specific expert z and state s, that is,EQUATIONNote that this objective corresponds to the standard reverse KL objective from variational inference (c.f. Eq. 4) , with an additional term log q(z|a, s), which acts as a repulsion force, keeping the individual components from concentrating on the same mode. Assuming that q is differentiable and following the logic in Section 4.1 it is apparent that the single component objective in Eq 11 can be optimized only by having access to scores ∇ a log π(a|s). Moreover, we can again leverage amortized stochastic VI, that ismin νz E µ(s) [U s z (ν z , q)] ≈ min νz M N si∼µ U si z (ν z , q). M-Step for Updating the Gating. The M-Step for the gating parameters, i.e., minimizing U (ϕ, q) with respect to ξ ⊂ ϕ is given byEQUATIONUsing gradient estimators such as reinforce [57] , requires evaluating U s z (ν z , q) which is not possible as we do not have access to log π(a|s). This motivates the need for a different optimization scheme. We note that q ξ is a categorical distribution and can approximate any distribution over z. It does therefore not suffer from problems associated with the forward KL divergence such as mode averaging due to limited complexity of q ξ . We thus propose using the following objective for optimizing ξ, i.e.,EQUATIONresulting in a cross-entropy loss that does not require evaluating the intractable log π(a|s). Moreover, we note that using the forward KL does not change the minimizer asEQUATIONand, therefore does not affect the convergence guarantees of our method. Please note that the forward KL requires samples from the teacher model π. In practice, if the dataset used for training the diffusion model is available, it can be used as a proxy and prevent costly data regeneration.Step: Tighening the Lower Bound. Using the properties of the KL divergence, it can easily be seen that the global minimizer of the E-Step, i.e., the optimization objective defined in Eq. 10 can be found by leveraging Bayes' rule, i.e., q(z|a, s) = q ϕ old (z|a, s) = q ν old z (a|s, z)q ξ old (z|s) z q ν old z (a|s, z)q ξ old (z|s)EQUATION)The superscript 'old' refers to the previous iteration. Numerically, ϕ old can easily be obtained by using a stop-gradient operation which is crucial as q is fixed during the subsequent M-step which requires blocking the gradients of q with respect to ϕ. As the KL is set to zero after this update, the upper bound is tight after every E-step ensuring E µ(s) D KL (q ϕ (a|s)∥π θ (a|s)) = U (ϕ, q). Hence, VDD has similar convergence guarantees to EM, i.e., every update step improves the original objective.In denoising diffusion models, the score function is usually characterized as a time-dependent function∇ at log π t (a t |s) = f θ (a t , s, t), where t is the diffusion timestep. Yet, the formulation in Section 4.1 only leverages the pretrained diffusion model at time t → 0. However, [47] showed that using an ensemble of scores from multiple diffusion time steps significantly improves performances. We, therefore, replace Eq. 11 with a surrogate objective that utilizes scores at different time steps, that is,EQUATIONwith p(t) being a distribution on [0, T ]. Furthermore, we provide an ablation study for different time step selection schemes and empirically confirm the findings from [47] . 
experiments: 
 We conducted imitation learning experiments by distilling two types of diffusion models: variance preserving (VP) [2, 12] and variance exploding (VE) [65, 4] . We selected DDPM as the representative for VP and BESO as the representative for VE. We adopt the choices of samplers and the number of denoising steps in [9] and [13] . Additional evaluation of teacher models with different numbers of denoising steps can be found in Appendix F. In the experiments, VP-1 and VE-1 denote the results when performing only one denoising step of the respective diffusion models during inference.VDD-VP and VDD-VE denote the results of distilled VDD Additionally, we consider the SoTA Consistency Distillation (CD) [42] and Consistency Trajectory Model (CTM) [44] as baselines for comparing VDD's performance in distillation. For CD and CTM, we distill from the VE following the original works. For CTM we adapt the implementation and design choices from Consistency Policy [45] , which are specialized for behavior learning. We compare VDD against MoE learning baselines, namely the widely-used Expectation-Maximization (EM) [66] approach as a representative of the maximum likelihood-based objective and the recently introduced SoTA method Information Maximizing Curriculum (IMC) as representative of the reverse KL-based objective. To make them stronger baselines, we extend them with the architecture described in Figure 5 and name the extended methods as EM-GPT and IMC-GPT, respectively. For a fair comparison, we used the same diffusion models as the origin model for all distillation methods that we have trained on seed 0. For a statistically significant comparison, all methods have been run on 4 random seeds, and the mean and the standard deviation are reported throughout the evaluation. Detailed descriptions regarding the baselines implementation and hyperparameters selection can be found in Appendix D and E.The evaluations are structured as follows. Firstly, we demonstrate that VDD is able to achieve competitive performance with the SoTA diffusion distillation method and the original diffusion models on two established datasets. Next, we proceed to a recently proposed challenging benchmark with human demonstrations, where VDD outperforms existing diffusion distillation and SoTA MoE learning approaches. We then highlight the faster inference time of VDD. Following this, a series of ablation studies reflect the importance of VDD's essential algorithmic properties. Finally, we provide a visualization to offer deeper insights into our method.We first demonstrate the effectiveness of VDD using two widely recognized imitation learning datasets: Relay Kitchen [67] and XArm Block Push [68] . A detailed description of these environments is provided in Appendix C. To ensure a fair comparison, we follow the same evaluation process as outlined in [9] . The environment rewards for Relay Kitchen and the success rate for XArm Block Push are presented in Table 1a with mean and standard deviation resulting from 100 environment rollouts. The results indicate that VDD achieves a performance comparable to CD in both tasks, with slightly better outcomes in the block push dataset. An additional interesting finding is that BESO, with only one denoising step (VE-1), already proves to be a strong baseline in these tasks, as the original models outperformed the distillation results in both cases. We attribute this interesting observation to the possibility that the Relay Kitchen and the XArm Block Push tasks are comparably easy to solve and do not provide diverse, multi-modal data distributions. We therefore additionally evaluate the methods on a more recently published dataset (D3IL) [13] which is explicitly generated for complex robot imitation learning tasks and provides task entropy measurements.The D3IL benchmark provides human demonstrations for several challenging robot manipulation tasks, focusing on evaluating methods in terms of both success rate and versatility, i.e. the different behaviors that solve the same task. This benchmark includes a versatility measure for each task, referred to as task entropy. Task entropy is a scalar value ranging from 0 to 1, where 0 indicates the model has only learned one way to solve the task, and 1 indicates the model has covered all the skills demonstrated by humans. Detailed descriptions of the environments and the calculation of task entropy are provided in the Appendix C. The task success rate of the distilled polices is presented in Table 1a . The results show VDD outperforms consistency distillation and 1-step variants of origin models in 6 out of 7 tasks except the Aligning. However, in the Aligning task VDD achieves higher task entropy, indicating more diverse learned behaviors. The task entropy is presented in Table 1b .The results demonstrate that VDD achieves higher task entropy compared to both consistency models (CD, CTM) and the 1-step diffusion models (VP-1, VE-1) in 4 out of 7 tasks, which shows that our method replicates high-quality versatile behaviors from diffusion policies.The previous evaluation demonstrated that VDD can effectively distill diffusion models into MoEs, preserving the performance and behavioral versatility. In this section, we discuss the necessity of using VDD instead of directly learning MoEs from scratch by comparing VDD against EM-GPT and IMC-GPT. Both methods train MoE models from scratch but differ in their objectives. While EM is based on the well-known maximum likelihood objective, IMC is based on a reverse KL objective. The results in Table 2 show that VDD consistently outperforms both, EM-GPT and IMC-GPT across a majority of all tasks in terms of both success rate and task entropy. We attribute the performance boost leveraging the generalization ability of diffusion models and the stable updates provided by our decomposed lower bound Eq.(11).Inference with MoE models do not require an iterative denoising process and are therefore faster in sampling. We evaluate the inference time of VDD against DDPM and BESO on the state-based pushing task and the image-based stacking task and report the average results from 200 predictions in Table 3 . In addition to the absolute inference time in milliseconds, we report the number of function 3 for better comparability. The results show that VDD is significantly faster than the original diffusion models in both cases, even when the diffusion model takes only one denoising step. For a fair comparison, all methods used an identical number of transformer layers. The predictions were conducted using the same system (RTX 3070 GPU, Intel i7-12700 CPU).We assess the importance of VDD's key properties on different environments by reporting the task performance and task entropy averaged over four different seeds.Number of experts matters for task entropy. We start by varying the number of experts of the MoE model while freezing all other hyperparameters on the avoiding task. Figure 3a shows the average task success rate and task entropy of MoE models trained with VDD. The success rate is almost constantly high for all numbers of experts, except for the single expert (i.e. a Gaussian policy) case which shows a slightly higher success rate. However, the single expert can only cover a single mode of the multi-modal behavior space and hence achieves a task entropy of 0. With an increasing number of experts, the task entropy increases and eventually converges after a small drop.Training a gating distribution boosts performance. Figure 3b shows the success rates when training a parameterized gating network q ξ (z|s) (red) and when fixing the probability of choosing expert z to q(z) = 1/N , where N is the number of experts (blue). While training a gating distribution increases the success rate over three different tasks, the task entropy (see Figure 3c ) slightly decreases in two out of three tasks. This observation makes sense as the MoE with a trained gating distribution leads to an input-dependent specialization of each expert, while the experts with a fixed gating are forced to solve the task in every possible input.Interval time step sampling increases task entropy. Here, we explore different time step distributions p(t), as introduced in Eq.( 16). We consider several methods in Fig. 3d : using the minimum time step, i.e., p(t) = lim t→0 δ(t), where δ denotes a Dirac delta distribution, the maximum time step p(t) = δ(T ), a uniform distribution on [0, T ] and on sub-intervals [t 0 , t 1 ] ⊂ [0, T ] with interval bounds t 0 , t 1 being hyperparameters. While success rates were comparable across the variants, interval sampling yielded the highest task entropy with very high success rates. Thus, interval time-step sampling is adopted as our default setting. The results were obtained from the avoiding task.We provide additional visualizations on the Avoiding task from the D3IL task suite aiming to provide further intuition on how VDD leverages the individual experts. Figure 4 illustrates the expert selection according to the likelihood of the gating distribution at a given state, offering several key insights. First, VDD effectively distills experts with distinct behaviors, e.g., z 1 typically moves downward, z 2 tends to move upward, while z 3 and z 4 tend to generate horizontal movements. Second, the gating mechanism effectively deactivates redundant experts (z 6 , z 7 , z 8 ) in most states, demonstrating that a larger number of components can be used without harming performance, as the gating mechanism deactivates redundant experts. Lastly, using a single component (Z = 1) can achieve a perfect success rate at the cost of losing behavior diversity. On the contrary, using many experts potentially results in a slightly lower success rate but increased behavior diversity. These qualitative results are consistent with the quantitative results from the ablation study presented in Figure 3a . Different colors indicate components with highest likelihood according to the learned gating network q ξ (z|s) at a state s. For each step we select the action by first sampling an expert from the categorical gating distribution and then take the mean of the expert prediction. We decompose the case Z = 8 and visualize the individual experts z i (bottom row). Diverse behavior emerges as multiple actions are likely given the same state. For example, moving to the bottom right (z 1 ) and top right (z 2 ). An extreme case of losing diversity is seen with Z = 1, where the policy is unable to capture the diverse behavior of the diffusion teacher, leading to deterministic trajectories.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: We will open source the codes in the near future once they are cleaned up and anomnymity is not a concern anymore. All the experiments we conducted were using open-source datasets. In the experiments section and appendix C we provide information to get access to the data.Guidelines:• The answer NA means that paper does not include experiments requiring code. • The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 This work introduced Variational Diffusion Distillation (VDD), a novel method that distills a diffusion model to an MoE. VDD enables the MoE to benefit from the diffusion model's properties like generalization and complex, multi-modal data representation, while circumventing its shortcomings like long inference time and intractable likelihood calculation. Based on the variational objective, VDD derives a lower bound that enables optimizing each expert individually. The lower-bound leads to a stable optimization and elegantly leverages the gradient of the pre-trained score function such that the overall MoE model effectively benefits from the diffusion model's properties. The evaluations on nine sophisticated behavior learning tasks show that VDD achieves on-par or better distillation performance compared to SOTA methods while retaining the capability of learning versatile skills. The ablation on the number of experts reveals that a single expert is already performing well, but can not solve the tasks in a versatile manner. Additionally, the results show that training the gating distribution greatly boosts the performance of VDD, but reduces the task entropy.Limitations. VDD is not straightforwardly applicable to generating very high-dimensional data like images due to the MoE's contextual mean and covariance prediction. Scaling VDD to images requires further extensions like prediction in a latent space. Additionally, the number of experts needs to be pre-defined by the user. However, a redundantly high number of experts could increase VDD's training time and potentially decrease the usage in post hoc fine-tuning using reinforcement learning. Similar to other distillation methods, the performance of VDD is bounded by the origin model. Broader Impact. Improving and enhancing imitation learning algorithms could make real-world applications like robotics more accessible, with both positive and negative impacts. We acknowledge that it falls on sovereign governments' responsibility to identify these potential negative impacts. q ϕ (a|s) = z q ξ (z|s)q νz (a|s, z) Recall that the marginal likelihood of the Mixture of Experts is given as q ϕ (a|s) = z q ξ (z|s)q νz (a|s, z), (17) where z denotes the latent variable, q ξ (z|s) and q νz (a|s, z) are referred to as gating and experts respectively.The expected reverse Kullback-Leibler (KL) divergence between q ϕ and π is given asEQUATIONwhereEQUATION)We can writeEQUATIONwhere we have used the definition of the marginal likelihood of the Mixture of Experts in Eq. 3.With the identity q ϕ (a|s) = q ξ (z|s)q νz (a|s, z) q ϕ (z|s, a)we can further writeJ(ϕ) = s µ(s) z q ξ (z|s)a q νz (a|s, z) log q ξ (z|s) + log q νz (a|s, z) -log q(z|s, a)-log π(a|s)] dads.We can now introduce the auxiliary distribution q(z|a, s) by adding and subtracting it asJ(ϕ) = s µ(s) z q ξ (z|s)a q νz (a|s, z) log q ξ (z|s) + log q νz (a|s, z) -log q ϕ (z|s, a)-log π(a|s) + log q(z|a, s) -log q(z|a, s)] dads.We can rearrange the terms such thatJ(ϕ) = s µ(s) z q ξ (z|s)a q νz (a|s, z) log q ξ (z|s) + log q νz (a|s, z) -log π(a|s)-log q(z|a, s)] dads + s µ(s) z q ξ (z|s) a q νz (a|s, z) [log q(z|a, s)-log q ϕ (z|s, a) dads.With can plug in the identity q νz (a|s, z) = q ϕ (a|s)q ϕ (z|s, a)EQUATIONinto the second sum and obtainEQUATIONz q ξ (z|s) a q νz (a|s, z) log q ξ (z|s) + log q νz (a|s, z) -log π(a|s) -log q(z|a, s) dads + s z a q ϕ (a|s)q ϕ (z|s, a) log q(z|a, s) -log q ϕ (z|s, a) dads.which is the expected negative KL asJ(ϕ) = s µ(s) z q ξ (z|s)a q νz (a|s, z) log q ξ (z|s) + log q νz (a|s, z) -log π(a|s) -log q(z|a, s) dads -E µ(s) E q ϕ (a|s) D KL (q ϕ (z|s, a)∥q(z|a, s)).We note thatU (ϕ, q) = s µ(s) z q ξ (z|s)a q νz (a|s, z) log q ξ (z|s) + log q νz (a|s, z) -log π(a|s)EQUATIONsuch that we arrive to the identical expression as in Eq. 8EQUATIONB.2 Derivation of the Gating Update (Eq. 13)First, we note thatEQUATIONas we are optimizing w.r.t. ϕ, we can write const. = E π(a|s) [log π(a|s)] D KL (π(a|s)∥q ϕ (a|s)) = -E π(a|s) log q ϕ (a|s) + const.We can now introduce the latent variable z D KL (π(a|s)∥q ϕ (a|s)) = -E π(a|s) z q(z|a, s) log q ϕ (a|s)EQUATIONWe use the identity in Eq. 22 to arrive at D KL (π(a|s)∥q ϕ (a|s)) = -E π(a|s) z q(z|a, s) log q νz (a|s, z) + log q ξ (z|s) -log q ϕ (z|s, a)= -E π(a|s) z q(z|a, s) log q νz (a|s, z) + log q ξ (z|s) -log q ϕ (z|s, a)+ log q(z|a, s) -log q(z|a, s)] + const.(= -E π(a|s) z q(z|a, s) log q νz (a|s, z) + log q ξ (z|s) -log q(z|a, s)EQUATIONSince D KL (q(z|a, s)∥q ϕ (z|s, a)) ≥ 0 we have the upper bound U (ϕ, q) = -E π(a|s) z q(z|a, s) log q νz (a|s, z) + log q ξ (z|s) -log q(z|a, s) .We can now write U (ϕ, q) = -E π(a|s) z q(z|a, s) log q νz (a|s, z) + E π(a|s) D KL (q(z|a, s)∥q ξ (z|s)) .Hence optimizing U (ϕ, q) with respect to ϕ = ξ ∪ {ν z } z is equivalent to optimizing D KL (π(a|s)∥q ϕ (a|s)). Specifically optimizing with respect to ξ boils down toEQUATIONNoting that the expectation concerning µ(s) does not affect the minimizer, concludes the derivation.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We provide hyperparameter lists for each of the algorithms, how they were chosen and type of optimizer in appendix E.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: The experiments describe the number of trials and show the deviations in the result tables. There are no deviations for the origin models as they are the used as the base of the distillation and hence set to a fixed seed 0.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Relay Kitchen: A multi-task kitchen environment with long-horizon manipulation tasks such as moving kettle, open door, and turn on/off lights. The dataset consists of 566 human-collected trajectories with sequences of 4 executed skills. We used the same experiment settings and the pre-trained diffusion models from [9] .XArm Block Push: We used the adapted goal-conditioned variant from [9] . The Block-Push Environment consists of an XARm robot that must push two blocks, a red and a green one, into a red and green squared target area. The dataset consists of 1000 demonstrations collected by a deterministic controller with 4 possible goal configurations. The methods got 0.5 credit for every block pushed into one of the targets with a maximum score of 1.0. We use the pretrained Beso model from [9] .D3IL [13] is a simulation benchmark with diverse human demonstrations, which aims to evaluate imitation learning models' ability to capture multi-modal behaviors. D3IL provides 7 simulation tasks consisting of a 7DoF Franka Emika Panda robot and various objects, where each task has different solutions and the robot is required to acquire all behaviors. Except for success rate, D3IL proposes to use task behavior entropy to quantify the policy's capability of learning multi-modal distributions. Given the predefined behaviors β for each task, the task behavior entropy is defined as,EQUATIONwhere s 0 refers to the initial state and S 0 refers to the number of samples from the initial state distribution p(s 0 ). During the simulation, we rollout the policy multiple times for each s 0 and use a Monte Carlo estimation to compute the expectation of the behavior entropy.Stacking Sorting In this paper, we evaluate our algorithm in Avoiding, Aligning, Pushing, and Stacking with statebased representations and Sorting and Stacking with image-based representations. The simulation environments can be found in Fig. 7 . The Avoiding task requires the robot to reach the green line without colliding with any obstacles. The task contains 96 demonstrations, consisting of 24 solutions with 4 trajectories for each solution. The Aligning task requires the robot to push the box to match the target position and orientation. The robot can either push the box from inside or outside, thus resulting in 2 solutions. This task contains 1000 demonstrations, 500 for each solution with uniformly sampled initial states. The Pushing task requires the robot to push two blocks to the target areas. The robot can push the blocks in different orders and to different target areas, which gives the task 4 solutions. This task contains 2000 demonstrations, 500 for each solution with uniformly sampled initial states. The Sorting task requires the robot to sort red and blue blocks to the corresponding box. The number of solutions is determined by the sorting order. D3IL provides Sorting 2, 4, and 6 boxes, here we only use the Sorting-4 task, which contains around 1054 demonstrations with 18 solutions. The Stacking task requires the robot to stack three blocks in the target zone. Additionally, the blue block needs to be stacked upright which makes it more challenging. This task contains 1095 demonstrations with 6 solutions.BESO [9] is a continuous time diffusion policy that uses a continuous stochastic-differential equation to represent the denoising process. We implement this method from the D3IL benchmark, following the default which predicts one-step action conditional on the past five-step observations. DDPM [2, 10] is a discrete diffusion policy. We do not directly use the code from DiffusionPolicy [10] which implements an encoder-decoder transformer structure. For fair comparison, we evaluate this model from the D3IL benchmark which shares the same architecture as in BESO.Consistency Distillation [42] is designed to overcome the slow generation of diffusion models. Consistency models can directly map noise to data using one-step and few-step generation and they can be trained either through distilling pre-trained diffusion models or as an independent generative model. Our implementation takes the main training part of the model by integrating a GPT-based diffusion policy as the backbone. [44] is an extension of the CD model originally used in image generation. It augments the performance by integrating additional CTM and GAN loss terms in consideration. Later, it is used in robot policy prediction in [45] without taking the GAN loss. Our implementation of CTM is extended from our CD implementation, by modifying the loss computation.IMC [14] is a curriculum-based approach that uses a curriculum to assign weights to the training data so that the policy can select samples to learn, aiming to address the mode-averaging problem in multimodal distributions. We implement the model using the official IMC code with a GPT structure.EM [66] is based on the maximum likelihood objective and follows an iterative optimization scheme, where the algorithm switches between the M-step and the E-step in each iteration.We executed a large-scale grid search to fine-tune key hyperparameters for each baseline method. For other hyperparameters, we choose the value specified in their respective original papers. Below is a list summarizing the key hyperparameters that we swept during the experiment phase.Consistency Distillation: µ: EMA decay rate, N: see Algorithm 2 in [42] . All the other hyperparameters reuse the ones from the diffusion policy (BESO), as CD requires to be initialized using a pre-trained diffusion model. Unlike DDPM, which uses a fixed set of timesteps, BESO learns a continuous-time representation of the scores. This continuous representation enables the use of various numerical integration schemes, which can impact the performance of the diffusion model. We conducted an evaluation on the BESO teacher we used with different denoising steps. The results are presented in Table 5 . Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 


Decision: This paper introduces Variational Diffusion Distillation (VDD), a method for distilling pre-trained diffusion policies into Mixture of Experts (MoE) models, combining the expressiveness of diffusion models with the fast inference of MoEs. VDD uses variational inference with a decomposed objective to train MoE experts efficiently. The method is evaluated on 9 behavior learning tasks, showing strong performance.

Reviewers appreciated the paper's technical approach, methodology, and extensive evaluation. Concerns included limited comparison to recent baselines and need for more analysis and visualizations. The authors addressed these in their rebuttal by adding comparisons, providing missing results, offering deeper analysis, and including additional visualizations. They also released code and provided details on hyperparameters and training costs.

Given the paper's strengths, thorough evaluation, and the authors' careful addressing of concerns, I recommend accepting this paper. The method enables practical use of expressive diffusion policies in domains requiring fast inference, representing a contribution to the field.

Reviews: review: paper_summary: paper_summary: This paper studies the knowledge distillation problem in diffusion models by distilling denoising diffusion policies into Mixtures of Experts (MoE) using variational inference. The goal is to combine the advantages of diffusion models, with the fast inference capabilities of MoE models. The authors construct an upper bound of a KL divergence-style loss function between the diffusion policy and the MoE policy and implement EM steps for variational inference to achieve the distillation. The proposed method, VDD (Variational Distillation of Diffusion), demonstrates superior performance compared to existing distillation and MoE methods.
strength_weakness: strength_weakness: 1. In the limitations section, the paper claims that "VDD is not straightforwardly applicable to generating very high-dimensional data like images." However, the experimental results in Table 1(a) show that the distilled model VDD-DDPM outperforms the original DDPM diffusion policy in the Sorting (Image) and Stacking (Image) tasks, which contradicts the discussion. What is the reason for this? Why does the distilled model outperform the original diffusion policy?

2. Regarding training time, the time cost for distillation models (like VDD-DDPM) is higher than for MoE models. I assume this refers only to the knowledge distillation time. If we consider the original training time for diffusion policies, the total time cost may be significantly larger. It would be beneficial to report this total time for a fair comparison between MoE methods.

3. Does the distilled MoE model have the same architecture as the EM-GPT and IMC-GPT models?

4. While I agree that, according to this paper, the distillation of MoE achieves better performance than training from scratch, why do we distill from diffusion models? How about distilling from large and robust MoE models into smaller MoE models? This could be especially relevant in low-dimensional state observation cases, where MoE models reportedly perform better.
questions: questions: 1. Why does the distilled model outperform the original diffusion policy in Sorting (Image) and Stacking (Image) tasks?

2. Does the distilled MoE model have the same architecture as the EM-GPT and IMC-GPT models?

3. Why do we distill from diffusion models? How about distilling from large and robust MoE models into smaller MoE models?
limitations: limitations: The authors have discussed the limitations of the paper.

Please refer to the [weakness] section for detailed concerns and questions about the paper.
paper_summary: paper_summary: This paper presents Variational Diffusion Distillation (VDD), a method that distills denoising diffusion policies into Mixtures of Experts (MoE) using variational inference. Diffusion Models excel in learning complex distributions for behavior learning but have drawbacks like slow inference times. MoEs address these issues but are hard to train. In VDD, each expert can be trained separately under the corresponding guidance from the diffusion teacher. VDD demonstrates convincing performance in distilling complex distributions and outperforming existing methods across nine behavior learning tasks.
strength_weakness: strength_weakness: - I like the idea of updating the gating of MoE, which has been proven to be crucial for the success of VDD in an ablation study. However, it appears a bit unnatural to me that the experts are learned with reverse KL with the gating is learned with forward KL, even though the minimizer for forward and reverse KLs should align. Authors may find it interesting to visit another perspective from a probably concurrent work (Xie et al., 2024), in which a new distillation framework is derived from forward KL. 

- A very relevant work on distilling diffusion models, Luo et al. 2023, is missing. 

- I observed some performance gaps between VDD-DDPM and VDD-BESO across various tasks even though the teachers perform similarly. I wonder if VDD is sensitive to the specific forward process used in the diffusion model.

- In my humble opinion, the author could try some more complex tasks to demonstrate the capabilities of VDD. For instance, implementing a fast and expressive model on a dexterous hand would benefit from an expressive policy like diffusion while also requiring speed. Prior work exists in directly learning an MoE image generation model with the EM algorithm from data (Yu et al., 2019). Given the power of diffusion teachers, I would expect a distillation method to perform better than directly learning from data. 

Xie et al. 2024, EM Distillation for One-step Diffusion Models

Luo et al. 2023, Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models

Yu et al. 2021, Unsupervised Foreground Extraction via Deep Region Competition
questions: questions: - It appears that in simpler tasks like kitchen and block push, the distilled model underperforms compared to the original model. However, in more complex tasks, many distilled models seem to outperform the original. I am curious if the authors have some insights on this.
limitations: limitations: The authors provided the limitation statement, mainly concerning the expressity of MoE.
paper_summary: paper_summary: This paper presents a variational inference method for distilling denoising diffusion policies into Mixture-of-Experts (MoE) policies. The primary motivation is to combine the strengths of both worlds - the ability to learn complex, multi-modal distributions of diffusion models - and the efficiency of MoEs offering faster inference and tractable likelihoods. Importantly, the decomposed upper bound of the variational objective enables separate, and thus more robust training for different experts. The proposed method is empirically evaluated on 9 behavior learning tasks. The author demonstrate its ability to distill complex distributions, outperforming existing distillation methods.
strength_weakness: strength_weakness: 1. While EM is arguably be able to handle some limitations, it may also suffer from more longer convergence time during training. Moreover, while the authors discuss inference time improvements, this paper does not analyze training cost. As the training involves optimizing multiple experts and a gating - when the “optimal” number of experts is difficult to be pre-determined, it could potentially be computationally intensive. Is it possible to take task similarity into account and perform kind of meta-training like policy distillation?
2. The authors do not explicitly discuss how well VDD can capture long-range temporal dependency that might be presented in the original diffusion model, provided that it improves the inference speed. However, this can be important for sequential decision-making tasks.
questions: questions: Please see Weaknesses section
limitations: limitations: The authors do have included discussion about the limitation and shed light on how they might be approaching in future studies.
paper_summary: paper_summary: This study introduces Variational Diffusion Distillation (VDD), a novel method that distills pre-trained diffusion models into Mixture of Experts (MoE) frameworks. VDD addresses diffusion models' drawbacks of intractable likelihoods and long inference times, while leveraging their ability to represent complex distributions. By leveraging a decompositional variational objective, VDD trains MoEs efficiently, enabling real-time applications. VDD outperforms existing distillation and MoE training methods in several complex behavior learning tasks.
strength_weakness: strength_weakness: 1) The challenges posed, although relevant to diffusion models and their potential difficulties in handling certain tasks or long inference times, lack sufficient evidence to conclusively state that these issues persist across all novel variants of diffusion models.
2) The use of identical phrasing in the abstract and introduction could potentially limit the reader's engagement, as it fails to provide a nuanced progression of ideas.
3) A missing period on line 218 detracts slightly from the overall readability and professionalism of the manuscript.
4) It appears that the experiments for the DDPM component in Table 1(a) are incomplete for the kitchen and block push datasets, which limits the comprehensiveness of the evaluation.
5) While the proposed method achieves commendable results, it does not consistently outperform all metrics, necessitating a deeper analysis to explain the observed variations and identify potential avenues for improvement.
6) The inconsistency in Table 3 left, where 2.16 is highlighted as better than 2.15 is confusing and suggests that the experiment may not have been fully completed.
7) The comparison to a limited number of methods may limit the ability to comprehensively assess the strengths and weaknesses of the proposed approach.
8) The lack of illustrative visualizations, such as good and bad case studies, in the main text hinders the reader's ability to fully grasp the performance and limitations of the method in real-world scenarios.
questions: questions: see weaknesses
limitations: limitations: see weaknesses
paper_summary: paper_summary: This paper introduces a theoretical method for extracting an MoE policy from a pretrained diffusion policy.  The advantage of MoE policy is that it allows faster sampling speed compared with the diffusion policy more stable training and better performance compared with methods that train MoE from scratch.
strength_weakness: strength_weakness: 1. One main advantage of MoE policy against deterministic policy is its diversity. However, this paper seems to only focus on evaluating the final performance of algorithms and does not compare the diversity of different algorithms. Regarding on performance solely, it seems to me Gaussian policy + reverse KL objective is good enough, in what cases do we need more diverse policies? This needs to be justified. 
2. Variantional distillation is a known technique in the 3D field. Although its application in RL is new and meaningful, this still indicates limited novelty. Could emphasize more on the new technique introduced in this paper when applying variational distillation for policy extraction.
questions: questions: 1. What is the mixture number n in your experiments?  Can you provide any experiment results ablating n? This could be very informative.
2. I highly suggest the authors provide the source code during rebuttal. Reproducibility is crucial for this kind of paper. (Even if it is just unsorted/raw code)
limitations: limitations: Above

Comments: title: Thank you for the insightful discussions and help to improve VDD!
Response
Regarding the unexpected lowering of the score
Please reply to the rebuttal.
Response to the rebuttal
We thank the reviewer for the positive feedback
Response
comment: Dear Reviewers,

We sincerely thank all of the reviewers for the thoughtful feedback and engaging discussions throughout the rebuttal process. Your insights have been invaluable in helping us improve our work on VDD. Going into the reviewer discussion, we would like to summarize the main improvements that resulted from your feedback during the rebuttal:

- **Related Work**: Expanded discussion on VDD's relation to recent diffusion distillation approaches to highlight our novelty.
- **Baseline Comparisons**: Added evaluations against Consistency Trajectory Model (CTM) to show the strong distillation performance of VDD.
- **Completed Results-Table:** Included results for DDPM and VDD-DDPM on kitchen and block push datasets.
- **Enhanced Evaluation on Teacher Models:** Added evaluations on diffusion teacher model with different samplers and different integration steps to show that VDD is able to maintain comparable performance with the teacher model.
- **Detailed Training Cost Information**: Provided clearer information on parameter counts and training times across varying number of experts. Reported the total training time of VDD including the training time of the diffusion teacher on kitchen and block push tasks.
- **Enhanced Visualization**: Added new visualization to offer better insights in to our method. It highlights the most likely experts according to the gating probability at a given state.
- **Open-Sourced Code-Base**: Provided a link to an anonymous Github repository containing the code for reproducing the results of the paper to the AC. Additionally, we provided a link to an IPython notebook for testing our method on a simple toy task.

These changes have significantly strengthened our paper and all discussed points will be included in the final version.

Thank you again for your time and valuable feedback and we hope you take these points into account for your final decision.

Sincerely,

The Authors
We sincerely thank the reviewer for their positive feedback and constructive suggestions, which helped us improve the paper further. We are also pleased that our response addressed their concerns.
Thank you for the author's response, which addressed most of my concerns and questions. I am raising my initial score from 6 to 7.
We are very surprised and saddened to see that the reviewer unexpectedly lowered their score without further explanation.

We thoroughly checked the review and our rebuttal again and fail to see which concerns have not been substantively addressed.

In particular, we want to emphasize that in response to the reviewer's concerns, we conducted additional experiments, evaluations and visualizations during the rebuttal period:
- We performed DDPM distillation experiments on the Kitchen and Block Push datasets (in response to **Weakness 4**).
- We introduced the recently proposed Consistency Trajectory Model (CTM) as a new baseline (in response to **Weakness 7**).
- We included an illustrative visualization to provide deeper insights into our method (in response to **Weakness 8**).

The remaining concerns have also been addressed in the original rebuttal:
- **Weakness 1** was addressed and we offered to include additional discussion and baselines if pointed towards respective work.

- **Weakness 2,3** and **6** were addressed by fixing typos and rephrasing sentences in the improved manuscript, which can not be uploaded as to rebuttal guidelines.

- **Weakness 5** has been addressed and helped us clarify the goals, limitations and future work of our approach

We sincerely hope there is an underlying misunderstanding and want to point out that the additional experiments and visualization are in the **PDF attached to the "Author Rebuttal" post** and not in the main manuscript (adhering to the rebuttal guidelines).

We are of course happy to further clarify and elaborate on any of these concerns.
We sincerely thank the reviewer for the positive comments on our work. We aimed to address all questions accurately and concisely, and we apologize if doing so left open questions. If the reviewer points us towards the respective issues, we would be happy to clarify and elaborate on them.
I appreciate the authors for the efforts in answering my questions. Although some issues were not substantively addressed, I maintain my positive score due to the overall quality of the paper.
Dear Reviewer,

As the discussion period is nearing its conclusion, we kindly ask you to engage in the discussion and provide notes on any concerns that have not yet been addressed, along with the reasons why.

Thank you for your attention to this matter.

AC.
> Why does the student VDD model outperform the teacher diffusion model? […] According to the paper, VDD-BESO is both more accurate and not less diverse than BESO. What could be the reason for this discrepancy?
> 

We would like to express our sincere gratitude to the reviewer for their insightful questions, which prompted us to conduct further investigations. Below is a summary of our findings.

Unlike DDPM, which uses a fixed set of timesteps, BESO learns a continuous-time representation of the scores. This continuous representation enables the use of various numerical integration schemes, which can impact the performance of the diffusion model.

In our original experiments, we used the Euler-Maruyama method with 16 integration steps, following [1]. However, upon further investigation, we observed that utilizing alternative integration schemes, such as DDIM [2] or the Heun method [3], and increasing the number of integration steps, produced different success rates and entropy values. Some of these results surpassed those of VDD-BESO, underscoring that the teacher model indeed serves as an upper bound for the student model's performance. These findings are detailed in the table below. In the table we follow the convention "Sampler name - Number of integration steps”.

|  | Euler_Maru-16 | Euler-16 | Heun-16 | DDIM-16 |
| --- | --- | --- | --- | --- |
| **Success Rate** | $0.70$ | $0.68$ | $0.72$ | $0.73$ |
| **Task Entropy** | $0.19$ | $0.25$ | $0.25$ | $0.25$ |
|  | **Euler_Maru-32** | **Euler_Maru-64** |  | **VDD-BESO** |
| **Success Rate** | $0.76$ | $0.76$ |  | $0.76\pm0.04$ |
| **Task Entropy** | $0.23$ | $0.24$ |  | $0.22\pm0.03$ |

We, therefore, attribute VDD-BESO’s superior performance over BESO to the timestep selection strategy (see Section 4.3) of the learned continuous-time score functions.

In response to the reviewer's comments, we will include the results of various integration schemes across all tasks and provide additional clarifications in our revised manuscript.

---

> Could this be related to the use of image observations?
> 

Regarding the performance gap observed primarily in image-based tasks, we suspect that this may be due to the use of the pre-trained vision encoder from the diffusion model for VDD, which likely enhances distillation performance. However, to confirm this hypothesis and provide a definitive explanation, further investigations are required.

---

> How to choose the optimal number of experts in real-world control tasks?
> 
> 
> The results indicate that reducing the number of experts can increase the success rate, albeit at the cost of skill diversity. While this conclusion seems reasonable, it raises the question: how should we determine the optimal number of experts in a Mixture of Experts (MoE) model for a given real-world control task or set of tasks? In many safety-critical control problems, success rate is often the primary concern, leading to the possibility of opting for a smaller number of experts. Could you provide guidance on this decision, or better yet, demonstrate through experiments any potential drawbacks of this strategy?
> 

We understand that when the reviewer refers to "safety-critical control problems," they are addressing situations where failures or incorrect operations in control systems could lead to significant harm to humans or robots.

To provide a meaningful response, we believe it is important to have more specific details about the particular task or setting.

For example, in a static task environment where achieving high success rates is the primary objective, it might be advantageous to select a small number of components.

On the other hand, if the task involves constantly changing dynamics, a policy that incorporates a diverse set of skills is likely to be more robust [4]. For instance, in the avoidance task depicted in Figure R1 of the PDF attached to our rebuttal, if the path learned by a single-expert policy (Z=1) becomes obstructed, the success rate would drop to zero. In such a dynamic environment, diverse policies could maintain some level of success in completing the task, making a larger number of experts the more suitable option.

If our interpretation of "safety-critical control problems" differs from the reviewer’s definition, we would be happy to provide further clarification.

---

[1] Jia X, Blessing D, Jiang X, Reuss M, Donat A, Lioutikov R, Neumann G. Towards diverse behaviors: A benchmark for imitation learning with human demonstrations. ICLR 2024.

[2] Song J, Meng C, Ermon S. Denoising diffusion implicit models. ICLR 2021.

[3] Karras T, Aittala M, Aila T, Laine S. Elucidating the design space of diffusion-based generative models. NIPS 2022.

[4] Eysenbach, B. and Levine, S., 2021. Maximum entropy RL (provably) solves some robust RL problems. ICLR 2022.
Dear Reviewer, 

Please reply to the rebuttal. 

AC.
I appreciate the author's detailed responses, which have addressed most of my concerns. 
However, there are still some aspects that remain unclear to me.

1. Why does the student VDD model outperform the teacher diffusion model?   
The author suggests that this may be related to the trade-off between mastering a single skill very well and learning multiple skills with slightly less accuracy. However, when looking at the success rates in Table 1(a), the student model only outperforms in the Sorting (Image) and Stacking (Image) tasks, but not in others. Given a fixed teacher model and a fixed number of experts, it seems difficult to attribute this phenomenon to the acquisition of multiple skills or increased diversity. Could this be related to the use of image observations? Interestingly, the task entropy of VDD-BESO is comparable to BESO in the Sorting (Image) task. According to the paper, VDD-BESO is both more accurate and not less diverse than BESO. What could be the reason for this discrepancy?

2. How to choose the optimal number of experts in real-world control tasks?  
The results indicate that reducing the number of experts can increase the success rate, albeit at the cost of skill diversity. While this conclusion seems reasonable, it raises the question: how should we determine the optimal number of experts in a Mixture of Experts (MoE) model for a given real-world control task or set of tasks? In many safety-critical control problems, success rate is often the primary concern, leading to the possibility of opting for a smaller number of experts. Could you provide guidance on this decision, or better yet, demonstrate through experiments any potential drawbacks of this strategy?
We thank the reviewer for the positive evaluation and valuable suggestions, and we are pleased that our response addressed their concerns.  We are happy to address additional questions and suggestions to improve our work further.
I appreciate the authors for the efforts in answering my questions and I have no more questions at this stage. I will keep my score and lean towards positive outcome.
We sincerely appreciate the reviewer's positive evaluation and valuable suggestions. We are happy to address any additional questions and suggestions to further improve our work and its assessment.
Thank the authors for their response. I will keep my accepting rating.
We thank the reviewer for the positive feedback and are pleased that our response addressed their concerns. We fully agree that the community needs new benchmarks that require and evaluate diversity in RL settings. We are happy to address any other concerns that might arise during discussion session.
I thank the authors for the very detailed responses, which have resolved most of my concerns. I raise the score from 5 to 6. 

However, I would like to share one of my very personal opinions: Most (if not all) current public RL benchmarks and their evaluation metrics actually lack a strong need for policy diversity, which hinders the necessity of diversified policies. I believe designing new tasks and more difficult benchmarks is of critical importance for designing these algorithms.