SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Convergence of the Mixture of Experts Objective (Proposition 3.3). As we tighten the lower bound  $L$  in every E-step, it remains to show that  $L(\psi^{(i)}, q) \geq L(\psi^{(i-1)}, q)$  to prove that  $J(\psi)$  is guaranteed to converge to a stationary point, with  $\psi = \{p(z), \{p(\mathbf{o}, \mathbf{a}|z)\}_{z}, \{\theta_z\}_{z}\}$ . To that end, we again perform a coordinate ascent on  $L$  to show a monotonic improvement in every coordinate. Note that we find the unique  $p(z)$  and  $\{p(\mathbf{o}, \mathbf{a}|z)\}_{z}$  that maximize  $L$  via Proposition 3.2 and Equation 5. Analogously to the proof of Proposition 3.1 we can show monotonic improvement of  $L$  in  $\{\theta_z\}_{z}$  under suitable assumptions on  $\log p_{\theta_z}$ .
Remark. Although stochastic gradient ascent does not guarantee strictly monotonic improvements in the objective function  $J$ , our empirical observations suggest that  $J$  indeed tends to increase monotonically in practice as shown by Figure 6.
Figure 6: Convergence of the Mixture of Experts Objective. Objective function value  $J$  for 100 training iterations  $i$  on Franka Kitchen.
Expanding the entropy in Equation 5 we obtain
$$
p ^ {*} (z) \propto \exp \left(\mathbb {E} _ {p ^ {*} (\mathbf {o}, \mathbf {a} | z)} \left[ R _ {z} (\mathbf {o}, \mathbf {a}) / \eta - \log p ^ {*} (\mathbf {o}, \mathbf {a} | z) \right]\right).
$$
Using  $p^{*}(\mathbf{o},\mathbf{a}|z) = \tilde{p} (\mathbf{o},\mathbf{a}|z) / \sum_{n}\tilde{p} (\mathbf{o}_{n},\mathbf{a}_{n}|z)$  yields
$$
p ^ {*} (z) \propto \exp \left(\mathbb {E} _ {p ^ {*} (\mathbf {o}, \mathbf {a} | z)} [ R _ {z} (\mathbf {o}, \mathbf {a}) / \eta - \log \tilde {p} (\mathbf {o}, \mathbf {a} | z) + \log \sum_ {n} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) ]\right).
$$
Next, leveraging that  $\log \tilde{p} (\mathbf{o},\mathbf{a}|z) = R_z(\mathbf{o}_n,\mathbf{a}_n) / \eta$  we see that
$$
p ^ {*} (z) \propto \exp \left(\mathbb {E} _ {p ^ {*} (\mathbf {o}, \mathbf {a} | z)} [ \log \sum_ {n} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) ]\right) = \sum_ {n} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z),
$$
A.3 Proof of Corollary 3.2.1
We start by rewriting the lower bound as  $L(\psi, q) =$
$$
\mathbb {E} _ {p ^ {*} (z)} \left[ \mathbb {E} _ {p ^ {*} (\mathbf {o}, \mathbf {a} | z)} \left[ R _ {z} (\mathbf {o}, \mathbf {a}) - \eta \log p ^ {*} (\mathbf {o}, \mathbf {a} | z) \right] - \eta \log p ^ {*} (z) \right].
$$
Using  $p^*(\mathbf{o}, \mathbf{a}|z) \propto \tilde{p}(\mathbf{o}, \mathbf{a}|z)$  and Proposition 3.2 we obtain
$$
\begin{array}{l} L (\boldsymbol {\psi}, q) = \mathbb {E} _ {p ^ {*} (z)} \left[ \mathbb {E} _ {p ^ {*} (\mathbf {o}, \mathbf {a} | z)} \left[ R _ {z} (\mathbf {o}, \mathbf {a}) - \eta \log \tilde {p} (\mathbf {o}, \mathbf {a} | z) \right. \right. \\ + \eta \log \sum_ {n} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) ] - \eta \log \sum_ {n} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) + \eta \log \sum_ {z} \sum_ {n} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) ] \\ \end{array}
$$
With  $\eta \log \tilde{p} (\mathbf{o},\mathbf{a}|z) = R_z(\mathbf{o}_n,\mathbf{a}_n)$  all most terms cancel, giving
$$
L (\boldsymbol {\psi}, q) = \mathbb {E} _ {p ^ {*} (z)} \left[ \eta \log \sum_ {z} \sum_ {n} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) \right] = \eta \log \sum_ {z} \sum_ {n} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z),
$$
A.4 Proof of Corollary 3.2.2
$$
\begin{array}{l} \min _ {\phi} \mathbb {E} _ {p (\mathbf {o})} D _ {\mathrm {K L}} (p (z | \mathbf {o}) \| g _ {\phi} (z | \mathbf {o})) = \max _ {\phi} \int_ {\mathcal {O}} \sum_ {z} p (\mathbf {o}, z) \log g _ {\phi} (z | \mathbf {o}) \mathrm {d} \mathbf {o} \\ = \max _ {\phi} \int_ {\mathcal {O}} \int_ {\mathcal {A}} \sum_ {z} p (\mathbf {o}, \mathbf {a}, z) \log g _ {\phi} (z | \mathbf {o}) \mathrm {d} \mathbf {a} \mathrm {d} \mathbf {o} = \max _ {\phi} \int_ {\mathcal {O}} \int_ {\mathcal {A}} p (\mathbf {o}, a) \sum_ {z} p (z | \mathbf {o}, \mathbf {a}) \log g _ {\phi} (z | \mathbf {o}) \mathrm {d} \mathbf {a} \mathrm {d} \mathbf {o} \\ = \min  _ {\phi} \mathbb {E} _ {p (\mathbf {o}, \mathbf {a})} D _ {\mathrm {K L}} \big (p (z | \mathbf {o}, \mathbf {a}) \| g _ {\phi} (z | \mathbf {o}) \big). \\ \end{array}
$$
Expanding the expected KL divergence, we get
$$
\underset {\boldsymbol {\phi}} {\min} \mathbb {E} _ {p (\mathbf {o}, \mathbf {a})} D _ {\mathrm {K L}} \big (p (z | \mathbf {o}, \mathbf {a}) \| g _ {\boldsymbol {\phi}} (z | \mathbf {o}) \big) = \underset {\boldsymbol {\phi}} {\min} \sum_ {n} p (\mathbf {o} _ {n}, \mathbf {a} _ {n}) \sum_ {z} p (z | \mathbf {o} _ {n}, \mathbf {a} _ {n}) \log \frac {p (z | \mathbf {o} _ {n} , \mathbf {a} _ {n})}{g _ {\boldsymbol {\phi}} (z | \mathbf {o} _ {n})}.
$$
Noting that  $p(z|\mathbf{o}_n, \mathbf{a}_n)$  is independent of  $\phi$  we can rewrite the objective as
$$
\max _ {\phi} \sum_ {n} p (\mathbf {o} _ {n}, \mathbf {a} _ {n}) \sum_ {z} p (z | \mathbf {o} _ {n}, \mathbf {a} _ {n}) \log g _ {\phi} (z | \mathbf {o} _ {n}).
$$
Using that  $p(z|\mathbf{o},\mathbf{a}) = \tilde{p} (\mathbf{o},\mathbf{a}|z) / \sum_z\tilde{p} (\mathbf{o},\mathbf{a}|z)$  together with  $p(\mathbf{o},\mathbf{a}) = \sum_{z}p^{*}(z)p^{*}(\mathbf{o},\mathbf{a}|z)$  yields
$$
\max _ {\boldsymbol {\phi}} \sum_ {n} \sum_ {z} p ^ {*} (z) p ^ {*} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) \sum_ {z} \frac {\tilde {p} (\mathbf {o} _ {n} , \mathbf {a} _ {n} | z)}{\sum_ {z} \tilde {p} (\mathbf {o} _ {n} , \mathbf {a} _ {n} | z)} \log g _ {\boldsymbol {\phi}} (z | \mathbf {o} _ {n}).
$$
Using Proposition 3.2 we can rewrite  $p^*(z)p^*(\mathbf{o}, \mathbf{a}|z)$  as  $\tilde{p}(\mathbf{o}, \mathbf{a}|z)/\sum_{z}\sum_{n}\tilde{p}(\mathbf{o}_n,\mathbf{a}_n|z)$ . Since the constant factor  $1/\sum_{z}\sum_{n}\tilde{p}(\mathbf{o}_n,\mathbf{a}_n|z)$  does not affect the optimal value of  $\phi$  we obtain
$$
\max _ {\boldsymbol {\phi}} \sum_ {n} \sum_ {z} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) \sum_ {z} \frac {\tilde {p} (\mathbf {o} _ {n} , \mathbf {a} _ {n} | z)}{\sum_ {z} \tilde {p} (\mathbf {o} _ {n} , \mathbf {a} _ {n} | z)} \log g _ {\boldsymbol {\phi}} (z | \mathbf {o} _ {n}) = \max _ {\boldsymbol {\phi}} \sum_ {n} \sum_ {z} \tilde {p} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) \log g _ {\boldsymbol {\phi}} (z | \mathbf {o} _ {n}),
$$
B.1 Lower Bound Decomposition
To arrive at Equation 4 by marginalizing over the latent variable  $o$  for the entropy of the joint curriculum, i.e.,
$$
\begin{array}{l} \mathcal {H} (\mathbf {o}, \mathbf {a}) = - \sum_ {n} p (\mathbf {o} _ {n}, \mathbf {a} _ {n}) \log p (\mathbf {o} _ {n}, \mathbf {a} _ {n}) \\ = - \sum_ {n} p (\mathbf {o} _ {n}, \mathbf {a} _ {n}) \sum_ {z} p (z | \mathbf {o} _ {n}, \mathbf {a} _ {n}) \log p (\mathbf {o} _ {n}, \mathbf {a} _ {n}) \\ \end{array}
$$
Next, we use Bayes' theorem, that is,  $p(\mathbf{o}_n, \mathbf{a}_n) = p(z)p(\mathbf{o}_n, \mathbf{a}_n|z) / p(z|\mathbf{o}_n, \mathbf{a}_n)$ , giving
$$
\mathcal {H} (\mathbf {o}, \mathbf {a}) = - \sum_ {n} p (\mathbf {o} _ {n}, \mathbf {a} _ {n}) \sum_ {z} p (z | \mathbf {o} _ {n}, \mathbf {a} _ {n}) \big (\log p (z) + \log p (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) - \log p (z | \mathbf {o} _ {n}, \mathbf {a} _ {n}) \big).
$$
Moreover, we add and subtract the log auxiliary distribution  $\log q(z|\mathbf{o}_n,\mathbf{a}_n)$  which yields
$$
\begin{array}{l} \mathcal {H} (\mathbf {o}, \mathbf {a}) = - \sum_ {n} p \left(\mathbf {o} _ {n}, \mathbf {a} _ {n}\right) \sum_ {z} p (z | \mathbf {o} _ {n}, \mathbf {a} _ {n}) \left(\log p (z) + \log p \left(\mathbf {o} _ {n}, \mathbf {a} _ {n} \mid z\right) \right. \\ \left. - \log p (z | \mathbf {o} _ {n}, \mathbf {a} _ {n}) + \log q (z | \mathbf {o} _ {n}, \mathbf {a} _ {n}) - \log q (z | \mathbf {o} _ {n}, \mathbf {a} _ {n})\right). \\ \end{array}
$$
Rearranging the terms leads and writing the sums in terms of expectations we arrive at
$$
\mathcal {H} (\mathbf {o}, \mathbf {a}) = \mathbb {E} _ {p (z)} \left[ \mathbb {E} _ {p (z | \mathbf {o}, \mathbf {a})} [ \log q (z | \mathbf {o}, \mathbf {a}) ] + \mathcal {H} (\mathbf {o}, \mathbf {a} | z) \right] + \mathcal {H} (z) + D _ {\mathrm {K L}} \big (p (z | \mathbf {o}, \mathbf {a}) \| q (z | \mathbf {o}, \mathbf {a}) \big).
$$
Lastly, multiplying  $\mathcal{H}(\mathbf{o},\mathbf{a})$  with  $\eta$  and adding  $\mathbb{E}_{p(z)}\mathbb{E}_{p(\mathbf{o},\mathbf{a}|z)}[\log p_{\theta_z}(\mathbf{y}|\mathbf{x},z)]$  we arrive at Equation 4 which concludes the derivation.
B.2 M-Step Objectives
Closed-Form Curriculum Updates. In order to derive the closed-form solution to Equation 5 (RHS) we solve
$$
\max  _ {p (\mathbf {o}, \mathbf {a} | z)} J _ {z} (p (\mathbf {o}, \mathbf {a} | z), \boldsymbol {\theta} _ {z}) = \max  _ {p (\mathbf {o}, \mathbf {a} | z)} \mathbb {E} _ {p (\mathbf {o}, \mathbf {a} | z)} [ R _ {z} (\mathbf {o}, \mathbf {a}) ] + \eta \mathcal {H} (\mathbf {o}, \mathbf {a} | z) \quad \text {s u b j e c t t o} \quad \sum_ {n} p (\mathbf {o}, \mathbf {a}) = 1.
$$
Following the procedure of constrained optimization, we write down the Lagrangian function [9] as
$$
\mathcal {L} (p, \lambda) = \sum_ {n} p (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) R _ {z} (\mathbf {o} _ {n}, \mathbf {a} _ {n}) - \eta \sum_ {n} p (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) \log p (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) + \lambda (\sum_ {n} p (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) - 1),
$$
where  $\lambda$  is the Lagrangian multiplier. As  $p$  is discrete, we solve for the optimal entries of  $p(\mathbf{o}_n,\mathbf{a}_n|z)$  that is,  $p^{\prime}(\mathbf{o}_n,\mathbf{a}_n,\lambda |z) = \arg \max_p\mathcal{L}(p,\lambda)$ . Setting the partial derivative of  $\mathcal{L}(p,\lambda)$  with respect to  $p$  zero, i.e.,
$$
\frac {\partial}{\partial p \left(\mathbf {o} _ {n} , \mathbf {a} _ {n} \mid z\right)} \mathcal {L} (p, \lambda) = R _ {z} \left(\mathbf {o} _ {n}, \mathbf {a} _ {n}\right) - \eta \log p \left(\mathbf {o} _ {n}, \mathbf {a} _ {n} \mid z\right) - \eta + \lambda \stackrel {!} {=} 0.
$$
yields  $p^{\prime}(\mathbf{o}_n,\mathbf{a}_n,\lambda |z) = \exp \left(R_z(\mathbf{o}_n,\mathbf{a}_n) - \eta +\lambda\right) / \eta .$
Plugging  $p'$  back in the Lagrangian gives the dual function  $g(\lambda)$ , that is,
$$
g (\eta) = \mathcal {L} (p ^ {\prime}, \lambda) = - \eta + \eta \sum_ {n} \exp \left(R _ {z} (\mathbf {o} _ {n}, \mathbf {a} _ {n}) - \eta + \lambda\right) / \eta .
$$
Solving for  $\lambda^{*} = \arg \min_{\lambda >0}g(\lambda)$  equates to
$$
\begin{array}{l} \frac {\partial}{\partial \lambda} g (\lambda) = - 1 + \eta \sum_ {n} \exp \left(R _ {z} \left(\mathbf {o} _ {n}, \mathbf {a} _ {n}\right) - \eta + \lambda\right) / \eta \stackrel {!} {=} 0 \\ \Longleftrightarrow \lambda^ {*} = - \log \left(\eta \sum_ {n} \exp \left(R _ {z} (\mathbf {o} _ {n}, \mathbf {a} _ {n}) - \eta\right) / \eta\right). \\ \end{array}
$$
Finally, substituting  $\lambda^{*}$  into  $p^{\prime}$  we have
$$
p ^ {*} (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) = p ^ {\prime} (\mathbf {o} _ {n}, \mathbf {a} _ {n}, \lambda^ {*} | z) = \frac {\exp \left(R _ {z} (\mathbf {o} _ {n} , \mathbf {a} _ {n}) / \eta\right)}{\sum_ {n} \exp \left(R _ {z} (\mathbf {o} _ {n} , \mathbf {a} _ {n}) / \eta\right)},
$$
which concludes the derivation. The derivation of the optimal mixture weights  $p^{*}(z)$  works analogously.
Expert Objective. In order to derive the expert objective of Equation 6 we solve  $\max_{\pmb{\theta}_z} J_z(p(\mathbf{o}, \mathbf{a}|z), \pmb{\theta}_z) =$
$$
\max _ {\pmb {\theta} _ {z}} \sum_ {n} p (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) \Big (\log p _ {\pmb {\theta} _ {z}} (\mathbf {a} _ {n} | \mathbf {o} _ {n}, z) + \eta \log q (z | \mathbf {o} _ {n}, \mathbf {a} _ {n}) - \eta \log p (\mathbf {o} _ {n}, \mathbf {a} _ {n} | z) \Big).
$$
Noting that  $q(z|\mathbf{o}_n,\mathbf{a}_n)$  and  $p(\mathbf{o}_n,\mathbf{a}_n|z)$  are independent of  $\theta_z$  and  $p(\mathbf{o}_n,\mathbf{a}_n|z) = \tilde{p}(\mathbf{o}_n,\mathbf{a}_n|z)/\sum_n \tilde{p}(\mathbf{o}_n,\mathbf{a}_n|z)$  we find that
$$
\max _ {\pmb {\theta} _ {z}} J _ {z} (p (\mathbf {o}, \mathbf {a} | z), \pmb {\theta} _ {z}) = \max _ {\pmb {\theta} _ {z}} \sum_ {n} \frac {\tilde {p} (\mathbf {o} _ {n} , \mathbf {a} _ {n} | z)}{\sum_ {n} \tilde {p} (\mathbf {o} _ {n} , \mathbf {a} _ {n} | z)} \log p _ {\pmb {\theta} _ {z}} (\mathbf {a} _ {n} | \mathbf {o} _ {n}, z).
$$
Noting that  $\sum_{n}\tilde{p} (\mathbf{o}_n,\mathbf{a}_n|z)$  is a constant scaling factor concludes the derivation.
C Experiment Setup
C.1 Environments and Datasets
C.1.1 Obstacle Avoidance
Figure 7: The left figure shows 6 out of 24 ways of completing the obstacle avoidance task. The middle figure shows 30 randomly sampled initial block configurations for the block pushing task. The right figure visualizes the Franka kitchen environment.
Dataset. The obstacle avoidance dataset contains 96 trajectories resulting in a total of  $7.3\mathrm{k}$  ( $\mathbf{o}, \mathbf{a}$ ) pairs. The observations  $\mathbf{o} \in \mathbb{R}^4$  contain the end-effector position and velocity in Cartesian space. Please note that the height of the robot is fixed. The actions  $\mathbf{a} \in \mathbb{R}^2$  represent the desired position of the robot. The data is recorded such that there are an equal amount of trajectories for all 24 ways of avoiding the obstacles and reaching the target line. For successful example trajectories see Figure 7.
Performance Metrics. The success rate indicates the number of end-effector trajectories that successfully reach the target line (indicated by green color in Figure 2). The entropy
$$
\mathcal {H} (\beta) = - \sum_ {\beta} p (\beta) \log_ {2 4} p (\beta),
$$
is computed for successful trajectories, where each behavior  $\beta$  is one of the 24 ways of completing the task. To assess the model performance, we simulate 1000 end-effector trajectories. We count the number of successful trajectories for each way of completing the task. From that, we calculate a categorical distribution over behaviors  $p(\beta)$  which is used to compute the entropy. By the use of  $\log_{24}$  we make sure that  $\mathcal{H}_{24}(\tau) \in [0,1]$ . If a model is able to discover all modes in the data distribution with equal probability, its entropy will be close to 1. In contrast,  $\mathcal{H}(\beta) = 0$  if a model only learns one solution.
C.1.2 Block Pushing
Dataset. The block pushing dataset contains 500 trajectories for each of the four push sequences (see Figure 8) resulting in a total of 2000 trajectories or  $463\mathrm{k}$  ( $\mathbf{o}, \mathbf{a}$ ) pairs. The observations  $\mathbf{o} \in \mathbb{R}^{16}$  contain the desired position and velocity of the robot in addition to the position and orientation of the green and red block. Please note that the orientation of the blocks is represented as quaternion number system and that the height of the robot is fixed. The actions  $\mathbf{a} \in \mathbb{R}^2$  represent the desired position of the robot. This task is similar to the one proposed in [17]. However, they use a deterministic controller to record the data whereas we use human demonstrators which increases the difficulty of the task significantly due to the inherent versatility in human behavior.
Performance Metrics. The success rate indicates the number of end-effector trajectories that successfully push both blocks to different target zones. To assess the model performance on nonsuccessful trajectories, we consider the distance error, that is, the Euclidean distance from the blocks to the target zones at the final block configuration of an end-effector trajectory. As there are a total of four push sequences  $\beta$  (see Figure 2) we use the expected entropy
$$
\mathbb {E} _ {p (\mathbf {o} _ {0})} \left[ \mathcal {H} (\beta | \mathbf {o} _ {0}) \right] \approx - \frac {1}{N _ {0}} \sum_ {\mathbf {o} _ {0} \sim p (\mathbf {o} _ {0})} \sum_ {\beta} p (\beta | \mathbf {o} _ {0}) \log_ {4} p (\beta | \mathbf {o} _ {0}),
$$
to quantify a model's ability to cover the modes in the data distribution. Please note that we use  $\log_4$  for the purpose of enhancing interpretability, as it ensures  $\mathcal{H}(\beta) \in [0,1]$ . An entropy value
Figure 8: Block pushing: Top view of four different push sequences. Starting from the black cross, the gray line visualizes the end-effector trajectory of the robot manipulator. The small rectangles indicate different box configurations in the push sequence, the big rectangles mark the target zones.
of 0 signifies a policy that consistently executes the same behavior, while an entropy value of 1 represents a diverse policy that executes all behaviors with equal probability and hence matches the true behavior distribution by design of the data collection process. Furthermore, we set  $p(\mathbf{o}_0) = 1 / 30$  as we sample 30 block configurations uniformly from a configuration space (see Figure 7). For each  $\mathbf{o}_0$  we simulate 16 end-effector trajectories. For a given configuration, we count how often each of the four push-sequences is executed successfully and use the result to calculate a categorical distribution  $p(\beta | \mathbf{o}_0)$ . Once repeated for all 30 configurations, we compute  $\mathbb{E}_{p(\mathbf{o}_0)}[\mathcal{H}(\beta | \mathbf{o}_0)]$ .
C.1.3 Table Tennis
Dataset. The table tennis dataset contains 5000  $(\mathbf{o},\mathbf{a})$  pairs. The observations  $\mathbf{o}\in \mathbb{R}^4$  contain the coordinates of the initial and target ball position as projection on the table. Movement primitives (MPs) [31] are used to describe the joint space trajectories of the robot manipulator using two basis functions per joint and thus  $\mathbf{a}\in \mathbb{R}^{14}$ .
Metrics. To evaluate the different algorithms on the demonstrations recorded using the table tennis environment quantitatively, we employ two performance metrics: The success rate and the distance error. The success rate is the percentage of strikes where the ball is successfully returned to the opponent's side. The distance error, is the distance between the target position and landing position of the ball for successful strikes.
C.1.4 Human Subjects for Data Collection
For the obstacle avoidance as well as the block pushing experiments we used data collected by humans. We note that all human subjects included in the data collection process are individuals who are collaborating on this work. The participants did, therefore, not receive any financial compensation for their involvement in the study.
C.2 IMC Details and Hyperparameter
IMC employs a parameterized inference network and conditional Gaussian distributions to represent experts. For the latter, we use a fixed variance of 1 and parameterize the means as neural networks. For both inference network and expert means we use residual MLPs [57]. For all experiments, we use batch-size  $|\mathcal{B}| = |\mathcal{D}|$ , number of components  $N_z = 50$  and expert learning rate equal to  $5 \times 10^{-4}$ . Furthermore, we initialized all curriculum weights as  $p(\mathbf{o}_n, \mathbf{a}_n | z) = 1$ . For the table tennis and obstacle avoidance task, we found the best results using a multi-head expert parameterization (see Section E.2) where we tested  $1 - 4$  layer neural networks. We found that using 1 layer with 32 neurons performs best on the table tennis task and 2 layer with 64 neurons for the obstacle avoidance task. For the block pushing and Franka kitchen experiments, we obtained the best results using a single-head parameterization of the experts. We used 6 layer MLPs with 128 neurons for both tasks. For the inference network, we used a fixed set of parameters that are listed in Table 2. For the entropy scaling factor  $\eta$  we performed a hyperparameter sweep using Bayesian optimization. The respective values are  $\eta = 1/30$  for obstacle avoidance,  $\eta = 2$  for block pushing and Franka kitchen and  $\eta = 1$  for table tennis.
C.3 Baselines and Hyperparameter
We now briefly mention the baselines and their hyperparameters. We used Bayesian optimization to tune the most important hyperparameters.
Table 2: IMC & EM Hyperparameter.
PARAMETER	VALUE
EXPERT LEARNING RATE	10-4
EXPERT BATCHSIZE	1024
EXPERT VARIANCE (σ2)	1
INFERENCE NET HIDDEN LAYER	6
INFERENCE NET HIDDEN UNITS	256
INFERENCE NET EPOCHS	800
INFERENCE NET LEARNING RATE	10-3
INFERENCE NET BATCHSIZE	1024
Mixture of Experts trained with Expectation-Maximization (EM). The architecture of the mixture of experts model trained with EM [48] is identical to the one optimized with IMC: We employ a parameterized inference network and conditional Gaussian distributions to represent experts with the same hyperparameters as shown in Table 2. Furthermore, we initialized all responsibilities as  $p(z|\mathbf{o}) = 1 / N_z$ , where  $N_z$  is the number of components.
Mixture Density Network (MDN). The mixture density network [8] uses a shared backbone neural network with multiple heads for predicting component indices as well as the expert likelihood. For the experts, we employ conditional Gaussians with a fixed variance. The model likelihood is maximized in an end-to-end fashion using stochastic gradient ascent. We experimented with different backbones and expert architectures. However, we found that the MDN is not able to partition the input space in a meaningful way, often resulting in sub-optimal outcomes, presumably due to mode averaging. To find an appropriate model complexity we tested up to 50 expert heads. We found that the number of experts heads did not significantly influence the results, further indicating that the MDN is not able to utilize multiple experts to solve sub-tasks. We additionally experimented with a version of the MDN that adds an entropy bonus to the objective [58] to encourage more diverse and multimodal solutions. However, we did not find significant improvements compared to the standard version of the MDN. For a list of hyperparameter choices see 3.
Table 3: MDN Hyperparameter. The 'Value' column indicates sweep values for the obstacle avoidance task, the block pushing task, the Franka kitchen task and the table tennis task (in this order).
PARAMETER	SWEEP	VALUE
EXPERT HIDDEN LAYER	{1,2}	1,1,1,1
EXPERT HIDDEN UNITS	{30,50}	50,30,30,50
BACKBONE HID. LAYER	{2,3,4,6,8,10}	3,2,4,3
BACKBONE HID. UNITS	{50,100,150,200}	200,200,200,200
LEARNING RATE ×10-3	[0.1,1]	5.949,7.748,1.299,2.577
EXPERT VARIANCE (σ2)	-	1
MAX. EPOCHS	-	2000
BATCHSIZE	-	512
Denoising Diffusion Probabilistic Models (DDPM). We consider the denoising diffusion probabilistic model proposed by [24]. Following common practice we parameterize the model as residual MLP [27] with a sinusoidal positional encoding [54] for the diffusion steps. Moreover, we use the cosine-based variance scheduler proposed by [59]. For further details on hyperparameter choices see Table 4.
Normalizing Flow (NF). For all experiments, we build the normalizing flow by stacking masked autoregressive flows [60] paired with permutation layers [18]. As base distribution, we use a conditional isotropic Gaussian. Following common practice, we optimize the model parameters by maximizing its likelihood. See Table 5 for a list of hyperparameters.
Conditional Variational Autoencoder (CVAE). We consider the conditional version of the autoencoder proposed in [20]. We parameterize the encoder and decoder with a neural network with mirrored
Table 4: DDPM Hyperparameter. The 'Value' column indicates sweep values for the obstacle avoidance task, the block pushing task, the Franka kitchen task, and the table tennis task (in this order).
PARAMETER	SWEEP	VALUE
HIDDEN LAYER	{4,6,8,10,12}	6,6,8,6
HIDDEN UNITS	{50,100,150,200}	200,150,200,200
DIFFUSION STEPS	{5,15,25,50}	15,15,15,15
VARIANCE SCHEDULER	-	COSINE
LEARNING RATE	-	10-3
MAX. EPOCHS	-	2000
BATCHSIZE	-	512
Table 5: NF Hyperparameter. The 'Value' column indicates sweep values for the obstacle avoidance task, the block pushing task, the Franka kitchen task and the table tennis task (in this order).
PARAMETER	SWEEP	VALUE
NUM. FLOWS	{4,6,8,10,12}	6,6,4,4
HIDDEN UNITS PER FLOW	{50,100,150,200}	100,150,200,150
LEARNING RATE ×10-4	[0.01,10]	7.43,4.5,4.62,7.67
MAX. EPOCHS	-	2000
BATCHSIZE	-	512