SlowGuess/NeurIPS_2023_OCR huggingface dataset:
1: Require: Data  $\mathcal{D} = \{(\mathbf{o}_n,\mathbf{a}_n)\}_{n = 1}^N$
2: Require: Number of components  $N_{z}$
3: Require: Curriculum pacing  $\eta$
4: while  $|\Delta L| \leq \epsilon$  do
5: Perform E-step:
6:  $q(z|\mathbf{o}_n,\mathbf{a}_n)\gets \tilde{p} (\mathbf{o}_n,\mathbf{a}_n|z) / \sum_z\tilde{p} (\mathbf{o}_n,\mathbf{a}_n|z)\forall n$
7: for  $z\gets 1,\dots,N_z$  do
8: Perform M-step for curricula:
9:  $R_{z}(\mathbf{o}_{n},\mathbf{a}_{n})\gets \log p_{\theta_{z}}(\mathbf{a}_{n}|\mathbf{o}_{n},z) + \eta \log q(z|\mathbf{o}_{n},\mathbf{a}_{n})\forall n$
0:  $\tilde{p} (\mathbf{o}_n,\mathbf{a}_n|z)\gets \exp \left(R_z(\mathbf{o}_n,\mathbf{a}_n) / \eta\right)\forall n$
11: Perform M-step for experts:
12:  $\pmb{\theta}_{z}^{*}\gets \arg \max_{\pmb{\theta}_{z}}\sum_{n}\tilde{p} (\mathbf{o}_{n},\mathbf{a}_{n}|z)\log p_{\pmb{\theta}_{z}}(\mathbf{a}_{n}|\mathbf{o}_{n},z)$
13: end for
14: end while
known as self-paced learning (SPL). Pioneering work in SPL was done in [47] which is related to our work in that the authors propose to update the curriculum as well as model parameters iteratively. However, their method is based on maximum likelihood which is different from our approach. Moreover, their algorithm is restricted to latent structural support vector machines. For a comprehensive survey on curriculum learning, the reader is referred to [34].
We briefly outline the key aspects of our experimental setup.
Experimental Setup. For all experiments, we employ conditional Gaussian expert policies, i.e.,  $p_{\theta_z}(\mathbf{a}|\mathbf{o},z) = \mathcal{N}(\mathbf{a}|\boldsymbol{\mu}_{\theta_z}(\mathbf{o}),\sigma^2\mathbf{I})$ . Please note that we parameterize the expert means  $\boldsymbol{\mu}_{\theta_z}$  using neural networks. For complex high-dimensional tasks, we share features between different expert policies by introducing a deep neural network backbone. Moreover, we use a fixed variance of  $\sigma^2 = 1$  and  $N_z = 50$  components for all experiments and tune the curriculum pacing  $\eta$ . For more details see Appendix C.2.
Baselines. We compare our method to state-of-the-art generative models including denoising diffusion probabilistic models (DDPM) [24], normalizing flows (NF) [18] and conditional variational autoencoders (CVAE) [20]. Moreover, we consider energy-based models for behavior learning (IBC) [17] and the recently proposed behavior transformer (BeT) [22]. Lastly, we compare against mixture of experts trained using expectation maximization (EM) [48] and backpropagation (MDN) [8] and the ML-Cur algorithm [30]. We extensively tune the hyperparameters of the baselines using Bayesian optimization [49] on all experiments.
Figure 2: Obstacle Avoidance: Visualization of 100 end-effector trajectories for all trained models. Every method is trained on diverse demonstration data (left). MDN, NF, and IBC mostly fail to successfully solve the task. Other methods (EM, DDPM, CVAE, BET, ML-Cur) either fail to have high success rates or disregard modes in the data. Only IMC is able to perform well on both metrics.
Table 1: Performance Table: We compare IMC against strong baselines on three complex robot manipulation tasks. The best results use bold formatting.  $\uparrow$  indicates that higher and  $\downarrow$  that lower values are better. For further details, please refer to the accompanying text.
OBSTACLE AVOIDANCE	BLOCK PUSHING	TABLE TENNIS
SUCCEEDRAT(↑)	ENTROPY(↑)	SUCCEEDRAT(↑)	ENTROPY(↑)	DISTANCE ERROR(↓)	SUCCEEDRAT(↑)	DISTANCE ERROR(↓)
MDN	0.200±0.421	0.000±0.000	0.000±0.000	0.000±0.000	0.360±0.005	0.031±0.013	0.549±0.056
EM	0.675±0.033	0.902±0.035	0.458±0.048	0.707±0.042	0.127±0.011	0.725±0.042	0.220±0.012
DDPM	0.719±0.075	0.638±0.079	0.516±0.036	0.713±0.043	0.123±0.007	0.866±0.010	0.185±0.007
NF	0.313±0.245	0.349±0.208	0.001±0.001	0.000±0.000	0.346±0.034	0.422±0.035	0.371±0.013
CVAE	0.853±0.113	0.465±0.183	0.505±0.089	0.162±0.071	0.123±0.027	0.620±0.050	0.320±0.010
IBC	0.379±0.411	0.098±0.131	0.482±0.026	0.758±0.022	0.123±0.007	0.567±0.030	0.310±0.010
BET	0.504±0.076	0.837±0.066	0.374±0.041	0.607±0.037	0.151±0.008	0.758±0.025	0.235±0.011
ML-CUR	0.454±0.223	0.035±0.024	0.000±0.000	0.000±0.000	0.408±0.030	0.836±0.020	0.181±0.011
IMC	0.855±0.053	0.930±0.031	0.521±0.045	0.654±0.041	0.120±0.014	0.870±0.017	0.153±0.007
Evaluation. For all experiments we perform multiple simulations using the trained policies to compute the performance metrics. Firstly, we report the success rate, which is the fraction of simulations that led to successful task completion, reflecting the susceptibility to mode averaging. Secondly, we compute a task-specific (categorical) distribution over pre-defined behaviors that lead to successful task completion. Thereafter, we compute the entropy over this distribution to quantify the diversity in the learned behaviors and therefore the ability to cover multiple modes present in the data distribution. An entropy of 0 implies that a model executes the same behavior, while higher entropy values indicate that the model executes different behaviors. Further details are provided in the descriptions of individual tasks.
We report the mean and the standard deviation over ten random seeds for all experiments. For a detailed explanation of tasks, data, performance metrics, and hyperparameters see Appendix C. For further experimental results see Appendix E. The code is available online  $^{2}$ .
5.1 Obstacle Avoidance
The obstacle avoidance environment is visualized in Figure 1 (left) and consists of a seven DoF Franka Emika Panda robot arm equipped with a cylindrical end effector simulated using the MuJoCo physics engine [50]. The goal of the task is to reach the green finish line without colliding with one of the six obstacles. The dataset contains four human demonstrations for all ways of avoiding obstacles and completing the task. It is collected using a game-pad controller and inverse kinematics (IK) in the xy-plane amounting to  $7.3\mathrm{k}$  ( $\mathbf{o}, \mathbf{a}$ ) pairs. The observations  $\mathbf{o} \in \mathbb{R}^4$  contain the end-effector position and velocity of the robot. The actions  $\mathbf{a} \in \mathbb{R}^2$  represent the desired position of the robot. There are 24 different ways of avoiding obstacles and reaching the green finish line, each of which we define as a different behavior  $\beta$ . At test time, we perform 1000 simulations for computing the success rate and the entropy  $\mathcal{H}(\beta) = -\sum_{\beta} p(\beta) \log_{24} p(\beta)$ . Please note that we use  $\log_{24}$  for the purpose of
$^{2}$ https://github.com/ALRhub/imc
(a)  $N_{z} = 1$
(b)  $N_{z} = 3$
Figure 3: Curriculum Visualization: Visualization of the curricula  $p(\mathbf{o}, \mathbf{a} | z)$  for a different number of components  $N_z$  on the obstacle avoidance task. The color indicates different components  $z$  and the size of the dots is proportional to  $p(\mathbf{o}_n, \mathbf{a}_n | z)$ . For  $N_z = 1$  we observe that the model ignores most samples in the dataset as a single expert is not able to achieve high log-likelihood values  $p_{\theta_z}$  on all samples (Figure 3a). Adding more components to the model results in higher coverage of samples as shown in Figure 3b-3d. Figure 3e visualizes all samples contained in the dataset.
(c)  $N_{z} = 5$
(d)  $N_{z} = 10$
(e) Dataset
enhancing interpretability, as it ensures  $\mathcal{H}(\beta) \in [0,1]$ . An entropy value 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. The results are shown in Table 1. Additionally, we provide a visualization of the learned curriculum in Figure 3. Further details are provided in Appendix C.1.1.
5.2 Block Pushing
The block pushing environment is visualized in Figure 1 (middle) and uses the setup explained in Section 5.1. The robot manipulator is tasked to push blocks into target zones. Having two blocks and target zones amounts to four different push sequences (see e.g. Figure 8), each of which we define as a different behavior  $\beta$ . Using a gamepad, we recorded 500 demonstrations for each of the four push sequences with randomly sampled initial block configurations  $\mathbf{o}_0$  (i.e., initial positions and orientations), amounting to a total of  $463\mathrm{k}$  ( $\mathbf{o},\mathbf{a}$ ) pairs. The observations  $\mathbf{o} \in \mathbb{R}^{16}$  contain information about the robot's state and the block configurations. The actions  $\mathbf{a} \in \mathbb{R}^2$  represent the desired position of the robot. We evaluate the models using three different metrics: Firstly, the success rate which is the proportion of trajectories that manage to push both boxes to the target zones. Secondly, the entropy, which is computed over different push sequences  $\beta$ . Since high entropy values can be achieved by following different behaviors for different initial block configurations  $\mathbf{o}_0$  in a deterministic fashion, the entropy of  $p(\beta)$  can be a poor metric for quantifying diversity. Hence, we evaluate the expected entropy conditioned on the initial state  $\mathbf{o}_0$ , i.e.,  $\mathbb{E}_{p(\mathbf{o}_0)}[\mathcal{H}(\beta|\mathbf{o}_0)] \approx -\frac{1}{N_0}\sum_{\mathbf{o}_0 \sim p(\mathbf{o}_0)}\sum_{\beta}p(\beta|\mathbf{o}_0)\log_4p(\beta|\mathbf{o}_0)$ . If, for the same  $\mathbf{o}_0$ , all behaviors can be achieved, the expected entropy is high. In contrast, the entropy is 0 if the same behavior is executed for the same  $\mathbf{o}_0$ . Here,  $p(\mathbf{o}_0)$  and  $N_0$  denote the distribution over initial block configurations and the number of samples respectively. See Appendix C.1.2 for more details. Lastly, we quantify the performance on non-successful trajectories, via distance error, i.e., the distance from the blocks to the target zones at the end of a trajectory. The success rate and distance error indicate whether a model is able to avoid averaging over different behaviors. The entropy assesses the ability to represent multimodal data distributions by completing different push sequences. The results are reported in Table 1 and generated by simulating 16 evaluation trajectories for 30 different initial block configurations  $\mathbf{o}_0$  per seed. The difficulty of the task is reflected by the low success rates of most models. Besides being a challenging manipulation task, the high task complexity is caused by having various sources of multimodality in the data distribution: First, the inherent versatility in human behavior. Second, multiple human demonstrators, and lastly different push sequences for the same block configuration.
5.3 Franka Kitchen
The Franka kitchen environment was introduced in [51] and uses a seven DoF Franka Emika Panda robot with a two DoF gripper to interact with a simulated kitchen environment. The corresponding dataset contains 566 human-collected trajectories recorded using a virtual reality setup amounting to  $128\mathrm{k}$  ( $\mathbf{o}, \mathbf{a}$ ) pairs. Each trajectory executes a sequence completing four out of seven different tasks. The observations  $\mathbf{o} \in \mathbb{R}^{30}$  contain information about the position and orientation of the task-relevant objects in the environment. The actions  $\mathbf{a} \in \mathbb{R}^{9}$  represent the control signals for
Figure 4: Franka Kitchen: Performance comparison between various policy learning algorithms, evaluating the success rate and entropy for a varying number of task completions. IMC is able to outperform the baselines in both metrics, achieving a high success rate while showing a higher diversity in the task sequences. Performance comparison between various policy learning algorithms. IMC is more successful in completing tasks (success rate) while at the same time having the highest diversity in the sequence of task completions (entropy).
the robot and gripper. To assess a model's ability to avoid mode averaging we again use the success rate over the number of tasks solved within one trajectory. For each number of solved tasks  $\in \{1,2,3,4\}$ , we define a behavior  $\beta$  as the order in which the task is completed and use the entropy  $\mathcal{H}(\beta) = -\sum_{\beta}p(\beta)\log p(\beta)$  to quantify diversity. The results are shown in Figure 5 and are generated using 100 evaluation trajectories for each seed. There are no results reported for IBC, MDN and ML-Cur as we did not manage to obtain reasonable results.
5.4 Table Tennis
The table tennis environment is visualized in Figure 1 (right) and consists of a seven DOF robot arm equipped with a table tennis rack and is simulated using the MuJoCo physics engine. The goal is to return the ball to varying target positions after it is launched from a randomized initial position. Although not collected by human experts, the 5000 demonstrations are generated using a reinforcement learning (RL) agent that is optimized for highly multimodal behavior such as backhand and forehand strokes [52]. Each demonstration consists of an observation  $\mathbf{o} \in \mathbb{R}^4$  defining the initial and target ball position. 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}$ . We evaluate the model performance using the success rate, that is, how frequently the ball is returned to the other side. Moreover, we employ the distance error, i.e., the Euclidean distance from the landing position of the ball to the target position. Both metrics reflect if a model is able to avoid averaging over different movements. For this experiment, we do not report the entropy as we do not know the various behaviors executed by the RL agent. The results are shown in Table 1 and are generated using 500 different initial and target positions. Note that the reinforcement learning agent used to generate the data achieves an average success rate of 0.91 and a distance error of 0.14 which is closely followed by IMC.
5.5 Ablation Studies
Additionally, we compare the performance of IMC with EM for a varying number of components on the obstacle avoidance and table tennis task. The results are shown in Figure 4 and highlight the properties of the moment and information projection: Using limited model complexity, e.g. 1 or 5 components, EM suffers from mode averaging, resulting in poor performances (Figure 5a and Figure 5b). This is further illustrated in Figure 5c. In contrast, the zero forcing property of the information projection allows IMC to avoid mode averaging (see Figure 5d) which is reflected in the success rates and distance error for a small number of components. The performance gap between EM and IMC for high model complexities suggests that EM still suffers from averaging problems. Moreover, the results show that IMC needs fewer components to achieve the same performance as EM.
(a) Obstacle avoidance performance comparison.
(c) Obstacle avoidance end-effector trajectories for EM.
Figure 5: Obstacle Avoidance: IMC and EM improve the success rate and entropy with an increasing number of components (a). For a small number of components, IMC archives a high success rate, as it allows the policy to focus on data that it can represent. In contrast, the policy trained with EM fails as it is forced to cover the whole data. This is visualized in the end-effector trajectories  $(c + d)$ . Similar observations can be made for Table Tennis: Both performance metrics increase with a higher number of components. IMC manages to achieve good performance with a small number of components.
(b) Table tennis performance comparison.
(d) Obstacle avoidance end-effector trajectories for IMC.
We presented Information Maximizing Curriculum (IMC), a novel approach for conditional density estimation, specifically designed to address mode-averaging issues commonly encountered when using maximum likelihood-based optimization in the context of multimodal density estimation. IMC's focus on mitigating mode-averaging is particularly relevant in imitation learning from human demonstrations, where the data distribution is often highly multimodal due to the diverse and versatile nature of human behavior.
IMC uses a curriculum to assign weights to the training data allowing the policy to focus on samples it can represent, effectively mitigating the mode-averaging problem. We extended our approach to a mixture of experts (MoE) policy, where each mixture component selects its own subset of the training data for learning, allowing the model to imitate the rich and versatile behavior present in the demonstration data.
Our experimental results demonstrate the superior performance of our method compared to state-of-the-art policy learning algorithms and mixture of experts (MoE) policies trained using competing optimization algorithms. Specifically, on complex multimodal simulated control tasks with data collected from human demonstrators, our method exhibits the ability to effectively address two key challenges:  $i$  ) avoiding mode averaging and  $ii$  ) covering all modes present in the data distribution.
Limitations. While our current approach achieves state-of-the-art performance, there are still areas for improvement in parameterizing our model. Presently, we employ simple multilayer perceptrons to parameterize the expert policies. However, incorporating image observations would require a convolutional neural network (CNN) [53] backbone. Additionally, our current model relies on the Markov assumption, but relaxing this assumption and adopting history-based models like transformers [54] could potentially yield significant performance improvements. Lastly, although this work primarily concentrates on continuous domains, an intriguing prospect for future research would be to explore the application of IMC in discrete domains.
Broader Impact. Improving imitation learning algorithms hold the potential to enhance the accessibility of robotic systems in real-world applications, with both positive and negative implications. We acknowledge that identifying and addressing any potential adverse effects resulting from the deployment of these robotic systems is a crucial responsibility that falls on sovereign governments.
[1] Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, Jan Peters, et al. An algorithmic perspective on imitation learning. Foundations and Trends® in Robotics, 7(1-2):1-179, 2018.
[2] Leslie Pack Kaelbling, Michael L Littman, and Andrew W Moore. Reinforcement learning: A survey. Journal of artificial intelligence research, 4:237-285, 1996.
[3] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[4] Kevin P Murphy. Machine learning: a probabilistic perspective. MIT press, 2012.
[5] Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathematical statistics, 22(1):79-86, 1951.
[6] Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, volume 4. Springer, 2006.
[7] Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the Royal Statistical Society: Series B (Methodological), 39(1):1-22, 1977.
[8] Christopher M Bishop. Mixture density networks. 1994.
[9] Stephen Boyd, Stephen P Boyd, and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
[10] Oleg Arenz, Gerhard Neumann, and Mingjun Zhong. Efficient gradient-free variational inference using policy search. In International conference on machine learning, pages 234-243. PMLR, 2018.
[11] Brenna D Argall, Sonia Chernova, Manuela Veloso, and Brett Browning. A survey of robot learning from demonstration. Robotics and autonomous systems, 57(5):469-483, 2009.
[12] Kevin Zakka, Andy Zeng, Pete Florence, Jonathan Tompson, Jeannette Bohg, and Debidatta Dwibedi. Xirl: Cross-embodiment inverse reinforcement learning. In Conference on Robot Learning, pages 537–546. PMLR, 2022.
[13] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaii, volume 8, pages 1433-1438. Chicago, IL, USA, 2008.
[14] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
[15] Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforcement learning. In International Conference on Learning Representations, 2018.
[16] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural information processing systems, pages 2672–2680, 2014.
[17] Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pages 158-168. PMLR, 2022.
[53] Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard E Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4):541-551, 1989.
[54] Ashish Vavwani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[55] Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909, 2018.
[56] Dimitri P Bertsekas. Nonlinear programming. Journal of the Operational Research Society, 48(3):334-334, 1997.
[57] Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32, 2019.
[58] You Zhou, Jianfeng Gao, and Tamim Asfour. Movement primitive learning and generalization: Using mixture density networks. IEEE Robotics & Automation Magazine, 27(2):22-32, 2020.
[59] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[60] George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density estimation. Advances in neural information processing systems, 30, 2017.
[61] Irina Higgins, Loic Matthew, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In International conference on learning representations, 2017.
[62] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
A.1 Proof of Proposition 3.1 and 3.3
Convergence of the Single Expert Objective (Proposition 3.1). We perform coordinate ascent on  $\tilde{J}$  which is guaranteed to converge to a stationary point if updating each coordinate results in a monotonic improvement of [9]. For fixed expert parameters we find the unique  $p(\mathbf{o},\mathbf{a})$  that maximizes  $\tilde{J}$  [55] (see Section 3.1) and hence we have  $\tilde{J} (p(\mathbf{o},\mathbf{a})^{(i)},\pmb {\theta})\geq \tilde{J} (p(\mathbf{o},\mathbf{a})^{(i - 1)},\pmb {\theta})$  where  $i$  denotes the iteration. Under suitable assumptions  $(\log p_{\theta}$  is differentiable, its gradient is  $L$  -Lipschitz,  $\theta_z$  is updated using gradient ascent and the learning rate is chosen such that the descent lemma [56] holds), it holds that  $\tilde{J} (p(\mathbf{o},\mathbf{a}),\pmb{\theta}^{(i)})\geq \tilde{J} (p(\mathbf{o},\mathbf{a}),\pmb{\theta}^{(i - 1)})$  . Hence, we are guaranteed to converge to a stationary point of  $\tilde{J}$