Original: Inferring Versatile Plans from Demonstrations by Matching Geometric Features
Revised: Inferring Versatile Behavior from Demonstrations by Matching Geometric Descriptors

Original: Anonymous Author(s)

Original: AffiliationAddress email Abstract: Humans intuitively solve tasks in versatile ways, varying their behaviorin terms of trajectory-based planning and for individual steps.
Revised: Abstract: Humans intuitively solve tasks in versatile ways, varying their behavior in terms of trajectory-based planning and for individual steps.

Original: Thus, they caneasily generalize and adapt to new and changing environments.
Revised: Thus, they can easily generalize and adapt to new and changing environments.

Original: Current Imitation
Revised: Current Imitation

Original: Learning algorithms often only consider unimodal expert demonstrations andact in a state-action-based setting, making it difficult for them to imitate humanbehavior in case of versatile demonstrations.
Revised: Learning algorithms often only consider unimodal expert demonstrations and act in a state-action-based setting, making it difﬁcult for them to imitate human behavior in case of versatile demonstrations.

Original: Instead, we combine a mixtureof movement primitives with a distribution matching objective to learn versatileplans that match the expert’s behavior and versatility.
Revised: Instead, we combine a mixture of movement primitives with a distribution matching objective to learn versatile behaviors that match the expert’s behavior and versatility.

Original: To facilitate generalizationto novel task configurations, we do not directly match the agent’s and expert’strajectory distributions but rather work with generic geometric representationswhich generalize well to unseen task configurations.
Revised: To facilitate generalization to novel task conﬁgurations, we do not directly match the agent’s and expert’s trajectory distributions but rather work with concise geometric descriptors which generalize well to unseen task conﬁgurations.

Original: We empirically validate ourmethod on various robot tasks using versatile human demonstrations and compareto imitation learning algorithms in a state-action setting as well as a trajectory-basedsetting.
Revised: We empirically validate our method on various robot tasks using versatile human demonstrations and compare to imitation learning algorithms in a state-action setting as well as a trajectory-based setting.

Original: We find that the geometric representation greatly helps in generalizing tonew configurations and that combining this representation with our distributionmatching objective is crucial for representing and reproducing versatile behavior.
Revised: We ﬁnd that the geometric descriptors greatly help in generalizing to new task conﬁgurations and that combining them with our distribution-matching objective is crucial for representing and reproducing versatile behavior.

Original:  Introduction
Revised:  Introduction

Original: Imitation Learning (IL)
Revised: Imitation Learning (IL)

Original: [1, 2, 3] from human demonstrations is challenging as humans often solvetasks in versatile ways.
Revised: [1, 2, 3] from human demonstrations is challenging as humans often solve tasks in versatile ways.

Original: Even the same person might solve a task differently when confronted withit multiple times.
Revised: Even the same person might solve a task differently when confronted with it multiple times.

Original: Behavior can be versatile in terms of task-level decisions, such as planning aroute to a target, and individual actions, such as randomly pausing during a movement to thinkabout the next steps.
Revised: Behavior can be versatile in terms of task-level decisions, such as planning a route to a target and individual actions, such as randomly pausing during a movement to think about the next steps.

Original: Most recent IL approaches [4, 5, 6, 7] model the behavior in state-actionspace using Gaussian policies, which assumes the behavior is unimodal and cannot capture thisplanning versatility.
Revised: Most recent IL approaches [4, 5, 6, 7] model the behavior in state-action space using Gaussian policies, which assumes the behavior is unimodal and cannot capture this planning versatility.

Original: Yet, this assumption is violated by most human expert datasets,causing theseapproaches to often generalize poorly to human demonstrations [8].
Revised: Yet, this assumption is violated by most human expert datasets, often causing poor generalization to human demonstrations [8].

Original: Additionally,while humansexcel at transferring learned behavior to previously unseen tasks or contexts, IL approaches oftenneed immense amounts of data to achieve generalization [9, 8].Such data-efficient generalizationis crucial for applications that involve dynamically changing or unstructured environments.
Revised: Additionally, IL approaches often need immense amounts of data to achieve generalization [9, 8].

Original: Wetackle these challenges by a feature matching approach to IL, matching context-insensitive geometricrepresentations to facilitate generalization to novel contexts.
Revised: We tackle these challenges by a feature matching approach to IL that is able to generalize to novel contexts from a small amount of expert demonstrations.

Original: Here, we use the term context asa synonym for ’task descriptor’, e.g., as  goal locations to reach.More precisely, we match thefeatures of planned trajectories with features of the long-term behavior of the expert.
Revised: Such contexts are (typically low-dimensional) vectors that describe a speciﬁc task conﬁguration within a family of related tasks, such as e.g., the coordinates of goal locations to reach.

Original: The resultingapproach, Versatile Imitation from Geometrically Observed Representations (VIGOR), modelsplanned trajectories  using Gaussian Mixture Models (GMMs) of Probabilistic Motion Primitives(ProMPs) [10] which capture the versatility in the demonstrated behavior. We enable generalizationto new contexts by using geometric descriptors instead of directly matching the agent’s and expert’strajectory distributions.
Revised: The resulting approach, Versatile Imitation from Geometrically Observed Representations (VIGOR), models distributions that match expert trajectories in terms of concise geometric behavioral descriptors. VIGOR represents distributions using Gaussian Mixture Models (GMMs) over Probabilistic Motion Primitives (ProMPs) [10], and utilizes descriptors that are designed to abstract away from concrete contexts and thus facilitate generalization to novel contexts in a sample-efﬁcient way.

Original: An example of such a descriptor would be using the distance to a target
Revised: An example of such descriptors would be geometric features in the form of (only) the distance between a target

Original: + ⇒ 1.
Revised: + ⇒ 1.

Original: TrainContext 2.
Revised: TrainContext 2.

Original: [Figure]
Revised: [Figure]

Original: GeneralizationtoNovelTestContext
Revised: GeneralizationtoNovelTestContext

Original: [Figure]
Revised: [Figure]

Original: D i s t a n ce t o T a r g e t
Revised: D i s t a n ce t o T a r g e t

Original: [Figure]
Revised: [Figure]

Original: instead of its absolute position.
Revised: and the robot’s end-effector, plus features that denote the smoothness of the robot. While similar behavioral descriptors have found use in Inverse Optimal Control [11, 12] to recover cost functions from demonstrations, common IL approaches instead use state representations that include both relative and absolute information [5, 13].

Original: The exact form of this representation is task-dependent and allowsthe user to include domain knowledge.
Revised: The exact form of this set of behavioral descriptors is task-dependent and allows the user to include domain knowledge.

Original: Figure 1 shows how this works on an exemplary reachingtask.
Revised: Figure 1 shows how VIGOR works on an exemplary reaching task.

Original: For training the GMMs, we rely on recent approaches [11, 12, 13] for distribution matchingand variational inference with GMMs.
Revised: For training the GMMs, we rely on recent approaches [14, 15, 16] for distribution matching and Variational Inference.

Original: For inference, we train individual GMMs for each trainingcontext and an arbitrary number ofunseen test contexts.
Revised: For inference, we train individual GMMs for each training context and an arbitrary number of test contexts.

Original: VIGOR generalizes well already from asmall number of demonstrated configurations and is capable of reproducing the versatility inherent tothe expert demonstrations.

Original: We evaluate VIGOR on a suite of versatile robot coordination tasks, where demonstrations arecollected from human experts through teleoperation.
Revised: We evaluate VIGOR on a suite of versatile robotic tasks, where demonstrations are collected from human experts through teleoperation.

Original: We compare VIGOR to standard Behavioral Cloning [14], Generative Adversarial Imitation Learning (GAIL) [5] and Inverse Reinforcement Learning (IRL) techniques [7].
Revised: We compare our method to standard Behavioral Cloning (BC) [17], BC with a GMM policy [13, 18], Generative Adversarial Imitation Learning (GAIL) [5] and Inverse Reinforcement Learning (IRL) techniques [7].

Original: We find that these methods either can not extract goal-directedbehavior from the versatile demonstrations as they try to  average over multiple solutions or that theycan not capture the full versatility of the demonstrations.
Revised: We ﬁnd that these methods either fail to learn useful behavior because they average over multiple solutions, or that they are unable to capture the full versatility of the demonstrations.

Original: Contrary to this, VIGORextracts behaviorof the highest performance that is also highly versatile.
Revised: Contrary to this, VIGOR accurately models highly versatile behavior from already a small number of trajectories demonstrated by human experts.

Revised: We perform extensive ablation studies to showcase the importance of different parameters and design choices.

Original: Expert demonstrations and code will beprovided for the camera-ready version.
Revised: Datasets and code can be found at https://www.github.com/NiklasFreymuth/VIGOR . To summarize, our list of contributions is as follows: • We infer multi-modal distributions of desired trajectories from highly versatile human expert trajectories by matching distributions over behavioral descriptors between learner and expert. • Our approach is, to the best of our knowledge, the ﬁrst Adversarial Imitation Learning method to utilize concise behavioral descriptors to facilitate generalization to novel contexts from already a small number of demonstrations.

Revised: • We conduct experiments with human demonstrations in simulation and on a real robot and ﬁnd that VIGOR accurately models highly versatile human behavior, outperforming various Imitation Learning baselines.

Original:  Related Work
Revised:  Related Work

Original: Imitation Learning for Skills.
Revised: Imitation Learning for Skills.

Original: A common way to represent skills over trajectories is via Movement
Revised: A common way to represent skills over trajectories is via Movement

Original: Primitives (MPs) [15, 10].
Revised: Primitives (MPs) [19, 10].

Original: While learning individual MPs from human demonstrations is a simpleregression problem, modeling versatile behavior requires a more sophisticated model that canrepresent such multi-modality, e.g., a mixture model.
Revised: While learning individual MPs from human demonstrations is a simple regression problem, modeling versatile behavior requires a more sophisticated model that can represent such multi-modality, e.g., a mixture model.

Original: Using Expectation Maximization (EM) [16],both [17] and [18] fit a single GMM over  multiple demonstrations for different contexts and use themto generalize to novel contexts by conditioning.
Revised: Using Expectation Maximization (EM) [20], both [21] and [22] ﬁt a single GMM over MP parameters of multiple demonstrations for different contexts and use them to generalize to novel contexts by conditioning.

Original: Yet, as they only fit a single GMM for all contexts,they do not explicitly focus on representing versatility.
Revised: Yet, as they only ﬁt a single GMM for all contexts, they do not explicitly focus on representing versatility.

Original: To alleviate this problem, recent work [19]instead learns a mixture over contextualized MPs using Gaussian Mixture Regression.
Revised: Recent work [23] instead learns a mixture over contextualized MPs using Gaussian Mixture Regression.

Original: This has beenextended to non-linear relations between MP parameters and context [20] by using Mixture Density
Revised: This has been extended to non-linear relations between MP parameters and context [18] by using Mixture Density

Original: Networks (MDNs) [21].
Revised: Networks (MDNs) [24].

Original: While the above approaches work for tasks withfew modes, they tend toscale poorly to versatile demonstrations with a high number of modes. Moreover, optimizingan MDNs can be challenging in the case of many modes in the demonstrations.
Revised: While the above approaches work for tasks with a small number of modes, optimizing MDNs can be challenging in the case of many modes in the demonstrations.

Original: Osa et al.
Revised: Osa et al.

Original: [22] useplanning algorithms to reproduce human demonstrations in a trajectory-based setting.
Revised: [25] use planning algorithms to reproduce human demonstrations in a trajectory-based setting.

Original: They tune theparameters of a cost function such that the planned trajectories match the demonstrations and use theextracted cost function to guide the trajectory optimization process.
Revised: They tune the parameters of a cost function such that the planned trajectories match the demonstrations and use the extracted cost function to guide the trajectory optimization process.

Original: Yet, this approach is limited touni-modal demonstrations and requires engineered cost functions.
Revised: Yet, this approach is limited to uni-modal demonstrations and requires engineered cost functions.

Original: Imitation Learning by Distribution Matching.
Revised: Imitation Learning by Distribution Matching.

Revised: Classical IL approaches [26, 27] employed feature expectation matching to learn a policy from behavioral descriptors of expert demonstrations.

Revised: Similarly, a recent body of work [11, 12] in Inverse Optimal Control utilizes geometric behavioral descriptors that are similar to ours to learn cost functions for manipulation tasks from a few demonstrations.

Original: Distribution matching for IL has been used in moreclassical approaches [23], and recently gained popularity with the advent of Generative Adversarial
Revised: Whereas these approaches linearly match the moments/expectations of their features, our method instead matches a versatile distribution over non-linear features using the reverse Kullback-Leibler Divergence (KL) [28]. Directly matching the distribution rather than its moments means that our method can represent complex and multi-modal distributions. More recently, a new class of distribution matching approaches has gained popularity with the advent of Generative Adversarial

Original: Nets [24].
Revised: Nets [29].

Original: Starting from GAIL [5], a whole class of distribution matching based IL approacheswas developed [4, 25, 26, 27, 28, 29].
Revised: Starting from GAIL [5], a whole class of distribution matching based IL approaches was developed [4, 30, 31, 32, 33, 34].

Original: These approaches commonly work in a step-based setting,minimizing some divergence between the state-action occupancy or state marginals.
Revised: All of these approaches commonly work in a step-based setting and minimize some divergence between the state-action occupancy or state marginals.

Original: They generallydo not consider versatile behavior, and additionallyhave the major disadvantage that they need tointeract with the environment during training in order to generate states to match.
Revised: They generally do not consider versatile behavior, and additionally need to interact with the environment during training in order to generate states to match.

Original: We instead work ina trajectory-based setting, abstracting away the system’s dynamics. This abstraction allows us to learnentirely offline, i.e., without any environment interactions.
Revised: We instead work in a trajectory-based setting , which allows us to abstract away the environment’s dynamics and learn entirely ofﬂine, i.e., without any environment interactions.

Original: So far, distribution matching approacheshave only found limited applicability in this setting, arguably due to the simplicity of maximumlikelihood approaches for skills.

Original: Yet, as previously discussed, maximum likelihood methods oftenperform poorly on versatile data.
Revised: Common methods in this setting usually employ a maximum likelihood approach, which, as previously discussed, performs poorly on versatile data.

Original: As a solution to this problem, Becker et al.
Revised: As a solution to this problem, Becker et al.

Original: [13] propose using the Information-Projection [30] instead, which is able to focus on individual modes of data rather thanaveraging over all modes of data.
Revised: [16] propose using the Information-Projection [35] instead, which is able to focus on individual modes of data rather than averaging over all modes of data.

Original: This approach has been extended to IRL in [31].
Revised: This approach has been extended to IRL [36].

Original: Here, we build onthese methods, generalizing them to sequential data and geometric representations.
Revised: Here, we build on these methods, generalizing them to sequential data and geometric behavioral descriptors.

Original: Further Relevant Approaches.
Revised: Comparison-Based Approaches.

Original: Brown et al.
Revised: Brown et al.

Original: [6] perform IRL from ranked demonstrations byutilizing a comparison-based loss.
Revised: [6] perform IRL by utilizing provided rankings of demonstrations to train a discriminator on a comparison-based loss.

Original: While originally requiring access to these rankings, Disturbancebased Reward Extrapolation (D-REX)
Revised: As the discriminator learns to favor samples with a high rank, it can be used as a reward function after training.

Original: [7] has  extended this approach to instead automatically generaterankings by fitting a Behavioral Cloning (BC) policy on the demonstrations and subsequently drawingsamples with different levels of noise from this policy.
Revised: This has been extended by Disturbance-based Reward Extrapolation (D-REX) [7], which automatically generates rankings by ﬁtting a BC policy on the demonstrations and subsequently draws samples with different levels of noise from this policy.

Original: Here, we adapt this method to our trajectorybased settingto serve as a baseline. 
Revised: We adapt D-REX to our trajectory-based setting as a baseline. Similar to VIGOR, this adapted method trains a model on expert demonstrations on training contexts to optimize GMM policies on novel test contexts. However, while D-REX learns a reward function to do so, VIGOR iteratively (re-)trains a discriminator on the expert demonstrations and policy samples.

Original:  Foundations
Revised:  Foundations

Original: In this section, we briefly cover the foundations and previous work that our method builds on.
Revised: In this section, we brieﬂy cover the foundations and previous work that our method builds on.

Original: Probabilistic Movement Primitives.
Revised: Probabilistic Movement Primitives.

Original: For a trajectory τ = ( τ 1 , · · · , τ T ) , ProMPs represent theelements as τ t =
Revised: For a trajectory τ = ( τ 1 , · · · , τ T ) , ProMPs represent the elements as τ t =

Original: Φ ( t ) T w  .
Revised: Φ ( t ) T w , where τ t represents a vector of the desired joint angles at time-step t [18] Note that w does not depends on t and that τ can thus be computed at an arbitrary resolution T .

Original: Here Φ ( t ) are time-dependent features, usually radial basis functionscentered around different time points, and w denotes the parameter vector.
Revised: Here Φ ( t ) are time-dependent features, usually radial basis functions centered around different time points, and w denotes the parameter vector.

Original: For a single demonstration,the parameters w are fitted using simple linear regression and compactly represent the entire trajectory. 
Revised: For a single demonstration, the parameters w are ﬁtted using simple linear regression and compactly represent the entire trajectory. [Figure]

Original: Variational Inference for Gaussian Mixture Models.
Revised: Variational Inference for Gaussian Mixture Models.

Original: We build on a recent class of distributionmatching approaches for GMMs using the Information-Projection [30], i.e., these approaches minimize the reverse Kullback-Leibler divergence (KL) [32] KL (cid:0) q ( w ) || p ( w ) (cid:1) between a model q ( w ) anda target distribution p ( w ) .
Revised: We build on a recent class of efﬁcient Variational Inference approaches that allow modeling versatile behavior with GMMs [16, 14]. These approaches use the Information-Projection [35] to minimize the reverse Kullback-Leibler Divergence KL (cid:0) q ( w ) || p ( w ) (cid:1) between a model q ( w ) and a target distribution p ( w ) .

Original: In our case, p ( w ) is the distribution of expert trajectories.
Revised: In our case, p ( w ) is the distribution of expert trajectories.

Original: In variationalinference (VI), p ( w ) is assumed to be given up to a normalizer.

Original: Efficient methods exist to use VI foroptimizing GMMs, which allow modeling versatile behavior [13, 11].

Original: Arenz et al.
Revised: Arenz et al.

Original: [11] introduced anupper-bound objective based on a variational decomposition.
Revised: [14] introduced an upper-bound objective based on a variational decomposition.

Original: The resulting approach, Variational Inference by Policy Search (VIPS) [11, 12], repeatedly solves the objective
Revised: The resulting approach, Variational Inference by Policy Search (VIPS) [14, 15], repeatedly solves the objective

Original: [Equation]
Revised: [Equation]

Original: where ˆ q denotes the model from the previous iteration.
Revised: where ˆ q denotes the model from the previous iteration.

Original: This objective decomposes into individualoptimizations for the components and gating distributions, which Arenz et al.
Revised: This objective decomposes into individual optimizations for the components and categorical distributions, which Arenz et al.

Original: [11] solve usingtrust-region methods from policy search [33, 34]. [Figure]
Revised: [14] solve using trust-region methods from policy search [37, 38].

Original: Distribution Matching for Gaussian Mixture Models.
Revised: Distribution Matching for Gaussian Mixture Models.

Original: We cannot directly work with VIPS, as wedo not know p ( w ) but only have the expert demonstrations.
Revised: We cannot directly work with VIPS, as we do not know p ( w ) but only have the expert demonstrations.

Original: To address this issue, Becker et al.
Revised: To address this issue, Becker et al.

Original: [13]propose using density ratio estimation [35] approaches to approximate the log(ˆ q ( w ) /p ( w )) term in
Revised: [16] propose using density ratio estimation [39] approaches to approximate the log(ˆ q ( w ) /p ( w )) term in

Original: Equation 1, removing the dependency on p ( w ) .
Revised: Equation 1, removing the dependency on p ( w ) .

Original: Specifically, they use logistic regression to estimatethe density ratio.
Revised: Speciﬁcally, they use logistic regression to estimate the density ratio.

Original: The resulting approach, Expected Information Maximization (EIM), resembles
Revised: The resulting approach, Expected Information Maximization (EIM), resembles

Original: Generative Adversarial Nets [24] where the discriminator effectively also estimates a density ratio.
Revised: Generative Adversarial Nets [29] where the discriminator effectively also estimates a density ratio.

Original:  Inferring Versatile Plans by Matching GeometricBehavioral Descriptors
Revised:  Inferring Versatile Behaviors by Matching Geometric Descriptors

Original: Our approach estimates versatile plans in form of GMM distributions over ProMP weights for alltraining and test configurations of the given task.
Revised: Our approach estimates versatile behavior in form of GMM distributions over ProMP weights for all training and test contexts of the given task.

Revised: We note that while this setup requires the test contexts in advance, it does not need interaction with the real environment since the behavioral descriptors are directly computed from the proposed trajectory and its context.

Original: To obtain a good match with the demonstrator’sdistribution, we utilize full-covariance GMM components.
Revised: To capture correlations in the demonstrations, we utilize full-covariance GMM components.

Original: For each plan, we compute geometricfeatures such as distances to key-points, like via-points the agent has to reach, or the edges of anobject to manipulate.
Revised: For each planned trajectory, we compute concise descriptors such as distances to key-points that isolate the performance of the trajectory from its context to allow for sample-efﬁcient generalization to novel contexts.

Original: We directly build on EIM [13] for inferring GMMs from demonstrations bymatching the agent’s and demonstrator’s trajectory distributions in the given geometric feature space.
Revised: We then build on EIM [16] and use a discriminator to infer GMMs from demonstrations by matching the learner’s and expert’s distributions in the resulting geometric feature space.

Original: Towards this end, we generalize EIM to sequential data and modify the policy parameterizationto accommodate for our setting.
Revised: Towards this end, we generalize EIM to sequential data and modify the policy parameterization to accommodate for our setting.

Original: Figure 2 illustrates an overview of our approach.
Revised: Figure 2 illustrates an overview of our approach.

Original: We providepseudocode in Appendix A.
Revised: We provide pseudocode in Appendix A.

Original: Context-Specific Mixture Policies We want to model the expert’s behavior for a given small set oftraining configurations of the task, denoted as training contexts c train , and a set of unseen test contextsc test .
Revised: Context-Speciﬁc Mixture Policies We want to model the expert’s behavior for a given small set of training conﬁgurations of the task, denoted as training contexts c train , and a set of unseen test contexts c test .

Original: We consider the set of all contexts as c = c train ∪ c test .
Revised: We consider the set of all contexts as c = c train ∪ c test .

Original: As context-dependent GMMs 1 are hardto learn and we assume a relatively small amount of training contexts, 6 in all our experiments, withmultiple demonstrations per train context, we resort to maintaining a single GMM-Policy q c ( w ) foreach context c .
Revised: As context-dependent GMMs 2 are hard to learn and we assume a relatively small amount of training contexts, 6 in all our experiments, with multiple demonstrations per training context, we resort to maintaining a single GMM-Policy q c ( w ) for each context c .

Original: This non-amortized formulation simplifies the policy updates as we can optimize thepolicy for each context individually, given the density ratio estimator, using highly efficient, tailoredmethods for full-rank GMM approximations [11, 12].
Revised: This non-amortized formulation simpliﬁes the policy updates as we can optimize the policy for each context individually, given the density ratio estimator, using highly efﬁcient, tailored methods for full-rank GMM approximations [14, 15].

Original: Distribution Matching of Behavioral Descriptors We assume access toa behavioral descriptor
Revised: Distribution Matching of Behavioral Descriptors We assume access to behavioral descriptors

Original: O = f  ( w )  .
Revised: O = f c ( w ) that represent features encoded in the parameter vector w with respect to the context c of w .

Original: For simplicity, we will assume that the mapping f  between w
Revised: For simplicity, we will assume that the mapping f c between w

Original: and O is deterministic,such that we can easily evaluatethe behavioral descriptor O for any desired plan w .
Revised: and O is deterministic, such that we can easily evaluate O for any desired plan w .

Original: We want tomatch the distribution of behavioral descriptors of the demonstrator while optimizing for q ( w ) , i.e,
Revised: We want to match the distribution of behavioral descriptors of the demonstrator while optimizing for q ( w ) , i.e,

Original: [Equation]
Revised: [Equation]

Original: where p ( O | w ) is a Dirac delta distribution defined by the mapping f  and q ( O )
Revised: where p ( O | w ) is a Dirac delta distribution deﬁned by the mapping f c and q ( O )

Original: = (cid:82) p ( O | w ) q ( w ) d w .
Revised: = (cid:82) p ( O | w ) q ( w ) d w .

Original: As the mapping between w and O is deterministic, the EIM algorithm can be directly applied to this setup with the difference that the discriminator is trained using O instead of w .
Revised: As the mapping between w and O is deterministic, the EIM algorithm can be directly applied to this setup with the difference that the discriminator is trained using O instead of w .

Original: Moreover, while welearn different distributions q c ( w ) for each context c , we use the same discriminator for all contexts.
Revised: Moreover, while we a learn different distributions q c ( w ) for each context c , we use the same discriminator for all contexts.

Original: This way, we can infer distributions  q ctest ( w )  , i.e. inferring versatile plans for unseen scenarios.
Revised: This way, we can infer distributions { q c ( w ) | c ∈ c test } , i.e. infer versatile trajectories for unseen scenarios.

Original: Density Ratio Estimation for Sequential Behavioral Descriptors
Revised: Density Ratio Estimation for Sequential Behavioral Descriptors

Original: As in our case, w representsthe desired trajectory, the behavioral descriptors are typically computed per time-step, i.e., O =( o 1 , . ..
Revised: As in our case, w encodes the desired trajectory τ , the behavioral descriptors are typically computed per time-step, i.e., O = ( o 1 , . .

Original: , o T ) T .
Revised: , o T ) T , where each o t is a feature vector.

Original: To enable the classifier to deal with sequential data, we consider a sequentialneural network ϕ ( O )=
Revised: To enable the classiﬁer to deal with sequential data, we consider a sequence-to-sequence neural network φ ( O )

Original: ϕ
Revised: =

Original:  (( o 1 , . . .
Revised: φ (( o 1 , . . .

Original: , o T ) T ) =
Revised: , o T ) T ) =

Original: ( y 1 , . . .
Revised: ( y 1 , . . .

Original: , y T ) T that receives a sequence of inputs
Revised: , y T ) T . The network receives a sequence of inputs

Original: O =
Revised: O =

Original: ( o 1 , . .
Revised: ( o 1 , . .

Original: , o T ) T and outputs a sequence of values ( y 1 , . .
Revised: , o T ) T and outputs a sequence of values ( y 1 , . .

Original: , y T ) T , where each y i depends onmultiple
Revised: , y T ) T , where each y i may depend on multiple o j .

Original: o j depending on the size of the receptive field of the network.

Original: In each iteration of ouroptimization, we (re-)train this network to discriminate between the provided demonstrations O ( p )of c train , and an equal number of policy samples O ( q ) drawn uniformly from c .
Revised: In each iteration of our optimization, we (re-)train this network to discriminate between the provided demonstrations O ( p ) of c train , and an equal number of policy samples O ( q ) drawn uniformly from c .

Original: We first considerthe case where we want to discriminate full trajectories as in EIM.
Revised: We ﬁrst consider the case where we want to discriminate full trajectories as in EIM.

Original: In this case, we can consider thesummed output h ( O )

Original: = (cid:80)

Original: [-Ti =
-]
Revised: [-Ti =
-]

Original: ϕ i ( O ) with ϕ

Original: i

Original: ( O ) = y i of thesequential network for discrimination. The discriminator cantherefore straightforwardly be trained on a binary cross-entropy loss of thesum of sequence values ˆ y = (cid:80) Tt =

Revised: Denoting the sigmoid function as σ , the discriminator can straightforwardly be trained on a binary cross-entropy loss of the sum of sequence values ˆ y = (cid:80) Tt =

Original: y t , i.e., by minimizing
Revised: y t , i.e., by minimizing

Original: BCE

Original: [Equation]
Revised: [Equation]

Original: , w.r.t.
Revised: w.r.t.

Original: ϕ ( O ) , recovering the log density ratio ϕ ( O ) =
Revised: φ ( O ) . Equation 2 recovers the log density ratio φ ( O ) =

Original: log O ( p ) − log O ( q ) at convergence [35].
Revised: log O ( p ) − log O ( q ) at convergence [39].

Original: However, this cost function only provides a single training sample per trajectory and is therefore hardto use for a small number of demonstrations.
Revised: However, this cost function only provides a single training sample per trajectory and is therefore hard to use for a small number of demonstrations.

Original: By utilizing the sequential structure of the trajectories,the classification error can be computed for each time-step, i.e. we minimizeBCE step (cid:0) ϕ [Equation] O ) , O
Revised: By utilizing the sequential structure of the trajectories, the classiﬁcation error can be computed for each time-step, i.e. we minimize

Original: [Equation]

Original: p ) , O

Original: [Equation]

Original: q ) (cid:1)

Original: [Equation]
Revised: [Equation]

Original: w.r.t.
Revised: w.r.t.

Original: ϕ ( O ) .
Revised: φ ( O ) .

Original: Finally, we train an ensemble of discriminators instead of a single one, using theiraverage logit as the log density ratio estimate for the policy updates.
Revised: Finally, we train an ensemble of discriminators instead of a single one, using their average logit as the log density ratio estimate for the policy updates.

Original: For the network architecture, wefind that  1 d convolutional neural networks ( 1 d -CNNs) work best in our experiments.
Revised: For the network architecture, we ﬁnd that simple 2 − 4 layer 1 d convolutional neural networks ( 1 d -CNNs) work best in our experiments.

Original: We investigatethis choice  in Section 5.
Revised: We compare this choice to other sequential network architectures the ablations in Appendix D.

Original: Geometric Behavioral Descriptors VIGOR utilizes geometric descriptors that are insensitive to thetask configuration to generalize to new configurations.
Revised: Geometric Behavioral Descriptors VIGOR utilizes geometric descriptors that abstract away from a concrete task conﬁguration to generalize to new conﬁgurations.

Original: To this end, we encode the current state alongthe desired trajectory with respect to geometric, context-independent features rather than absolutevalues.
Revised: To this end, we encode the current state along the desired trajectory with respect to relative geometric features rather than absolute values.

Revised: The resulting geometric descriptors cause trajectories of similar performance to appear similar in feature space regardless of their context, allowing for sample-efﬁcient generalization to novel contexts.

Original: Such a descriptor space can often be straightforwardly constructed from the geometry of thetask. For example, in object manipulation, the behavioral descriptor can be composed of distances ofthe end-effector to key-points if the object tomanipulate, e.g., the corners of a box to push.
Revised: We note that such a descriptor space can often be straightforwardly constructed from the geometry of the task by composing distances of the end-effector to key-points of the objects in a scene. For example, in object manipulation, this space can be composed of distances of the end-effector to the corners of a box to push.

Original:  Experiments
Revised:  Experiments

Original: All experiments use human demonstrations.
Revised: All experiments use human demonstrations.

Revised: The tasks were selected due to their highly multi-modal nature and because they allow for an easy collection of human demonstrations.

Original: We instructed the demonstrators to solve the task invarying ways to create versatile solutions.
Revised: We instructed the demonstrators to solve the task in varying ways to create versatile solutions.

Original: As a preprocessing step, we fit a ProMP for each humandemonstration, and filter the resulting ProMPs such that only successful demonstrations remain.
Revised: As a preprocessing step, we ﬁt a ProMP for each human demonstration and ﬁlter the resulting ProMPs such that only successful demonstrations remain.

Original: Thisstep ensures that the human demonstrations can be imitated by the learner. It also comes with thebenefit of compressing all demonstrations to a fixed length in a principled fashion.
Revised: This step ensures that the human demonstrations can be imitated by the learner and conveniently allows to compress all demonstrations to a ﬁxed length in a principled fashion.

Original: Unless notedotherwise, all experiments use 6 training and 6 test contexts. We repeat each experiment for 10random seeds.
Revised: We use the resulting preprocessed expert demonstrations for all experiments unless otherwise stated. Similarly, unless noted otherwise, all experiments use 6 training and 6 test contexts, and

Original: We train a mixture with 5 components for each configuration, method and environment.
Revised: GMM components for each conﬁguration and method.

Revised: Experimentally, this number of components is sufﬁcient for matching the distribution of expert descriptors for the considered tasks, whereas more components generally only lead to spurious improvements.

Original:  Unless noted otherwise, we do not train the weight of these components, using a uniform weightinginstead.
Revised: We experiment with less components in the ablations in Appendix D. Unless noted otherwise, we do not train the categorical distribution of the GMM components, using a uniform distribution instead.

Original: We evaluate this mixture using a number of samples for each component of each test context,and then rank these components according to  their performance.
Revised: We evaluate the GMM policies using a number of samples for each component of each test context, and then rank these components according to the average performance of their samples.

Original: We then average over thesamples ofthe best componentfor each test context, and report statistics of the resulting value over the randomseeds, which allows us to obtain statistics about the best to the fifth-best learned components.
Revised: We then report statistics of the best component of the resulting value over the random seeds.

Original: Wereport the best component in order to avoid outlier components disrupting our evaluations.
Revised: Note that reporting the best component leads to a fair comparison to unimodal approaches.

Original: Forbaselines that do not use mixture policies, we instead draw samples of the trained policies for eachtest context and averageover the performance of all of them.
Revised: For baselines that do not use GMM policies, we instead draw samples of the trained policies for each test context and average the performance of all of them.

Revised: For each experiment, we report a measure of how well the task is performed in the main paper, and additional success rates in Appendix C. We repeat each experiment for 10 random seeds.

Original: Appendix E shows all hyperparametersand training details.V I GO
Revised: Appendix E shows all hyperparameters and training details.

Original: R E M + D - R E X
Revised: Baselines.

Original: M BC ( T ) BC ( E ) GA I L M BC ( S )
Revised: We compare VIGOR to a number of IL baselines, using three state-action baselines, namely GAIL [5], behavioral cloning (BC(S)) and behavioral cloning with a GMM policy (BCGMM(S)) [13].

Original: BC ( S ) D i s t a n ce T o B ound a r i e s 0 .0 . V I GO
Revised: For all experiments, the state-action baselines receive geometric descriptors of the current state and a time-step as state information, and output the velocities of the robot joints as actions.

Original: R E M + D - R E X
Revised: Further, we also consider three trajectory-based methods.

Original: M BC ( T ) BC ( T ) GA I L M BC ( S )
Revised: These are standard behavioral cloning (BC(T)), behavioral cloning with a GMM policy (BC-GMM(T))

Original: BC ( S ) 0 .0 .0 .0 . D i s t a n ce T o B ound a r i e s [ m ] Demonstrations 0 .0 .
Revised: (see e.g., Zhou et al.

Original: Baselines. We compare VIGOR to a number of IL baselines, using three trajectory-based methods,namely, standard behavioral cloning (BC(T)), behavioral cloning with a mixture model (MBC(T)), anda modified version of D-REX (EM+D-REX).Further, we also consider three state-action baselines.
Revised: [18]), and a modiﬁed version of D-REX (EM+D-REX).

Original: These are GAIL [5], as well as behavioral cloning with and without a multi-modal policy (BC(S),and MBC(S)).
Revised: The latter is chosen due to its similarities to VIGOR, which are detailed in Appendix B.1.

Original: We refer to  B for details on the baselines and to  C for furtherdetails and results on allexperiments described below.
Revised: We refer to Appendix B for details on the baselines and to Appendix C for further results on all experiments described below.

Original: Planar Reacher.
Revised: Planar Reacher.

Original: We consider a planar robot with 5 Degrees of Freedom (DoF) and joints of length 1 . Here, the goal is to reach an intermediate target area with the end-effector  before ending the trajectoryin a final target area.The left and middle panels of Figure 1 depict example demonstrations.
Revised: First, we consider the introductory task shown in Figure 1. In this task, the goal is to reach an intermediate target area with the end-effector of a 5 Degrees of Freedom (DoF) robot with joints of length 1 before ending the trajectory in a ﬁnal target area.

Original: Thegeometric representation for this task is given by the euclidean distance of the end-effector to bothtargets, paired with the average velocity and acceleration of the joints to encode the smoothness ofthe motion.
Revised: The geometric descriptors for this task are given by the euclidean distance of the end-effector to both targets, paired with the average velocity and acceleration of the joints to encode the smoothness of the motion.

Original: This yields a total of 4 features per time-step.
Revised: This yields a total offeatures per time-step.

Revised: The left and middle panels of Figure 1 depict example demonstrations for two training contexts (top) and the corresponding distances between the end-effector and targets over time (bottom).

Original: We explore other choices for the geometricrepresentation in Section 5. For evaluation, we use the Distance to Boundaries of the target areas.
Revised: We explore other choices for the geometric descriptors in the ablations in Appendix D. For evaluation, we use the Distance to Boundaries of the target areas.

Original: Forthe first target, we consider the minimum distance of the end-effector to the boundary; for the secondtarget, we consider the distance at the last time-step.
Revised: For the ﬁrst target, we consider the minimum distance of the end-effector to the boundary; for the second target, we consider the distance at the last time-step.

Original: We train all methods on 5 demonstrations pertraining context.
Revised: We train all methods on 5 demonstrations per training context.

Original: The left of Figure 3 shows results on test contexts.7 -D Panda Reacher.
Revised: The left of Figure 3 shows results on test contexts.

Revised: We ﬁnd that only VIGOR and EM+D-REX can consistently get close to both targets.

Revised: Yet, EM+D-REX can only learn a single mode and fails to reproduce the versatility of the demonstrations.

Revised: We investigate this in more detail in Appendix C.1.

Original: Next, we repeat the above point-reaching task using a 7 DoF Franka Emika Panda robot and 3 -D intermediate and goal positions.
Revised: 7 -D Panda Reacher. We repeat the above point-reaching task using a 7 DoF Franka Emika Panda robot and 3 -D intermediate and goal positions.

Original: Since we only care about the end-effector’sposition but not its rotation,we can ignore the last DoF, leading to an  action dimension of 6 .
Revised: Since we do not care about the end-effector’s rotation, the last DoF can be ignored, leading to an effective action dimension of 6 .

Revised: We use 8 basis functions per action dimension for the MPs, leading to a total of MP dimension of 48 and thus a high-dimensional problem space.

Original: Theexpert demonstrations are collected using a teleoperation setup with a virtual twin that records jointvalues over time.
Revised: The expert demonstrations are collected using a teleoperation setup with a virtual twin that records joint values over time.

Original: Appendix C.2 explains the setup in more detail.
Revised: Appendix C.2 explains the setup in more detail.

Original: We use the same geometricrepresentation as for the planar reacher task.
Revised: We use the same geometric descriptors as for the planar reacher task.

Original: The right of Figure 3 shows results on test contexts.
Revised: The right of Figure 3 shows target distances on test contexts.

Original: Similar to the previous task, only VIGOR and EM+D-REX can consistently get close to bothtargets. Yet, EM+D-REX can only learn a single mode and fails to reproduce the versatility of thedemonstrations.
Revised: Similar to the previous task, only VIGOR and EM+D-REX can consistently get close to both targets while the other methods regularly fail to pursue at least one of the two targets.

Original: We investigate this in more detail in the appendix.

Original: Box Pusher.
Revised: Box Pusher.

Original: We also evaluate our approach on a complex robot manipulation task.
Revised: We also evaluate our approach on a contact-rich robot manipulation task.

Original: In this task, asimulated Franka Emika Panda robot has a rod attached to its end-effector and needs to use it to pusha rectangular box to a given goal position and orientation.
Revised: In this task, a simulated Franka Emika Panda robot has a rod attached to its end-effector and needs to use it to push a rectangular box to a given goal position and orientation.

Original: The box always starts at the same position,and the task context is given by the desired ( x, y ) translation of the box, plus its desired rotation indegrees.
Revised: The box always starts at the same position, and the task context is given by the desired ( x, y ) translation of the box, plus its desired rotation in degrees.

Original: For simplicity, the task is learned in task space with a fixed height.
Revised: For simplicity, the task is learned in task space with a ﬁxed height.

Original: The demonstrationscontain the starting position of the robot as well as the trajectory, yielding versatile solutions asthe box can be pushed from both inside and outside of the box.
Revised: The demonstrations contain the starting position of the robot as well as its trajectory.

Original: The geometric representation forthis task considers the ( x, y ) distances of the end-effector to the initial and desired final positionof the corners of the box. To facilitate generalization between contexts, we translate the basis of
Revised: Here, the solutions are versatile

Original: V I GO
Revised: V I GO

Original: R
Revised: R

Original: E M +
Revised: E M +

Original: D - R E
Revised: D - R E

Original: X M BC ( T )
Revised: X BCG MM ( T )

Original: BC ( T )
Revised: BC ( T )

Original: GA I L
Revised: GA I L

Revised: BCG

Original: M BC ( S )
Revised: MM ( S ) BC ( S )

Revised: D i s t a n ce

Revised: T o B ound a r i e s 0 .0 .

Revised: V I GO

Revised: R

Revised: E M +

Revised: D - R E

Revised: X BCG MM ( T )

Revised: BC ( T )

Revised: GA I L BCG MM ( S )

Original: BC ( S ) 0 .0 .0  .
Revised: BC ( S ) 0 .0 .0 .0 .

Original: D i s t a n cet o T  a r g e t [ m ]
Revised: D i s t a n ce T o B ound a r i e s [ m ]

Original: Demonstrations 
Revised: Demonstrations 0 .0 .

Original: Initial Box Position     10 −10 −
Revised: Similarly, using a discriminator to iteratively optimize the policies is advantageous on the reaching tasks, as seen from VIGOR and GAIL. Finally, the clear distribution matching objective of VIGOR and EM+D-REX seems crucial for imitation from a low number of versatile demonstrations. VIGOR uniquely combines the above properties and is able to reliably reach both targets for both tasks.

Original: Number of Training Contexts
Revised: [Figure] [Equation] [Figure]

Original: D i s t a n ce t o T a r g e t [ m ]
Revised: D i s t a n ce t o T a r g e t [ m ]

Original: VIGOR
Revised: [Figure]

Original: EIM+D-REX
Revised: because the box can be pushed both from the inside and outside.

Original: BC(T) GAIL these distances such that the desired final box aligns with the origin. This geometric representationof the task allows for IL without simulating or explicitly modeling the box, significantly speedingup the training process and facilitating generalization to both new contexts and real-world robots.
Revised: The geometric descriptors consist of the ( x, y ) distances of the end-effector to the initial and desired ﬁnal position of the corners of the box. To facilitate generalization between contexts, we translate the basis of these distances such that the desired ﬁnal box aligns with the origin of the coordinate system.

Original: We additionally add  end-effector velocity and acceleration as well as the time-step, resulting in 19features per time-step.
Revised: We additionally add the end-effector velocity and acceleration as well as the time-step, resulting in 19 features per time-step.

Revised: This geometric description of the task allows for IL without simulating or explicitly modeling the box, signiﬁcantly speeding up the training process and facilitating generalization to both new contexts and real-world robots.

Original: We use 3 demonstrations per train context. We evaluate the average distanceof the corners of the final box to that of the desired final position.
Revised: We use 3 demonstrations per training context and evaluate the average distance of the corners of the ﬁnal box to that of the desired ﬁnal position.

Original: Figure 4 shows the results forour experiments on the left and the effects of varying numbers of train-contexts on the performanceof VIGOR and several baselines on the right.
Revised: The contact with the box causes a mismatch between a planned trajectory and its execution on the robot, which we illustrate in Appendix C.3. Results are shown in the left of Figure 4. Additionally, the right of Figure 4 shows how the performance varies when changing the number of training contexts.

Original: Finally, we illustrate how the planned trajectories of VIGOR perform on a real robot.
Revised: Finally, we illustrate how the planned trajectories of VIGOR perform on a real robot.

Original: To this end, we create a real-world replica of the simulated boxthat the robot manipulates.
Revised: To this end, we create a real-world replica of the simulated box that the robot manipulates.

Original: Rollouts for different policy components of VIGOR trained on 24 traincontexts on an exemplary test context are shown in Figure 5.
Revised: Rollouts for different policy components of VIGOR trained on 24 training contexts on an exemplary test context are shown in Figure 5.

Original: Videos can be found in the supplement.

Original: Ablations.
Revised: Ablations.

Original: We conduct extensive ablation studies to investigate how our design choices affect themodel.
Revised: We conduct extensive ablation studies to investigate how our design choices affect our approach.

Original: In Figure 6 we show some of these results for VIGOR on the Planar Reaching task.

Original: For

Original: A ng l e s J

Original: o i n t w

Original: i s e G

Original: eo m e t r

Original: i c 10 − 10 − 10 − Raw ProMP MLP LSTM1d-CNN B C E S t ep w i s e B C E
Revised: We ﬁnd that a concise choice of behavioral descriptors, a preprocessing of the expert demonstrations to ProMPs, and a 1 d -CNN discriminator are crucial for performance.

Original: Component #4 # #2 # # D i s t an c e t o B ounda r i e s Representation Demonstrations
Revised: We also notice that the policy beneﬁts from additional components and that we see modest improvements for an ensemble of discriminators and the stepwise loss proposed in Equation 3.

Original: Architecture #Components #Discriminators

Original: Disc. Loss #Train Contexts more thoroughexplanations and extensive ablations of both VIGOR and EM+D-REX, we refer to Appendix D.
Revised: We refer to Appendix D for the full results and more thorough ablations of both VIGOR and EM+D-REX.

Original:  Limitations
Revised:  Limitations

Original: Scale.
Revised: Scale.

Original: We currently train and maintain a separatefull-rank GMM  for each training and test configuration.
Revised: We currently train and maintain a separate GMM over ProMP parameters for each context, leading to linear space and time complexity w.r.t.

Original: While we show that VIGOR can generalize well from a few configurations, it currentlyhas linear space and time complexity w.r.t. the total number ofconfigurations, impeding its use forsettings with large or continuous numbers of contexts.
Revised: the total number of contexts.

Original: One way to address this challenge would be toinstead train a joint Mixture Density Network q ( w | c ) over all configurations, which is left for futurework.
Revised: One way to address this challenge would be to instead train a joint Mixture Density Network q ( w | c ) over all contexts, which is left for future work.

Revised: Additionally, representing trajectories with a single ProMP limits them to single smooth movements.

Revised: To alleviate this, MP chaining [40, 41, 42] can be straightforwardly integrated into our approach to allow for the representation of more complex movements

Original: Geometric Representation.
Revised: Geometric Descriptors.

Original: VIGOR facilitates generalization to novel task configurations by usinggeometric task representations.
Revised: VIGOR facilitates generalization to novel task conﬁgurations by using geometric descriptors.

Original: While this allows the user to include domain knowledge, it alsorequires a clear idea about which aspects of the task are important.
Revised: While this allows the user to include domain knowledge, it also requires a clear idea about which aspects of the task are important.

Original: One natural extension wouldthus be to instead extract a compact geometric representation of the task, for example by using NNsprocessing point clouds [37, 38].
Revised: Instead, geometric descriptors could be automatically extracted from the task for example by using NNs processing point clouds [43, 44].

Original: Re-training.
Revised: Re-training.

Original: Finally, our method assumes that the test configurations c test are known in advance.
Revised: Finally, our method assumes that the test conﬁgurations c test are known in advance.

Original: Since the test configurations are optimized jointly with the training configurations, the only way tointegrate new test configurations is to re-train the algorithm.
Revised: Since the test conﬁgurations are optimized jointly with the training conﬁgurations, the only way to integrate new test conﬁgurations is to re-train the algorithm.

Original: In the future, we want to alleviate thisissue by recovering a reward representation from ranked intermediate policies of the given trainingconfigurations, essentially combining our approach with some of the ideas of D-REX.
Revised: In the future, we want to alleviate this issue by recovering a reward representation from ranked intermediate policies of the given training conﬁgurations, essentially combining our approach with some of the ideas of D-REX.

Original:  Conclusion
Revised:  Conclusion

Original: We proposed VIGOR, a novel feature-matching approach to Imitation Learning in environmentswith versatile solutions.
Revised: We proposed VIGOR, a novel feature-matching approach to Imitation Learning in environments with versatile solutions.

Original: Utilizing a combination of movement primitives, mixture policies andmatching geometric behavioral descriptors, our method can closely imitate the behavior distributionof human experts from a few demonstrations.
Revised: Utilizing a combination of movement primitives, mixture policies and matching geometric behavioral descriptors, our method can closely imitate the behavior distribution of human experts from a few demonstrations.

Original: We show the effectiveness of our approach on a suiteof challenging robot coordination tasks.
Revised: We show the effectiveness of our approach on a suite of challenging robot coordination tasks.

Original: The results show that VIGOR is able to closely match thedistribution of the demonstrator, outperforming the chosen baselines in all considered settings. 
Revised: The results show that VIGOR is able to closely match the distribution of the demonstrator, outperforming the chosen baselines in all considered settings. Acknowledgments

Revised: NS and GN were supported by the Carl Zeiss Foundation under the project JuBot (Jung Bleiben mit

Revised: Robotern).

Revised: The authors acknowledge support by the state of Baden-Württemberg through bwHPC.