Title: Information Maximizing Curriculum: A Curriculum-Based Approach for Learning Versatile Skills
Venue: NeurIPS 2023
Accepted: yes
Keywords: Imitation Learning

Verstile Skill Learning

Curriculum Learning
Abstract: Imitation learning uses data for training policies to solve complex tasks. However,
when the training data is collected from human demonstrators, it often leads
to multimodal distributions because of the variability in human actions. Most
imitation learning methods rely on a maximum likelihood (ML) objective to learn
a parameterized policy, but this can result in suboptimal or unsafe behavior due
to the mode-averaging property of the ML objective. In this work, we propose
Information Maximizing Curriculum, a curriculum-based approach that assigns
a weight to each data point and encourages the model to specialize in the data it
can represent, effectively mitigating the mode-averaging problem by allowing the
model to ignore data from modes it cannot represent. To cover all modes and thus,
enable versatile behavior, we extend our approach to a mixture of experts (MoE)
policy, where each mixture component selects its own subset of the training data
for learning. A novel, maximum entropy-based objective is proposed to achieve
full coverage of the dataset, thereby enabling the policy to encompass all modes
within the data distribution. We demonstrate the effectiveness of our approach on
complex simulated control tasks using versatile human demonstrations, achieving
superior performance compared to state-of-the-art methods.

Information Maximizing Curriculum: A Curriculum-Based Approach for Imitating Diverse Skills

Abstract

Imitation learning uses data for training policies to solve complex tasks. However, when the training data is collected from human demonstrators, it often leads to multimodal distributions because of the variability in human actions. Most imitation learning methods rely on a maximum likelihood (ML) objective to learn a parameterized policy, but this can result in suboptimal or unsafe behavior due to the mode-averaging property of the ML objective. In this work, we propose Information Maximizing Curriculum , a curriculum-based approach that assigns a weight to each data point and encourages the model to specialize in the data it can represent, effectively mitigating the mode-averaging problem by allowing the model to ignore data from modes it cannot represent. To cover all modes and thus, enable diverse behavior, we extend our approach to a mixture of experts (MoE) policy, where each mixture component selects its own subset of the training data for learning. A novel, maximum entropy-based objective is proposed to achieve full coverage of the dataset, thereby enabling the policy to encompass all modes within the data distribution. We demonstrate the effectiveness of our approach on complex simulated control tasks using diverse human demonstrations, achieving superior performance compared to state-of-the-art methods.

1 Introduction

Equipping agents with well-performing policies has long been a prominent focus in machine learning research. Imitation learning (IL) [1] offers a promising technique to mimic human behavior by leveraging expert data, without the need for intricate controller design, additional environment interactions, or complex reward shaping to encode the target behavior. The latter are substantial advantages over reinforcement learning techniques [2, 3] that rely on reward feedback. However, a significant challenge in IL lies in handling the multimodal nature of data obtained from human demonstrators, which can stem from differences in preferences, expertise, or problem-solving strategies. Conventionally, maximum likelihood estimation (MLE) is employed to train a policy on expert data. It is well-known that MLE corresponds to the moment projection [4], causing the policy to average over modes in the data distribution that it cannot represent. Such mode averaging can lead to unexpected and potentially dangerous behavior . We address this critical issue by introducing Information Maximizing Curriculum (IMC), a novel curriculum-based approach.

In IMC, we view imitation learning as a conditional density estimation problem and present a mathematically sound weighted optimization scheme. Data samples are assigned curriculum weights, which are updated using an information projection. The information projection minimizes the reverse

KL divergence, forcing the policy to ignore modes it cannot represent [4]. As a result, the optimized policy ensures safe behavior while successfully completing the task.

Yet, certain modes within the data distribution may remain uncovered with a single expert policy. To address this limitation and endow the policy with the diverse behavior present in the expert data, we extend our approach to a mixture of expert (MoE) policy. Each mixture component within the MoE policy selects its own subset of training data, allowing for specialization in different modes. Our objective maximizes the entropy of the joint curriculum, ensuring that the policy covers all data samples.

We show that our method is able to outperform state-of-the-art policy learning algorithms and MoE policies trained using competing optimization algorithms on complex multimodal simulated control tasks where data is collected by human demonstrators. In our experiments, we assess the ability of the models to i) avoid mode averaging and ii) cover all modes present in the data distribution.

2 Preliminaries

Our approach heavily builds on minimizing Kullback-Leibler divergences as well as mixtures of expert policies. Hence, we will briefly review both concepts.

2.1 Moment and Information Projection

The Kullback-Leibler (KL) divergence [5] is a similarity measure for probability distributions and is defined as D KL ( p ∥ p ′ ) = ∑ x p ( x ) log p ( x ) /p ′ ( x ) for a discrete random variable x . Due to its asymmetry, the KL divergence offers two different optimization problems for fitting a model distribution p to a target distribution p ∗ [4], that is,

arg min

p

D

KL

(

p

∗

∥

p

)

︸

︷︷

︸

M(oment)-Projection

p

p

∗

and

arg min

p

D

KL

(

p

∥

p

∗

)

︸

︷︷

︸

I(nformation)-Projection

.

p

p

∗

The M-projection - or equivalently maximum likelihood estimation (MLE) [6] - is probability forcing , meaning that the model is optimized to match the moments of the target distribution, causing it to average over modes that it cannot represent. In contrast, the I-projection is zero forcing which leads the model to ignore modes of the target distribution that it is not able to represent. The I-projection can be rewritten as maximum entropy problem, i.e.,

Using this formulation, it can be seen that the optimization balances between fitting the target distribution and keeping the entropy H ( x ) = -∑ x p ( x ) log p ( x ) high.

2.2 Mixtures of Expert Policies

Mixtures of expert policies are conditional discrete latent variable models. Given some observation o ∈ O and action a ∈ A , the marginal likelihood is decomposed into individual components, i.e.,

Here, O and A denote the observation and action space respectively, and z stands for the discrete latent variable that indexes distinct components within the mixture.

The gating p ( z | o ) is responsible for soft-partitioning the observation space O into sub-regions where the corresponding experts p ( a | o , z ) approximate the target density. Typically the experts and the gating are parameterized and learned by maximizing the marginal likelihood via expectationmaximization (EM) [7] or gradient ascent [8]. In order to sample actions, that is, a ′ ∼ p ( a | o ′ ) for some observation o ′ , we first sample a component index from the gating, i.e., z ′ ∼ p ( z | o ′ ) . The component index selects the respective expert to obtain a ′ ∼ p ( a | o ′ , z ′ ) .

3 Information Maximizing Curriculum

In this section, we propose Information Maximizing Curriculum (IMC), a novel algorithm for training mixtures of expert polices. We motivate our optimization objective using a single policy. Next, we generalize the objective to support learning a mixture of experts policy. Thereafter, we discuss the optimization scheme and provide algorithmic details.

3.1 Objective for a Single Expert Policy

We propose an objective that jointly learns a curriculum p ( o , a ) and a parameterized policy p θ ( a | o ) with parameters θ . The curriculum is a categorical distribution over samples of a dataset { ( o n , a n ) } N n =1 , assigning probability mass to samples according to the performance of the policy. To allow the curriculum to ignore samples that the policy cannot represent, we build on the I-projection (see Equation 1). We, therefore, formulate the objective function as

which is optimized for p ( o , a ) and θ in an alternating fashion using coordinate ascent [9]. We additionally introduce a trade-off factor η that determines the pacing of the curriculum. For η →∞ the curriculum becomes uniform, exposing all samples to the policy and hence reducing to maximum likelihood estimation for θ . In contrast, if η → 0 the curriculum concentrates on samples where the policy log-likelihood log p θ ( a | o ) is highest. The objective can be solved in closed form for p ( o , a ) (see Appendix B), resulting in

Maximizing the objective w.r.t θ reduces to a weighted maximum-likelihood estimation, that is,

The optimization is repeated until reaching a local maximum, indicating that the curriculum has selected a fixed set of samples where the policy attains high log-likelihoods log p θ ( a | o ) . Proposition 3.1 establishes convergence guarantees, the details of which are elaborated upon in Appendix A.1.

Proposition 3.1. Let ˜ J be defined as in Equation 2 and 0 < η < ∞ . Under mild assumptions on p θ and the optimization scheme for θ , it holds for all i , denoting the iteration index of the optimization process, that

where equality indicates that the algorithm has converged to a local optimum.

The capacity to disregard samples where the policy cannot achieve satisfactory performance mitigates the mode-averaging problem. Nevertheless, a drawback of employing a single expert policy is the potential for significantly suboptimal performance on the ignored samples. This limitation is overcome by introducing multiple experts, each specializing in different subsets of the data.

3.2 Objective for a Mixture of Experts Policy

Assuming limited complexity, a single expert policy is likely to ignore a large amount of samples due to the zero-forcing property of the I-projection. Using multiple curricula and policies that specialize to different subsets of the data is hence a natural extension to the single policy model. To that end, we make two major modifications to Equation 2: Firstly, we use a mixture model with multiple components z where each component has its own curriculum, i.e., p ( o , a ) = ∑ z p ( z ) p ( o , a | z ) . Secondly, we employ an expert policy per component p θ z ( a | o , z ) , that is paced by the corresponding curriculum p ( o , a | z ) . The resulting objective function is given by

where ψ summarizes the dependence on p ( z ) , { p ( o , a | z ) } z and { θ z } z . However, Equation 3 is difficult to optimize as the entropy of the mixture model prevents us from updating the curriculum of each component independently. Similar to [10], we introduce an auxiliary distribution q ( z | o , a ) to decompose the objective function into a lower bound L ( ψ , q ) and an expected D KL term, that is,

with p ( z | o , a ) = p ( o , a | z ) p ( z ) /p ( o , a ) and

̸

with R z ( o n , a n ) = log p θ z ( a n | o n , z ) + η log q ( z | o n , a n ) , allowing for independent updates for p ( o , a | z ) and θ z by maximizing the per-component objective function J z ( p ( o , a | z ) , θ z ) . Aderivation can be found in Appendix B.1. Since E p ( o , a ) D KL ( p ( z | o , a ) ∥ q ( z | o )) ≥ 0 , L is a lower bound on J for η ≥ 0 . Please note that the per-component objective function J z is very similar to Equation 2, with J z including an additional term, log q ( z | o , a ) , which serves the purpose of preventing different curricula from assigning probability mass to the same set of samples: Specifically, a component z is considered to 'cover' a datapoint ( o n , a n ) when q ( z | o n , a n ) ≈ 1 . Since ∑ z q ( z | o n , a n ) = 1 , it follows that for other components z ′ = z it holds that q ( z ′ | o n , a n ) ≈ 0 . Consequently, log q ( z ′ | o n , a n ) →-∞ , implying that R z ′ ( o n , a n ) →-∞ . As a result, the other curricula effectively disregard the datapoint, as p ( o n , a n | z ′ ) ∝ exp R z ′ ( o n , a n ) ≈ 0 .

We follow the optimization scheme of the expectation-maximization algorithm [7], that is, we iteratively maximize (M-step) and tighten the lower bound (E-step) L ( ψ , q ) .

3.3 Maximizing the Lower Bound (M-Step)

We maximize the lower bound L ( ψ , q ) with respect to the mixture weights p ( z ) , curricula p ( o , a | z ) and expert policy parameters θ z . We find closed form solutions for both, p ( z ) and p ( o , a | z ) given by where ˜ p ( o n , a n | z ) are the optimal unnormalized curricula, such that holds p ∗ ( o n , a n | z ) = ˜ p ( o n , a n | z ) / ∑ n ˜ p ( o n , a n | z ) . However, due to the hierarchical structure of L ( ψ , q ) we implicitly optimize for p ( z ) when updating the curricula. This result is frequently used throughout this work and formalized in Proposition 3.2. A proof can be found in Appendix A.2.

Proposition 3.2. Let p ∗ ( z ) and ˜ p ( o , a | z ) be the optimal mixture weights and unnormalized curricula for maximizing L ( ψ , q ) . It holds that

The implicit updates of the mixture weights render the computation of p ∗ ( z ) obsolete, reducing the optimization to computing the optimal (unnormalized) curricula ˜ p ( o , a | z ) and expert policy parameters θ ∗ z . In particular, this result allows for training the policy using mini-batches and thus greatly improves the scalability to large datasets as explained in Section 3.5. Maximizing the lower bound with respect to θ ∗ z results in a weighted maximum likelihood estimation, i.e.,

where the curricula ˜ p ( o n , a n | z ) assign sample weights. For further details on the M-step, including derivations of the closed-form solutions and the expert parameter objective see Appendix B.2.

3.4 Tightening the Lower Bound (E-Step)

Tightening of the lower bound (also referred to as E-step) is done by minimizing the expected Kullback-Leibler divergence in Equation 4. Using the properties of the KL divergence, it can easily be seen that the lower bound is tight if for all n ∈ { 1 , ..., N } q ( z | o n ) = p ( z | o n , a n ) holds. To obtain p ( z | o n , a n ) we leverage Bayes' rule, that is, p ( z | o n , a n ) = p ∗ ( z ) p ∗ ( o n , a n | z ) / ∑ z p ∗ ( z ) p ∗ ( o n , a n | z ) . Using Proposition 3.2 we find that

Please note that the lower bound is tight after every E-step as the KL divergence is set to zero. Thus, increasing the lower bound L maximizes the original objective J assuming that updates of θ z are not decreasing the expert policy log-likelihood log p θ z ( a | o , z ) .

3.5 Algorithmic Details

Convergence Guarantees. Proposition 3.3 establishes convergence guarantees for the mixture of experts policy objective J . The proof mainly relies on the facts that IMC has the same convergence guarantees as the EM algorithm and that Proposition 3.1 can be transferred to the per-component objective J z . The full proof is given in Appendix A.1.

Proposition 3.3. Let J be defined as in Equation 4 and 0 < η < ∞ . Under mild assumptions on p θ z and the optimization scheme for θ z , it holds for all i , denoting the iteration index of the optimization process, that

where equality indicates that the IMC algorithm has converged to a local optimum.

Stopping Criterion. We terminate the algorithm if either the maximum number of training iterations is reached or if the lower bound L ( ψ , q ) converges, i.e.,

with threshold ϵ and two subsequent iterations ( i ) and ( i -1) . The lower bound can be evaluated efficiently using Corollary 3.2.1.

Corollary 3.2.1. Consider the setup used in Proposition 3.2. For p ∗ ( z ) ∈ ψ and { p ∗ ( o , a | z ) } z ∈ ψ it holds that

See Appendix A.3 for a proof.

Inference. In order to perform inference, i.e., sample actions from the policy, we need to access the gating distribution for arbitrary observations o ∈ O which is not possible as p ( z | o , a ) is only defined for observations contained in the dataset o , a . We therefore leverage Corollary 3.2.2 to learn an inference network g ϕ ( z | o ) with parameters ϕ by minimizing the KL divergence between p ( z | o ) and g ϕ ( z | o ) under p ( o ) (see Appendix A.4 for a proof).

Corollary 3.2.2. Consider the setup used in Proposition 3.2. It holds that

Once trained, the inference network can be used for computing the exact log-likelihood as p ( a | o ) = ∑ z g ϕ ( z | o ) p θ z ( a | o , z ) or sampling a component index, i.e., z ′ ∼ g ϕ ( z | o ) .

Mini-Batch Updates. Due to Proposition 3.2, the M- and E-step in the training procedure only rely on the unnormalized curricula ˜ p ( o , a | z ) . Consequently, there is no need to compute the normalization constant ∑ n ˜ p ( o n , a n | z ) . This allows us to utilize mini-batches instead of processing the entire dataset, resulting in an efficient scaling capability for large datasets. Please refer to Algorithm 1 for a detailed description of the complete training procedure.

4 Related Work

Imitation Learning. A variety of algorithms in imitation learning [1, 11] can be grouped into two categories: Inverse reinforcement learning [12, 13], which extracts a reward function from demonstrations and optimizes a policy subsequently, and behavioral cloning, which directly extracts a policy from demonstrations. Many works approach the problem of imitation learning by considering behavior cloning as a distribution-matching problem, in which the state distribution induced by the policy is required to align with the state distribution of the expert data. Some methods [14, 15] are based on adversarial methods inspired by Generative Adversarial Networks (GANs) [16]. A policy is trained to imitate the expert while a discriminator learns to distinguish between fake and expert data. However, these methods are not suitable for our case as they involve interacting with the environment during training. Other approaches focus on purely offline training and use various policy representations such as energy-based models [17], normalizing flows [18, 19], conditional variational autoencoders (CVAEs) [20, 21], transformers [22], or diffusion models [23, 24, 25, 26, 27, 28]. These

initial position

target position

$..!

models can represent multi-modal expert distributions but are optimized based on the M-Projection, which leads to a performance decrease. Recent works [29, 30] have proposed training Mixture of Experts models with an objective similar to ours. However, the work by [29] requires environmentspecific geometric features, which is not applicable in our setting, whereas the work by [30] considers linear experts and the learning of skills parameterized by motion primitives [31]. For a detailed differentiation between our work and the research conducted by [30], please refer to Appendix D.1.

Curriculum Learning. The authors of [32] introduced curriculum learning (CL) as a new paradigm for training machine learning models by gradually increasing the difficulty of samples that are exposed to the model. Several studies followed this definition [33, 34, 35, 36, 37, 38, 39]. Other studies used the term curriculum learning for gradually increasing the model complexity [40, 41, 42] or task complexity [43, 44, 45, 46]. All of these approaches assume that the difficulty-ranking of the samples is known a-priori. In contrast, we consider dynamically adapting the curriculum according to the learning progress of the model which is

Algorithm 1 IMC training procedure

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].

5 Experiments

We briefly outline the key aspects of our experimental setup.

Experimental Setup. For all experiments, we employ conditional Gaussian expert policies, i.e., p θ z ( a | o , z ) = N ( a | µ θ z ( o ) , σ 2 I ) . Please note that we parameterize the expert means µ θ 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 σ 2 = 1 and N z = 50 components for all experiments and tune the curriculum pacing η . 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.

Demonstrations

MDN

EM

DDPM

NF

CVAE

IBC

BET

ML-Cur

IMC

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 k ( o , a ) pairs. The observations o ∈ R 4 contain the end-effector position and velocity of the robot. The actions a ∈ 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 β . At test time, we perform 1000 simulations for computing the success rate and the entropy H ( β ) = -∑ β p ( β ) log 24 p ( β ) . Please note that we use log 24 for the purpose of

(a)

N

z

= 1

(b)

N

z

= 3

(c)

N

z

= 5

(d)

N

z

= 10

(e) Dataset

enhancing interpretability, as it ensures H ( β ) ∈ [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 β . Using a gamepad, we recorded 500 demonstrations for each of the four push sequences with randomly sampled initial block configurations o 0 (i.e., initial positions and orientations), amounting to a total of 463 k ( o , a ) pairs. The observations o ∈ R 16 contain information about the robot's state and the block configurations. The actions a ∈ 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 β . Since high entropy values can be achieved by following different behaviors for different initial block configurations o 0 in a deterministic fashion, the entropy of p ( β ) can be a poor metric for quantifying diversity. Hence, we evaluate the expected entropy conditioned on the initial state o 0 , i.e., E p ( o 0 ) [ H ( β | o 0 ) ] ≈ -1 N 0 ∑ o 0 ∼ p ( o 0 ) ∑ β p ( β | o 0 ) log 4 p ( β | o 0 ) . If, for the same 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 o 0 . Here, p ( 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 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 k ( o , a ) pairs. Each trajectory executes a sequence completing four out of seven different tasks. The observations o ∈ R 30 contain information about the position and orientation of the task-relevant objects in the environment. The actions a ∈ R 9 represent the control signals for

EM

DDPM

NF

CVAE

BET

IMC

1

2

3

4

0

0

.

5

1

number of task completions

success rate

1

2

3

4

0

1

2

number 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 ∈ { 1 , 2 , 3 , 4 } , we define a behavior β as the order in which the task is completed and use the entropy H ( β ) = -∑ β p ( β ) log p ( β ) 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, MDNand 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 racket 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 o ∈ 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 a ∈ 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.

1

10

20

30

40

50

0

0

.

2

0

.

4

0

.

6

0

.

8

1

number of components

success rate

EM

IMC

1

10

20

30

40

50

0

0

.

2

0

.

4

0

.

6

0

.

8

1

number of components

entropy

(a) Obstacle avoidance performance comparison.

10

20

30

40

50

0

0

.

2

0

.

4

0

.

6

0

.

8

1

number of components

success rate

Expert

EM

IMC

10

20

30

40

50

0

.

2

0

.

4

0

.

6

number of components

distance error

(b) Table tennis performance comparison.

EM (

N

z

= 1)

EM (

N

z

= 5)

EM (

N

z

= 10)

IMC (

N

z

= 1)

IMC (

N

z

= 5)

IMC (

N

z

= 10)

(c) Obstacle avoidance end-effector trajectories for EM.

(d) Obstacle avoidance end-effector trajectories for IMC.

6 Conclusion

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-ofthe-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 holds 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.

Acknowledgments and Disclosure of Funding

This work was supported by funding from the pilot program Core Informatics of the Helmholtz Association (HGF). The authors acknowledge support by the state of Baden-Württemberg through bwHPC, as well as the HoreKa supercomputer funded by the Ministry of Science, Research and the Arts Baden-Württemberg and by the German Federal Ministry of Education and Research.

References

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.

Leslie Pack Kaelbling, Michael L Littman, and Andrew W Moore. Reinforcement learning: A survey. Journal of artificial intelligence research , 4:237-285, 1996.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Kevin P Murphy. Machine learning: a probabilistic perspective . MIT press, 2012.

Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathematical statistics , 22(1):79-86, 1951.

Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning , volume 4. Springer, 2006.

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.

Christopher M Bishop. Mixture density networks. 1994.

Stephen Boyd, Stephen P Boyd, and Lieven Vandenberghe. Convex optimization . Cambridge university press, 2004.

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.

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.

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.

Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai , volume 8, pages 1433-1438. Chicago, IL, USA, 2008.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems , 29, 2016.

Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse reinforcement learning. In International Conference on Learning Representations , 2018.

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.

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.

George Papamakarios, Eric T Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. J. Mach. Learn. Res. , 22(57):1-64, 2021.

Avi Singh, Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine. Parrot: Data-driven behavioral priors for reinforcement learning. In International Conference on Learning Representations , 2021.

Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems , 28, 2015.

Oier Mees, Lukas Hermann, and Wolfram Burgard. What matters in language conditioned robotic imitation learning over unstructured data. IEEE Robotics and Automation Letters , 7(4):11205-11212, 2022.

Nur Muhammad Mahi Shafiullah, Zichen Jeff Cui, Ariuntuya Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with one stone. arXiv preprint arXiv:2206.11251 , 2022.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations , 2021.

Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137 , 2023.

Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In The Eleventh International Conference on Learning Representations , 2023.

Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using score-based diffusion policies. arXiv preprint arXiv:2304.02532 , 2023.

Niklas Freymuth, Nicolas Schreiber, Aleksandar Taranovic, Philipp Becker, and Gerhard Neumann. Inferring versatile behavior from demonstrations by matching geometric descriptors. In 6th Annual Conference on Robot Learning , 2022.

Maximilian Xiling Li, Onur Celik, Philipp Becker, Denis Blessing, Rudolf Lioutikov, and Gerhard Neumann. Curriculum-based imitation of versatile skills. arXiv preprint arXiv:2304.05171 , 2023.

Alexandros Paraschos, Christian Daniel, Jan R Peters, and Gerhard Neumann. Probabilistic movement primitives. Advances in neural information processing systems , 26, 2013.

Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning , pages 41-48, 2009.

Valentin I Spitkovsky, Hiyan Alshawi, and Daniel Jurafsky. Baby steps: How 'less is more' in unsupervised dependency parsing. 2009.

Petru Soviany, Radu Tudor Ionescu, Paolo Rota, and Nicu Sebe. Curriculum learning: A survey. International Journal of Computer Vision , pages 1-40, 2022.

Xinlei Chen and Abhinav Gupta. Webly supervised learning of convolutional networks. In Proceedings of the IEEE international conference on computer vision , pages 1431-1439, 2015.

Radu Tudor Ionescu, Bogdan Alexe, Marius Leordeanu, Marius Popescu, Dim P Papadopoulos, and Vittorio Ferrari. How hard can it be? estimating the difficulty of visual search in an image. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 2157-2166, 2016.

Anastasia Pentina, Viktoriia Sharmanska, and Christoph H Lampert. Curriculum learning of multiple tasks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 5492-5500, 2015.

Yangyang Shi, Martha Larson, and Catholijn M Jonker. Recurrent neural network language model adaptation with curriculum learning. Computer Speech & Language , 33(1):136-154, 2015.

Wojciech Zaremba and Ilya Sutskever. Learning to execute. arXiv preprint arXiv:1410.4615 , 2014.

Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710.10196 , 2017.

Pietro Morerio, Jacopo Cavazza, Riccardo Volpi, René Vidal, and Vittorio Murino. Curriculum dropout. In Proceedings of the IEEE International Conference on Computer Vision , pages 3544-3552, 2017.

Samarth Sinha, Animesh Garg, and Hugo Larochelle. Curriculum by smoothing. Advances in Neural Information Processing Systems , 33:21653-21664, 2020.

Antoine Caubrière, Natalia Tomashenko, Antoine Laurent, Emmanuel Morin, Nathalie Camelin, and Yannick Estève. Curriculum-based transfer learning for an effective end-to-end spoken language understanding and domain portability. arXiv preprint arXiv:1906.07601 , 2019.

Carlos Florensa, David Held, Markus Wulfmeier, Michael Zhang, and Pieter Abbeel. Reverse curriculum generation for reinforcement learning. In Conference on robot learning , pages 482-495. PMLR, 2017.

William Lotter, Greg Sorensen, and David Cox. A multi-scale cnn and curriculum learning strategy for mammogram classification. In Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support , pages 169-177. Springer, 2017.

Nikolaos Sarafianos, Theodore Giannakopoulos, Christophoros Nikou, and Ioannis A Kakadiaris. Curriculum learning for multi-task classification of visual attributes. In Proceedings of the IEEE International Conference on Computer Vision Workshops , pages 2608-2615, 2017.

MKumar, Benjamin Packer, and Daphne Koller. Self-paced learning for latent variable models. Advances in neural information processing systems , 23, 2010.

Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation , 3(1):79-87, 1991.

Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. Advances in neural information processing systems , 25, 2012.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems , pages 5026-5033. IEEE, 2012.

Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. arXiv preprint arXiv:1910.11956 , 2019.

Onur Celik, Dongzhuoran Zhou, Ge Li, Philipp Becker, and Gerhard Neumann. Specializing versatile skill libraries using local mixture of experts. In Conference on Robot Learning , pages 1423-1433. PMLR, 2022.

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.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909 , 2018.

Dimitri P Bertsekas. Nonlinear programming. Journal of the Operational Research Society , 48(3):334-334, 1997.

Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems , 32, 2019.

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.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density estimation. Advances in neural information processing systems , 30, 2017.

Irina Higgins, Loic Matthey, 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.

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 Proofs

A.1 Proof of Proposition 3.1 and 3.3

Convergence of the Single Expert Objective (Proposition 3.1). We perform coordinate ascent on ˜ 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 ( o , a ) that maximizes ˜ J [55] (see Section 3.1) and hence we have ˜ J ( p ( o , a ) ( i ) , θ ) ≥ ˜ J ( p ( o , a ) ( i -1) , θ ) where i denotes the iteration. Under suitable assumptions ( log p θ is differentiable, its gradient is L -Lipschitz, θ z is updated using gradient ascent and the learning rate is chosen such that the descent lemma [56] holds), it holds that ˜ J ( p ( o , a ) , θ ( i ) ) ≥ ˜ J ( p ( o , a ) , θ ( i -1) ) . Hence, we are guaranteed to converge to a stationary point of ˜ J .

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 ( ψ ( i ) , q ) ≥ L ( ψ ( i -1) , q ) to prove that J ( ψ ) is guaranteed to converge to a stationary point, with ψ = { p ( z ) , { p ( o , a | z ) } z , { θ 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 ( o , 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 { θ z } z under suitable assumptions on log p θ 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.

0

20

40

60

80

100

10

10

.

5

11

11

.

5

i

J

(

ψ

)

A.2 Proof of Proposition 3.2

Expanding the entropy in Equation 5 we obtain

Using p ∗ ( o , a | z ) = ˜ p ( o , a | z ) / ∑ n ˜ p ( o n , a n | z ) yields

Next, leveraging that log ˜ p ( o , a | z ) = R z ( o n , a n ) /η we see that

which concludes the proof.

A.3 Proof of Corollary 3.2.1

We start by rewriting the lower bound as L ( ψ , q ) =

Using p ∗ ( o , a | z ) ∝ ˜ p ( o , a | z ) and Proposition 3.2 we obtain

With η log ˜ p ( o , a | z ) = R z ( o n , a n ) all most terms cancel, giving

L ( ψ , q ) = E p ∗ ( z ) [ η log ∑ z ∑ n ˜ p ( o n , a n | z ) ] = η log ∑ z ∑ n ˜ p ( o n , a n | z ) , which concludes the proof.

A.4 Proof of Corollary 3.2.2

Expanding the expected KL divergence, we get

Noting that p ( z | o n , a n ) is independent of ϕ we can rewrite the objective as

Using that p ( z | o , a ) = ˜ p ( o , a | z ) / ∑ z ˜ p ( o , a | z ) together with p ( o , a ) = ∑ z p ∗ ( z ) p ∗ ( o , a | z ) yields

Using Proposition 3.2 we can rewrite p ∗ ( z ) p ∗ ( o , a | z ) as ˜ p ( o , a | z ) / ∑ z ∑ n ˜ p ( o n , a n | z ) . Since the constant factor 1 / ∑ z ∑ n ˜ p ( o n , a n | z ) does not affect the optimal value of ϕ we obtain

B Derivations

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.,

Next, we use Bayes' theorem, that is, p ( o n , a n ) = p ( z ) p ( o n , a n | z ) /p ( z | o n , a n ) , giving

Moreover, we add and subtract the log auxiliary distribution log q ( z | o n , a n ) which yields

Rearranging the terms leads and writing the sums in terms of expectations we arrive at

Lastly, multiplying H ( o , a ) with η and adding E p ( z ) E p ( o , a | z ) [log p θ z ( y | 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

Following the procedure of constrained optimization, we write down the Lagrangian function [9] as

where λ is the Lagrangian multiplier. As p is discrete, we solve for the optimal entries of p ( o n , a n | z ) , that is, p ′ ( o n , a n , λ | z ) = arg max p L ( p, λ ) . Setting the partial derivative of L ( p, λ ) with respect to p zero, i.e.,

Plugging p ′ back in the Lagrangian gives the dual function g ( λ ) , that is, yields p ′ ( o n , a n , λ | z ) = exp ( R z ( o n , a n ) -η + λ ) /η.

Solving for λ ∗ = arg min λ ≥ 0 g ( λ ) equates to

Finally, substituting λ ∗ into p ′ we have

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 θ z J z ( p ( o , a | z ) , θ z ) =

Noting that q ( z | o n , a n ) and p ( o n , a n | z ) are independent of θ z and p ( o n , a n | z ) = ˜ p ( o n , a n | z ) / ∑ n ˜ p ( o n , a n | z ) we find that

Noting that ∑ n ˜ p ( o n , a n | z ) is a constant scaling factor concludes the derivation.

C Experiment Setup

C.1 Environments and Datasets

C.1.1 Obstacle Avoidance

Dataset. The obstacle avoidance dataset contains 96 trajectories resulting in a total of 7 . 3 k ( o , a ) pairs. The observations o ∈ R 4 contain the end-effector position and velocity in Cartesian space. Please note that the height of the robot is fixed. The actions a ∈ 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

is computed for successful trajectories, where each behavior β 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 ( β ) which is used to compute the entropy. By the use of log 24 we make sure that H 24 ( τ ) ∈ [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, H ( β ) = 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 k ( o , a ) pairs. The observations o ∈ 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 a ∈ 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 β (see Figure 2) we use the expected entropy

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 H ( β ) ∈ [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. Furthermore, we set p ( o 0 ) = 1 / 30 as we sample 30 block configurations uniformly from a configuration space (see Figure 7). For each 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 ( β | o 0 ) . Once repeated for all 30 configurations , we compute E p ( o 0 ) [ H ( β | o 0 ) ] .

C.1.3 Table Tennis

Dataset. The table tennis dataset contains 5000 ( o , a ) pairs. The observations o ∈ 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 a ∈ 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 |B| = |D| , number of components N z = 50 and expert learning rate equal to 5 × 10 -4 . Furthermore, we initialized all curriculum weights as p ( o n , 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 sigle-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 η we performed a hyperparameter sweep using Bayesian optimization. The respective values are η = 1 / 30 for obstacle avoidance, η = 2 for block pushing and Franka kitchen and η = 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.

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 | 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).

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

architecture. Moreover, we consider an additional scaling factor ( β ) for the KL regularization in the lower bound objective of the VAE as suggested in [61].

Implicit Behavior Cloning (IBC). IBC was proposed in [17] and uses energy-based models to learn a joint distribution over inputs and targets. Following common practice we parameterize the model as neural network. Moreover, we use the version that adds a gradient penalty to the InfoNCE loss [17]. For sampling, we use gradient-based Langevin MCMC [57]. Despite our effort, we could not achieve good results with IBC. A list of hyperparameters is shown in Table 7.

Behavior Transformer (BET). Behavior transformers were recently proposed in [22]. The model employs a minGPT transformer [62] to predict targets by decomposing them into cluster centers and residual offsets. To obtain a fair comparison, we compare our method to the version with no history. A comprehensive list of hyperparameters is shown in Table 8.

D Extended Related Work

D.1 Connection to ML-Cur

This section elaborates on the differences between this work and the work by [30], as both studies leverage curriculum learning paired with Mixture of Experts to facilitate the acquisition of diverse skills. However, it's essential to discern several noteworthy distinctions between our approach and theirs:

Firstly, Li et al. primarily focus on linear policies and a gating distribution with limited flexibility, constraining the expressiveness of their learned policies. In contrast, our approach allows for the use of arbitrarily non-linear neural network parameterizations for both, expert policies and gating.

Another divergence lies in the handling of mini-batches during training. While our algorithm accommodates the use of mini-batches, Li et al.'s method does not support this feature. The ability to work with mini-batches can significantly enhance the efficiency and scalability of the learning process, especially when dealing with extensive datasets or intricate environments.

Additionally, Li et al.'s evaluation primarily centers around relatively simple tasks, that do not require complex manipulations, indicating potential limitations in terms of task complexity and applicability. In contrast, our work is designed to address more intricate and challenging environments, expanding the range of potential applications and domains.

Lastly, it's worth noting that Li et al. specifically focus on the learning of skills parameterized by motion primitives [31]. In contrast, our framework offers the versatility to encompass various skill types and parameterizations, broadening the scope of potential applications and use cases.

D.2 Connection to Expectation Maximization

In this section we want to highlight the commonalities and differences between our algorithm and the expectation-maximization (EM) algorithm for mixtures of experts. First, we look at the updates of the variational distribution q . Next, we compare the expert optimization. Lastly, we take a closer look at the optimization of the gating distribution.

The EM algorithm sets the variational distribution during the E-step to

for all samples n and components z . In the M-step, the gating distribution p ( z | o ) is updated such that the KL divergence between q ( z | o ) and p ( z | o ) is minimized. Using the properties of the KL divergence, we obtain a global optimum by setting p ( z | o n ) = q ( z | o n ) for all n and all z . This allows us to rewrite Equation 7 using the recursion in q , giving

where ( i ) denotes the iteration of the EM algorithm. The update for the variational distribution of the IMC algorithm is given by

Consequently, we see that q ( z | o ) = q ( z | oa ) for η = 1 . However, the two algorithms mainly differ in the M-step for the experts: The EM algorithm uses the variational distribution to assign weights to samples, i.e.

whereas IMC uses the curricula as weights, that is,

This subtle difference shows the properties of moment and information projection: In the EM algorithm each sample o n contributes to the expert optimization as ∑ z q ( z | o n ) = 1 . However, if all curricula ignore the n th sample, it will not have impact on the expert optimization. Assuming that the curricula ignore samples that the corresponding experts are not able to represent, IMC prevents experts from having to average over 'too hard' samples. Furthermore, this results in reduced outlier sensitivity as they are likely to be ignored for expert optimization. Lastly, we highlight the difference between the gating optimization: Assuming that both algorithms train a gating network g ϕ ( z | o ) we have for the EM algorithm and

for IMC. Similar to the expert optimization, EM includes all samples to fit the parameters of the gating network, whereas IMC ignores samples where the unnormalized curriculum weights ˜ p ( o n , a n | z ) are zero for all components.

E Algorithm Details & Ablation Studies

E.1 Inference Details

We provide pseudocode to further clarify the inference procedure of our proposed method (see Algorithm 2).

E.2 Expert Design Choices

Distribution. In our mixture of experts policy, we employ Gaussian distributions with a fixed variance to represent the individual experts. This choice offers several benefits in terms of likelihood calculation, optimization and ease of sampling:

To perform the M-Step for the curricula (Section 3.3), exact log-likelihood computation is necessary. This computation becomes straightforward when using Gaussian distributions. Additionally, when Gaussian distributions with fixed variances are employed to represent the experts, the M-Step for the experts simplifies to a weighted squared-error minimization. Specifically, maximizing the weighted likelihood reduces to minimizing the weighted squared error between the predicted actions and the actual actions.

The optimization problem for the expert update can be formulated as follows:

This optimization problem can be efficiently solved using gradient-based methods. Lastly, sampling from Gaussian distributions is well-known to be straightforward and efficient.

Parameterization. We experimented with two different parameterizations of the Gaussian expert means µ θ z , which we dub single-head and multi-head : For single-head, there is no parameter sharing between the different experts. Each expert has its own set of parameters θ z . As a result, we learn N z different multi-layer perceptrons (MLPs) µ θ z : R |O| → R |A| , where N z is the number of mixture components. In contrast, the multi-head parameterization uses a global set of parameters θ for all experts and hence allows for feature sharing. We thus learn a single MLP µ θ : R |O| → R N z ×|A| .

To compare both parameterizations, we conducted an ablation study where we evaluate the MoE policy on obstacle avoidance, table tennis and Franka kitchen. In order to have a similar number of parameters, we used smaller MLPs for single-head, that is, 1 -4 layers whereas for multi-head we used a 6 layer MLP. The results are shown in Table 9 and are generated using 30 components for the obstacle avoidance and table tennis task. The remaining hyperparameters are equal to the ones listed in the main manuscript. For Franka kitchen, we report the cumulative success rate and entropy for a different number of completed tasks. We report the mean and standard deviation calculated across 10 different seeds. Our findings indicate that, in the majority of experiments, the single-head parameterization outperforms the mutli-head alternative. Notably, we observed a substantial performance disparity, especially in the case of Franka kitchen.

Expert Complexity. We conducted an ablation study to evaluate the effect of expert complexity on the performance of the IMC algorithm. The study involved varying the number of hidden layers in the single-head expert architecture while assessing the IMC algorithm's performance on the Franka kitchen task using the cumulative success rate and entropy. The results, presented in Figure 9, were obtained using the hyperparameters specified in the main manuscript. Mean and standard deviation were calculated across 5 different seeds. Our findings demonstrate a positive correlation between expert complexity and achieved performance.

0

2

4

6

8

10

0

2

4

# hidden layer

success rate

0

2

4

6

8

10

0

2

4

6

8

# hidden layer

entropy

E.3 Component Utilization

We further analyze how IMC utilizes its individual components. Specifically, we assess the entropy of the weights assigned to these components, denoted as H ( z ) and defined as H ( z ) = -∑ z p ( z ) log p ( z ) . Maximum entropy occurs when the model allocates equal importance to all components to solve a task, which implies that p ( z ) = 1 /N z . Conversely, if the model relies solely on a single component, the entropy H ( z ) equals zero. We conduct the study on the Franka kitchen task, evaluating it through cumulative success rates and entropy. The results are shown in Figure 10a. We computed the mean and standard deviation based on data collected from 5 different seeds. Our findings reveal that IMC exhibits substantial component utilization even when dealing with a large number of components, denoted as N z .

E.4 Number of Components Sensitivity

To investigate how the algorithm responds to changes in the number of components N z , we performed a comprehensive ablation study. Figure 10b shows the success rate of IMC on the Franka kitchen task using a varying number of components. Our findings indicate that the algorithm's performance remains stable across different values of N z . This robust behavior signifies that the algorithm can adeptly accommodate varying numbers of components without experiencing substantial performance fluctuations. Hence, opting for a higher number of components primarily entails increased computational costs, without leading to overfitting issues.

50

100

150

200

0

2

4

N

z

H

(

z

)

IMC

1

/N

z

(a)

0

50

100

150

200

2

.

5

3

3

.

5

4

N

z

success rate

(b)

E.5 Curriculum Pacing Sensitivity

To examine the algorithm's sensitivity to the curriculum pacing parameter η , we conducted an ablation study. Figure 11 presents the results obtained using 30 components for the obstacle avoidance and table tennis tasks, while maintaining the remaining hyperparameters as listed in the main manuscript. For the Franka kitchen task, we analyzed the cumulative success rate and entropy across varying numbers of completed tasks. The mean and standard deviation were calculated across 5 different seeds. Our findings reveal that the optimal value for η is dependent on the specific task. Nevertheless, the algorithm exhibits stable performance even when η values differ by an order of magnitude.

0

.

001

0

.

01

0

.

1

1

10

100

0

0

.

2

0

.

4

0

.

6

0

.

8

1

η

success rate

0

.

001

0

.

01

0

.

1

1

10

100

0

0

.

2

0

.

4

0

.

6

0

.

8

1

η

entropy

(a) Obstacle avoidance

0

.

001

0

.

01

0

.

1

1

10

100

0

0

.

2

0

.

4

0

.

6

0

.

8

1

η

success rate

0

.

001

0

.

01

0

.

1

1

10

100

0

0

.

2

0

.

4

0

.

6

0

.

8

1

η

distance error

(b) Table tennis

0

.

001

0

.

01

0

.

1

1

10

100

0

1

2

3

4

η

success rate

0

.

001

0

.

01

0

.

1

1

10

100

0

2

4

6

8

η

entropy

(c) Franka kitchen

Peer reviews:
comments: SUMMARY:
This paper proposes an imitation learning method IMC that can model multi-modal behaviors. IMC avoids the mode-averaging issue with an objective similar to reverse-KL. To cover all modes in the dataset, IMC further introduces a mixture model with multiple components, each focusing on different data distribution it specializes to. The mixture model is optimized with the EM algorithm. The authors provide extensive experiments over simulated control environments and demonstrate IMC's superior modeling abilities.

STRENGTHS:
1. IMC is a well-motivated method for multi-modal density estimation and provides an elegant reverse KL-based solution.
2. The experiments in low-dimensional control environments are extensive. IMC is compared against major generative models (with maximum likelihood objective) and addresses the mode-averaging issue better.
3. The visualization of learned trajectories clearly demonstrates the learned modes of different methods and the natural mode-ignorance ability of IMC.

WEAKNESSES:
The presentation in the experiment section could be improved. Currently, the subsections for different environments repeat mostly the same conclusions that IMC achieves high success rates and covers diverse modes. I think it is better to emphasize anything special for different environments or different baselines.

QUESTIONS:
1. How sensitive is the performance of IMC with respect to different choices of $\eta$?
2. Can IMC be scaled up to scenarios with high-dimensional input?

LIMITATIONS:
The authors have discussed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study a method, Information Maximizing Curriculum (IMC), that performs behavioral cloning by having the model selectively choose a learned weighing of the demonstration data for which the model is best at predicting (via minimizing the reverse KL divergence). To avoid the mode-seeking behavior of this approach, the authors extend the method to make use of K such model components, leading to a mixture of experts (MoE) approach, whereby each component selectively models the distribution in this way, while maximizing their joint information projection over the dataset (via simultaneously maximizing the entropy of the MoE distribution). Experiments in two robotics  simulators show it outperforms several other baselines based on generative models and basic MoE methods trained via expectation-maximization and backpropagation.\n\nSTRENGTHS:\n- The paper is overall well written. The method and its motivations are clearly communicated.\n- The authors compare to several important baselines, spanning generative models, energy-based models, and MoE methods, and show strong performance improvements.\n\nWEAKNESSES:\nThe main weakness I see is that this work seems almost identical to Li, et al (2023), which proposes largely the same method. Moreover, this prior work looks at a very similar set of environments as this paper. While the authors cite this prior work, they do not directly compare to it. Given the extremely close similarity between the two methods, it seems important to compare to this work to justify the contribution in this paper, which in a sense, is an extension of the method in Li, et al (2023) by incorporating neural networks.\n\nI see two main ways to show improvement over this prior work: The authors can either (1) show their method outperforms the approach in Li, et al (2023) on the tasks studied, or (2) Demonstrate clearly how IMC can scale to environments in which the method of Li, et al (2023) cannot, thereby clearly justifying the strengths of this extension. I imagine neither of these aims is too difficult, but such a comparison seems sorely missing, given the strong similarity between the two works.\n\nLastly, I find it odd that the authors couch their method as “curriculum learning,” ignore the field of active learning altogether, and proceed to imply that their method is novel in not requiring an a priori difficulty metric for each datapoint. The authors should relate their method to active learning, which is a mature field of study, and most active learning methods can be viewed as “curriculum learning” without any a priori notion of task difficulty.\n\n### References\n\nMaximilian Xiling Li, Onur Celik, Philipp Becker, Denis Blessing, Rudolf Lioutikov, and Gerhard Neumann. Curriculum-based imitation of versatile skills. arXiv preprint arXiv:2304.05171,\n2023.\n\nQUESTIONS:\n- This study primarily focuses on continuous control robotics simulators. Do you have a sense of whether IMC can be made to work in discrete control settings, e.g. Atari, chip design, web navigation.\n- Are there cases where maximizing mode coverage is disadvantageous? Often demonstration datasets contain a mix of approximately optimal and very suboptimal trajectories. Could IMC lead to learning MoE that perform worse in these settings?\n\nLIMITATIONS:\nThis work only focuses on fairly simple continuous control tasks. Scaling this to more complex imitation learning settings (e.g. from pixels or controlling a much more complex action space) as well as discrete control settings would add confidence in the utility of this approach. Relatedly, demonstrating success on these more challenging settings would be a great way to highlight why IMC is an important improvement over the Li, et al (2023) work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a curriculum based approach for imitation learning. Overall, the imitation learning problem is posed as a conditional density estimation problem. Given the multi-modal nature of underlying data, this paper proposes to learn a curriculum based mixture of expert policy. Intuitively, each expert is only responsible for learning a subset of the training data and learns the policy for this subset using reverse KL. Experimentally, the proposed approach is verified in 4 different environments and compared against both common and recently proposed approaches.\n\nSTRENGTHS:\nThe paper presents an interesting and grounded approach for learning from multimodal data distribution. The overall idea of using a curriculum to weight data samples based on how well they are represented under the expert policy seems to be generally useful. The approach is shown to work in diverse settings and seems competitive with recent results.\n\nWEAKNESSES:\n*Baseline results:* Overall, the proposed approach provides very little significant advantage over simpler baselines. For instance, in Figure 4 and Figure 5 we can see that for most tasks (Pusing, Tennis, Kitchen), a diffusion model based approach (DDPM) is highly competitive to the proposed approach (success rate difference is less than 0.05) across all tasks. Only in the Obstacle avoidance task, does DDPM approach suffer although CVAE still performs quite well (unclear why DDPM performs poorly here). \n\n*Benchmark tasks:* While the paper tries to evaluate the proposed approach on multiple datasets, most of the tasks/datasets are not commonly used across multimodal tasks (only Franka-Kitchen is the commonly used dataset). Given the large set of recent works in this area it would be much better to evaluate on tasks which other recent works evaluate. For instance, most recent works evaluate on RoboMimic dataset [1] and the block pushing task from IBC [2]. Both of these task suites have human demonstrations available and since the underlying data distribution is highly important for these proposed approaches, reusing datasets from previously proposed approaches will allow for a much fairer comparison.\n\n[1] Robomimic: https://github.com/ARISE-Initiative/robomimic\n[2] Florence et al. Implicit Behavior Cloning\n\n--- Post Rebuttal ---\n\nI don't see any issues with the paper the provided code also looks reasonable and reproduces the results so I would update my recommendation to Weak Accept.\n\nQUESTIONS:\nOverall the approach is broadly similar to an EM approach with mixture of policies. As noted, the main difference being that the proposed approach uses a curriculum weighted optimization objective to assign weights to samples instead of directly using the variational distribution. While algorithmically, this difference is not large in terms of empirical results the EM based approaches lags behind the proposed approach quite a lot (especially for the obstacle avoidance task). Do the authors have an intuition why the EM approach suffers so much in comparison?\n\nLIMITATIONS:\nLimitations are discussed", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a learning protocol that learns multiple policies (\"skills\"), distribution over skills, and a per-skill priority over experience buffer. This is achieved by maximizing a variational lower bound of a certain averaged regularized KL distance. Namely, the objective is a sum of two terms. The first term is a KL distance between the per-skill priority over previously seen actions and a policy, both conditioned on the skill and previously seen observation, and averaged over them. A second term is an entropy of the joint distribution over experience and skills, which acts as a regularizer, forcing it to be as diverse as possible.\n\nIn the experiment section, the skill policies share a common backbone and output a mean of an isotropic Gaussian distribution. The paper tests the method over several tasks, such as Obstacle Avoidance, Block Pushing, Franka Kitchen, or Table Tennis. The approach compares (mostly) favorably against seven baseline methods.\n\nSTRENGTHS:\nThe paper proposes a simple objective to train multiple policies at once, such that together they cover multiple modes in the dataset. The results show that the method (mostly) performs better than the considered baselines.\n\nWEAKNESSES:\nThe empirical part could be improved. Namely:\n* Performance metrics should be defined in the main body of the paper:\n\t* There could be a dedicated section for that purpose. This would improve the exposition and allow more discussion about the method in Sections 5.1-5.4.\n\t* This could also make reading the results (e.g., Table 1) more accessible. For instance, the definition of entropy varies across environments, as shown in Appendix C, which makes the numbers incomparable and possibly draw wrong conclusions.\n\t* The discussion on the range of entropy values should be moved from the Appendix to the main body. It provides grounding for the numbers in Table 1.\n\t* Additional place for this and a more in-depth discussion of the results could be achieved by moving Section 2.1 and Figure 1 to the Appendix (it is a well know property of KL). Similarly, some parts of Section 3 could be moved to the Appendix.\n* The paper does not provide numbers on how the algorithm mixes between skills, e.g., the entropy of the distribution over components ($p(z)$).\n* It seems that each environment required a different setup of the method. This is a limiting factor for the method. The paper, however, does not guide the potential user of the method in choosing the relevant hyperparameters.  For example, estimating how many modes are in the data can be a non-trivial task, and consequently, setting the number of components ($N_z$) selected for each experiment may be non-trivial. The paper has no information regarding the actual values of $N_z$ chosen for each experiment.\n* Descriptions in Sections 5.1-5.4 are mainly technical and deal with the setup. The discussion about the results of IMC is limited to one or two sentences that do not add more information than Table 1 or Figure 6. A reader would expect that most of such a section would provide real insights about the method (e.g., what skills were learned, how they were acquired during training, where all modes were discovered, how a distribution that mixes skills looks, etc.).\n* Figure 6: Shouldn't we expect that the performance is an increasing function of the number of components (the more modes covered, the higher the objective)?\n* There is no mention of what values $\\eta$ were chosen.\n* Other issues:\n\t* Table 1: The description is not self-explanatory. What is the selected number of components? What versions of entropy are used? What is the setup of experiments? Etc.\n\t* Sections 5.1, 5.2, and 5.4 refer to Table 13, which is not present in the paper. Most likely, it should refer to Table 1.\n\t* Figure 6 precedes Figure 5.\n\nThe quality of the technical part of the paper could be improved and more clearly agitated. In particular, \n* Section 2.2: \n\t* There is no definition of $\\mathcal O$, $\\mathcal A$, $p$, $z$.\n\t* Is $z$ a continuous or discrete random variable? From the context of the following sections, it would seem that the latter.\n* Section 3.1:\n\t* It seems more clear to write $p(o, a)$ instead of $p(\\mathcal D)$.\n\t* In equation (2), there should be $\\mathcal H(p(\\mathcal D))$ (or better $\\mathcal H(p(o,a))$, see the item above). Additionally $\\mathbb E_{p(\\mathcal D)}$ could be more transparent if written as $\\mathbb E_{o,a\\sim p(\\mathcal D)}$.\n\t* The formulas in lines 80, 85, and 86 should be better justified. There is only a cryptic comment in line 81 about optimization in an alternating fashion, which is not justified, nor a suitable reference is given.\n\t* The description in lines 87-93 is relatively informal and does not refer to the formulas in lines 80, 85, and 86. Additionally, the text uses colloquial terms such as \"representational capacity of the policy\", \"capacity [..] is exhausted\", or alludes to the convergence of curriculum, which was not proved.\n\t* There is no definition of $\\mathcal D_n$ (lines 85, 86, 119, 120, 124, 129, 130, etc.). Why not just write $(o_n, a_n)$?\n\t* In the proof of Proposition 3.1 (Appendix A.1.), the formulas read $p^*(o)$ where the Proposition refers to $p^*(z)$. A similar comment refers to other proofs in Appendix A.\n* Section 3.2\n\t* There is no description of the objective $J(\\psi)$ and its lower bound, showing how the individual terms promote desirable behavior. \n\t* In equation (3), it seems that the entropy term should be equal to $\\mathcal H(p(o, a, z))$; otherwise, equation (4) seems not to be valid.\n\t* It seems that in equation (4) and in the definition of $R_z$, there should be $q(z|o,a)$ (in place of $q(z|o)$ and $q(z|\\mathcal D)$, respectively). Additionally, it would make sense to define $R_z$ as a function of $(o, a)$ (instead of $\\mathcal D$).\n\t* Similar comments apply to Sections 3.3-3.5.\n\nEdit: After the Authors' rebuttal, I have raised the score (4->6).\n\nQUESTIONS:\nConsult the Weaknesses section.\n\nLIMITATIONS:\nThe paper includes a brief limitations section. What could also be mentioned is that the method requires setting several important hyperparameters (see the review), some of which require non-trivial knowledge about data. The method also was tested on continuous tasks, so a natural question (limitation or future research) would be to ask about performance on discrete domains, such as combinatorial puzzles like chess or video games like Atari benchmark.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}