Title: ADAPTIVE ADVERSARIAL IMITATION LEARNING

Abstract: We present the ADaptive Adversarial Imitation Learning (ADAIL) algorithm for learning adaptive policies that can be transferred between environments of varying dynamics, by imitating a small number of demonstrations collected from a single source domain. This is an important problem in robotic learning because in real world scenarios 1) reward functions are hard to obtain, 2) learned policies from one domain are difficult to deploy in another due to varying source to target domain statistics, 3) collecting expert demonstrations in multiple environments where the dynamics are known and controlled is often infeasible. We address these constraints by building upon recent advances in adversarial imitation learning; we condition our policy on a learned dynamics embedding and we employ a domain-adversarial loss to learn a dynamics-invariant discriminator. The effectiveness of our method is demonstrated on simulated control tasks with varying environment dynamics and the learned adaptive agent outperforms several recent baselines.

1. INTRODUCTION:

Humans and animals can learn complex behaviors via imitation. Inspired by these learning mechanisms, Imitation Learning (IL) has long been a popular method for training autonomous agents from human-provided demonstrations. However, human and animal imitation differs markedly from commonly used approaches in machine learning. Firstly, humans and animals tend to imitate the goal of the task rather than the particular motions of the demonstrator (Baker et al., 2007) . Secondly, humans and animals can easily handle imitation scenarios where there is a shift in embodiment and dynamics between themselves and a demonstrator. The first feature of human IL can be represented within the framework of Inverse Reinforcement Learning (IRL) (Ng et al., 2000; Abbeel & Ng, 2004; Ziebart et al., 2008) , which at a high level casts the problem of imitation as one of matching outcomes rather than actions. Recent work in adversarial imitation learning (Ho & Ermon, 2016; Finn et al., 2016) has accomplished this by using a discriminator to judge whether a given behavior is from an expert or imitator, and then a policy is trained using the discriminator expert likelihood as a reward. While successful in multiple problem domains, this approach makes it difficult to accommodate the second feature of human learning: imitation across shifts in embodiment and dynamics. This is because in the presence of such shifts, the discriminator may either simply use the embodiment or dynamics to infer whether it is evaluating expert behavior, and as a consequence fails to provide a meaningful reward signal.

In this paper we are concerned with the problem of learning adaptive policies that can be transferred to environments with varying dynamics, by imitating a small number of expert demonstrations collected from a single source domain. This problem is important in robotic learning because it is better aligned with real world constraints: 1) reward functions are hard to obtain, 2) learned policies from one domain are hard to deploy to different domains due to varying source to target domain statistics, and 3) the target domain dynamics oftentimes changes while executing the learned policy. As such, this work assumes ground truth rewards are not available, and furthermore we assume that expert demonstrations come from only a single domain (i.e. an instance of an environment where dynamics cannot be exactly replicated by the policy at training time). To the best of our knowledge, this is the first work to tackle this challenging problem formulation.

Our proposed method solves the above problem by building upon the GAIL (Ho & Ermon, 2016; Finn et al., 2016) framework, by firstly conditioning the policy on a learned dynamics embedding ("context variable" in policy search literature (Deisenroth et al., 2013) ). We propose two embedding approaches on which the policy is conditioned, namely, a direct supervised learning approach and a variational autoencoder (VAE) (Kingma & Welling, 2013) based unsupervised approach. Secondly, to prevent the discriminator from inferring whether it is evaluating the expert behavior or imitator behavior purely through the dynamics, we propose using a Gradient Reversal Layer (GRL) to learn a dynamics-invariant discriminator. We demonstrate the effectiveness of the proposed algorithm on benchmark Mujoco simulated control tasks.

The main contributions of our work include: 1) present a general and novel problem formulation that is well aligned with real world scenarios in comparison to recent literature 2) devise a conceptually simple architecture that is capable of learning an adaptive policy from a small number of expert demonstrations (order of 10s) collected from only one source environment, 3) design an adversarial loss for addressing the covariate shift issue in discriminator learning.

2. RELATED WORK:

Historically, two main avenues have been heavily studied for imitation learning: 1) Behavioral Cloning (BC) and 2) Inverse Reinforcement Learning (IRL). Though conceptually simple, BC suffers from compound errors caused by covariate shift, and subsequently, often requires a large quantity of demonstrations (Pomerleau, 1989) , or access to the expert policy (Ross et al., 2011) in order to recover a stable policy.

Recent advancements in imitation learning (Ho & Ermon, 2016; Finn et al., 2016) have adopted an adversarial formation that interleaves between 1) discriminating the generated policy against the expert demonstrations and 2) a policy improvement step where the policy aims to fool the learned discriminator.

Dynamics randomization (Tobin et al., 2017; Sadeghi & Levine, 2016; Mandlekar et al., 2017; Tan et al., 2018; Pinto et al., 2017; Peng et al., 2018; Chebotar et al., 2018; Rajeswaran et al., 2016) has been one of the prevailing vehicles for addressing varying simulation to real-world domain statistics. This avenue of methods typically involves perturbing the environment dynamics (often times adversarially) in simulation in order to learn an adaptive policy that is robust enough to bridge the "Reality Gap".

While dynamics randomization has been explored for learning robust policies in an RL setting, it has a critical limitation in the imitation learning context: large domain shifts might result in directional differences in dynamics, therefore, the demonstrated actions might no longer be admissible for solving the task in the target domain. Our method (Figure 2 ) also involves training in a variety of environments with different dynamics. However, we propose conditioning the policy on an explicitly learned dynamics embedding to enable adaptive policies based on online system ID. Yu et al. (2017) adopted a similar approach towards building adaptive policies. They learn an online system identification model and condition the policy on the predicted model parameters in an RL setting. In comparison to their work, we do not assume access to the ground truth reward signals or the ground truth physics parameters at evaluation time, which makes this work's problem formulation a harder learning problem, but with greater potential for real-world applications. We will compare our method with Yu et al. (2017) in the experimental section.

Third person imitation learning (Stadie et al., 2017 ) also employs a GRL (Ganin & Lempitsky, 2014) under a GAIL-like formulation with the goal of learning expert behaviors in a new domain. In comparison, our method also enables learning adaptive policies by employing an online dynamics identification component, so that the policies can be transferred to a class of domains, as opposed to one domain. In addition, learned policies using our proposed method can handle online dynamics perturbations.

Meta learning (Finn et al., 2017) has also been applied to address varying source to target domain dynamics (Duan et al., 2017; Nagabandi et al., 2018) . The idea behind meta learning in the context of robotic learning is to learn a meta policy that is "initialized" for a variety of tasks in simulation, and then fine-tune the policy in the real-world setting given a specific goal. After the meta-learning phase, the agent requires significantly fewer environment interactions to obtain a policy that solves the task. In comparison to meta learning based approaches, fine-tuning on the test environment is not required in our method, with the caveat being that this is true only within the target domain where the dynamics posterior is effective.

2.1. BACKGROUND:

In this section we will briefly review GAIL (Ho & Ermon, 2016) . Inspired by GANs, the GAIL objective is defined as:

min θ max ω E π E [log D ω (s, a)] + E π θ [log(1 -D ω (s, a))]

(1)

Where π E denotes the expert policy that generated the demonstrations; π θ is the policy to imitate the expert; D is a discriminator that learns to distinguish between π θ and π E with generated stateaction pairs. In comparison to GAN optimization, the GAIL objective is rarely differentiable since differentiation through the environment step is often intractable. Optimization is instead achieved via RL-based policy gradient algorithms, e.g., PPO (Schulman et al., 2017) or off policy methods, e.g., TD3 (Kostrikov et al., 2018) . Without an explicit reward function, GAIL relies on reward signals provided by the learned discriminator, where a common reward formulation is

r ω (s, a) = -log(1 -D ω (s, a)).

3. ADAPTIVE ADVERSARIAL IMITATION LEARNING (ADAIL):

3.1 PROBLEM DEFINITION Suppose we are given a class E of environments with different dynamics but similar goals, a domain generator g(c) which takes in a code c and generates an environment e c ∈ E, and a set of expert demonstrations {τ exp } collected from one source environment e exp ∈ E. In adaptive imitation learning, one attempts to learn an adaptive policy π θ that can generalize across environments within E. We assume that the ground truth dynamics parameters c, which are used to generate the simulated environments, are given (or manually sampled) during the training phase.

3.2. ALGORITHM OVERVIEW:

We allow the agent to interact with a class of similar simulated environments with varying dynamics parameters, which we call "adaptive training". To be able to capture high-level goals from a small set of demonstrations, we adopt a approach similar to GAIL. To provide consistent feedback signals during training across environments with different dynamics, the discriminator should be dynamicsinvariant. We enable this desirable feature by learning a dynamics-invariant feature layer for the discriminator by 1) adding another head D R (c|s, a) to the discriminator to predict the dynamics parameters, and 2) inserting a GRL in-between D R and the dynamics-invariant feature layer. The discriminator design is illustrated in Figure 1 . In addition, to enable adaptive policies, we introduced a dynamics posterior that takes a roll-out trajectory and outputs an embedding, on which the policy is conditioned. Intuitively, explicit dynamics latent variable learning endows the agent with the ability to identify the system and act differently against changes in dynamics. Note that a policy can learn to infer dynamics implicitly, without the need for an external dynamics embedding. However, we find experimentally that policies conditioned explicitly on the environment parameters outperform those that do not. The overall architecture is illustrated in Figure 2 . We call the algorithm Adaptive Adversarial Imitation Learning (ADAIL), with the following objective (note that for brevity, we for now omit the GRL term discussed in Section 3.4):

min θ max ω,φ E π E [log D ω (s, a)] + E π θ (•|c) [log(1 -D ω (s, a))] + E τ ∼π θ (•|c) [log Q φ (c|τ )] (2)

Where c is a learned latent dynamics representation that is associated with the rollout environment in each gradient step; τ is a roll-out trajectory using π θ (•|c) in the corresponding environment; Q(c|τ ) is a "dynamics posterior" for inferring the dynamics during test time; The last term in the objective,

E τ ∼π θ (•|c) [log Q φ (c|τ )]

, is a general form of the expected log likelihood of c given τ . Note that, the posterior training is on-policy, meaning that the rollouts are collected online using the current 

Gradient Reversal Layer:



Dynamics invariant features:



Gradients have sign negated:

Figure 1 : Discriminator with Gradients Reversal Layer (GRL). The red layer is the GRL which reverses the gradients during backprop. The yellow layer is a dynamics-invariant layer that is shared with the classification task.

policy, thereby the last term of the objective is dependent on θ. One can employ various supervised and unsupervised methods towards optimizing this term. We will explore a few methods in the following subsections.

The algorithm is outlined in Algorithm 1.

3.3. ADAPTIVE TRAINING:

Adaptive training is achieved through 1) allowing the agent to interact with a class of similar simulated environments within class E, and 2) learning a dynamics posterior for predicting the dynamics based on rollouts. The environment class E is defined as a set of parameterized environments with n degrees of freedom, where n is the total number of latent dynamics parameters that we can change.

We assume that we have access to an environment generator g(c) that takes in a sample of the dynamics parameters c and generates an environment. At each time when an on-policy rollout is initiated, we re-sample the dynamics parameters c based on a predefined prior distribution p(c).

3.4. LEARNING A DYNAMICS-INVARIANT DISCRIMINATOR:

GAIL learns from the expert demonstrations by matching an implicit state-action occupancy measure. However, this formulation might be problematic in our training setting, where on-policy rollouts are collected from environments with varying dynamics. In non-source environments, the discriminator can no longer provide canonical feedback signals. This motivates us to learn a dynamics-invariant feature space, where, the behavior-oriented features are preserved but dynamicsidentifiable features are removed. We approach this problem by assuming that the behavior-oriented characteristics and dynamics-identifiable characteristics are loosely coupled and thereby we can learn a dynamics-invariant representation for the discriminator. In particular, we employ a technique called a Gradient Reversal Layer (GRL) (Ganin & Lempitsky, 2014) , which is widely used in image domain adaptation (Bousmalis et al., 2016) . The dynamics-invariant features layer is shared with the original discriminator classification head, illustrated in Figure 1 .

3.5. DIRECT SUPERVISED DYNAMICS LATENT VARIABLE LEARNING:

Perhaps one of the best latent representations of the dynamics is the ground truth physics parameterization (gravity, friction, limb length, etc). In this section we explore supervised learning for inferring dynamics. A neural network is employed to represent the dynamics posterior, which is learned via supervised learning by regressing to the ground truth physics parameters given a replay buffer of policy rollouts. We update the regression network using a Huber loss to match environment dynamics labels. Details about the Huber loss can be found in appendix A.1. During training, we condition the learned policy on the ground truth physics parameters. During evaluation, on the other hand, the policy is conditioned on the predicted physics parameters from the posterior.

We use (state, action, next state) as the posterior's input, i.e., Q φ (c|s, a, s ), and a 3-layer fullyconnected neural network to output the N-dimensional environment parameters. Note that one can use a recurrent neural network and longer rollout history for modeling complex dynamic structures, however we found that this was not necessary for the chosen evaluation environments.

3.6. VAE-BASED UNSUPERVISED DYNAMICS LATENT VARIABLE LEARNING:

For many cases, the number of varying latent parameters of the environment is high, one might not know the set of latent parameters that will vary in a real world laboratory setting, or the latent parameters are oftentimes strongly correlated (e.g., gravity and friction) in terms of their effect on environment dynamics. In this case, predicting the exact latent parameterization is hard. The policy is mainly concerned with the end effector of the latent parameters. This motivates us to use a unsupervised tool to extract a latent dynamics embedding. In this section, we explore a VAE-based unsupervised approach similar to conditional VAE (Sohn et al., 2015) with an additional contrastive regularization loss, for learning the dynamics without ground truth labels.

With the goal of capturing the underlying dynamics, we avoid directly reconstructing the (state, action, next state) tuple, (s, a, s ). Otherwise, the VAE would likely capture the latent structure of the state space. Instead, the decoder is modified to take-in the state-action pair, (s, a), and a latent code, c, and outputs the next state, s . The decoder now becomes a forward dynamics predictive model. The unsupervised dynamics latent variable learning method is illustrated in Figure 3 .

The evidence lower bound (ELBO) used is: (2007) and Hsu & Kira (2015) , to avoid the encoder learning an identity mapping on s , we add the following KL-based contrastive regularization to the loss:

L contrastive = KL(Q φ (s 0 , a 0 , s 0 )||Q φ (s 1 , a 1 , s 1 )) -min{KL(Q φ (s 2 , a 2 , s 2 )||Q φ (s 3 , a 3 , s 3 )), D 0 }

Where (s 0 , a 0 , s 0 ) and (s 1 , a 1 , s 1 ) are sampled from the same roll-out trajectory; (s 2 , a 2 , s 2 ) and (s 3 , a 3 , s 3 ) are sampled from different roll-out trajectories. D 0 is a constant. We use this regularization to introduce additional supervision in order to improve the robustness of the latent posterior.

The overall objective for the dynamics learner is

min φ,ψ -ELBO + λL contrastive (4)

where λ is a scalar to control the relative strength of the regularization term. The learned posterior (encoder) infers the latent dynamics, which is used for conditioning the policy. The modified algorithm can be found in the appendix (Algorithm 2).

4. EXPERIMENTS 4.1 ENVIRONMENTS:

To evaluate the proposed algorithm we consider 4 simulated environments: CartPole, Hopper, HalfCheetah and Ant. The chosen dynamics parameters are specified in table 3, and an example of one such parameter (HalfCheetah gravity component x) is shown in Figure 4 . During training the parameters are sampled uniformly from the chosen range. Source domain parameters are also given in the table 3. For each source domain, we collect 16 expert demonstrations.

Gym CartPole-V0: We vary the force magnitude in continuous range [-1, 1] in our training setting.

Note that the force magnitude can take negative values, which flips the force direction.

3 Mujoco Environments: Hopper, HalfCheetah, and Ant: With these three environments, we vary 2d dynamics parameters: gravity x-component and friction.

Environment Paramater 1 Parameter 2 Source CartPole-V0 F m [-1,1] F m = 1.0 Hopper Gx [-1.0, 1.0] F r [1.5, 2.5] Gx = 0.0, F r = 2.0 HalfCheetah Gx [-3.0, 3.0] F r [0.0, 2.0] Gx = 0.0, F r = 0.5 Ant Gx [-5.0, 5.0] F r [0.0, 4.0] Gx = 0.0, F r = 1.0

Table 1: Environments. F m = Force magnitude; Gx=Gravity x-component; F r = Friction. For each environment, we collect 16 expert demonstrations from the source domain.

4.2. ADAIL ON SIMULATED CONTROL TASKS:

Is the dynamics posterior component effective under large dynamics shifts?

We first demonstrate the effectiveness of the dynamics posterior under large dynamics shifts on a toy Gym environment, Cartpole, by varying 1d force magnitude. As the direction of the force changes, blindly mimicking the demonstrations collected from the source domain (F m = 1.0) would not work on target domains with F m < 0.0. This result is evident when comparing ADAIL to GAIL with dynamics randomization. As shown in Figure 5a , GAIL with Dynamics Randomization failed to generalize to F m < 0.0, whereas, ADAIL is able to achieve the same performance as F m > 0.0.

We also put a comparison with ADAIL-rand, where the policy is conditioned on uniformly random values of the dynamics parameters, which completely breaks the performance across the domains.

How does the GRL help improve the robustness of performance across domains?

To demonstrate the effectiveness of GRL in the adversarial imitation learning formulation, we do a comparative study with and without GRL on GAIL with dynamics randomization in the Hopper environment. The results are shown in Figure 5b . We demonstrate the overall performance of ADAIL by applying it to three Mujoco control tasks: HalfCheetah, Ant and Hopper. For each of the Mujoco environments, we vary 2 continuous dynamics parameters and we compare the performance of ADAIL with a few baseline methods, including 1) the PPO expert which was used to collect demonstrations; 2) the UP-true algorithm of Yu et al. (2017) , which is essentially a PPO policy conditioned on ground truth physics parameters; and 3) GAIL with dynamics randomization, which is unmodified GAIL training on a variety of environments with varying dynamics. The results of this experiment are show in in Figure 6 .

HalfCheetah The experiments show that 1) as expected the PPO expert (Plot 6a) has limited adaptability to unseen dynamics. 2) UP-true (Plot 6b) achieves similar performance across test environments. Note that since UP-true has access to the ground truth reward signals and the policy is conditioned on ground truth dynamics parameters, the Plot 6b shows an approximate expected upper bound for our proposed method since we do not assume access to reward signals during policy training, or to ground truth physics parameters at policy evaluation time. 3) GAIL with dynamics randomization (Plot ??) can generalize to some extent, but failed to achieve the demonstrated performance in the source environment (gravity x = 0.0, friction = 0.5) 4) Plots 9f 9g show evaluation of the proposed method ADAIL with policy conditioned on ground truth physics parameters and predicted physics parameters respectively; ADAIL matches the expert performance in the source environment (gravity x = 0.0, friction = 0.5) and generalizes to unseen dynamics. In particular, when the environment dynamics favors the task, the adaptive agent was able to obtain even higher performance (around friction = 1.2, gravity = 2).

Ant and Hopper. We again show favorable performance on both Ant and Hopper in Figure 6 .

How does the algorithm generalize to unseen environments? To understand how ADAIL generalizes to environments not sampled at training time, we do a suite of studies in which the agent is only allowed to interact in a limited set of environments. Figure 7 shows the performance of ADAIL on different settings, where a 5×5 region of environment parameters including the expert source environment are "blacked-out". This case is particularly challenging since the policy is not allowed access the domain from which the expert demonstrations were collected, and so our dynamics-invariant discriminator is essential. For additional held out experiments see Appendix A.5.

The experiments show that, 1) without training on the source environment, ADAIL with the ground truth parameters tends to have performance drops on the blackout region but largely is able to generalize (Figure 7a ); 2) the posterior's RMSE raises on the blackout region (Figure 7c ); 3) consequently ADAIL with the predicted dynamics parameters suffers from the posterior error on the blackout region (Figure 7b ).

How does unsupervised version of the algorithm perform?

VAE-ADAIL on HalfCheetah. With the goal of understanding the characteristics of the learned dynamics latent embedding through the unsupervised method and its impact on the overall algorithm, as a proof of concept we apply VAE-ADAIL to HalfCheetah environment varying a 1D continuous dynamics, friction. The performance is shown in Figure 8 . 

5. CONCLUSION:

In this work we proposed the ADaptive Adversarial Imitation Learning (ADAIL) algorithm for learning adaptive control policies from a limited number of expert demonstrations. We demonstrated the effectiveness of ADAIL on two challenging MuJoCo test suites and compared against recent SoTA. We showed that ADAIL extends the generalization capacities of policies to unseen environments, and we proposed a variant of our algorithm, VAE-ADAIL, that does not require environment dynamics labels at training time. We will release the code to aid in reproduction upon publication. σ 1 is a positive constant (6)

EQUATION

∇ log P (c|τ ) =∇(log 

EQUATION

:



Figures and Tables:

environment class E. 3: Initial parameters of policy θ, discriminator ω, and posterior φ. 4: A set of expert demonstrations {τexp} on one of the environment eexp ∈ E. An environment generator g(c) that takes a code and generates an environment ec ∈ E. A prior distribution of p(c). 5: for i = 1, 2, .. do 6: Sample c ∼ p(c) and generate environment ec ∼ g(c) 7: Sample trajectories τi ∼ π θ (•|c) in ec and τ e i ∼ {τexp} 8: Update the discriminator parameters ω with the gradients: Ê(s,a)∼τ i [∇w log(Dw(s, a))] + Ê(s,a)∼τ e i [∇w log(1 -Dw(s, a)] 9:Update the discriminator parameters ω again with the following loss, such that the gradients are reversed when back-prop through the dynamics-invariant layer: -Ê(s,a)∼τ i [log(D R (c|s, a))] 10:Update the posterior parameters φ with gradients Êτ i [∇ φ log Q φ (c|τi))] 11:Update policy π θ (•|c) using policy optimization method (PPO) with: Ê(s,a)∼τ i [-log(1 -Dω(s, a))] 12: Output: Learned policy π θ , and posterior Q φ .

Figure 3: VAE-based unsupervised dynamics latent variable learning.

Figure 4: Vary x-component of gravity in HalfCheetah environment. The red arrows in the picture show the gravity directions.

Figure 5: (a): ADAIL on CartPole-V0. Blue: PPO Expert; green: GAIL with Dynamics Randomization; red: ADAIL with latent parameters from the dynamics posterior; light blue: ADAIL with uniformly random latent parameters. (b): GAIL with Dynamics Randomization without (left, 2024.89 ± 669.39) or with (right, 2453.63 ± 430.51) GRL on Hopper

Figure 6: Comparing ADAIL with baselines on Mujoco tasks. Each plot is a heatmap that demonstrates the performance of an algorithm in environments with different dynamics. Each cell of the plot shows 10 episodes averaged cumulative rewards on a particular 2D range of dynamics. Note that to aid visualization, we render plots for Ant in log scale.

Figure 7: Generalization of our policy to held out parameters on the HalfCheetah environment. The red rectangles in plots show the blackout regions not seen during policy training.

Figure 8: VAE-ADAIL performance on HalfCheetah

LEMMA 1 PROOF Proof. For |c -Q φ (τ )| < δ,

Figure9: Comparing ADAIL with a few baselines on HalfCheetah. Each plot is a heatmap that demonstrates the performance of an algorithm in environments with different dynamics. Each cell of the plot shows 10 episodes averaged cumulative rewards on a particular 2D range of dynamics.

Peer review (iclr2020):
Summary: The submission considers the problem of imitation learning when the dynamics of the expert are not known to the agent and the dynamics of the agent may change frequently. It is however assumed that the agent has access to a parameterized simulator that can simulate the expert dynamics. The parameters for the simulator are not known but are assumed to be drawn from a known distribution. The proposed method is based on GAIL but uses several modifications: - A contextual policy is trained that also takes the dynamics-parameters as additional input. At each iteration, a new environment is sampled for performing the policy-rollout. - A *posterior* prediction network is trained to maximize the likelihood of the parameters that were used for the different roll-outs. This network is used for the test case, where the true dynamics of the agent are not known. - The discriminator might use features of the state-action input that correlate with the corresponding dynamics. Classifying based on such features may be undesirable because the discriminator might no longer produce useful rewards. In order to address this problem, an additional head is added to the discriminator that outputs a prediction of the dynamic parameters. The prediction error is trained by backpropagation, however by flipping the sign of the gradient at the last shared layer, the features of the discriminator are optimized to be unsuited for predicting the dynamic parameters (the technique is known as Gradient Reversal Layer). - A VAE-based method for learning latent dynamic parameters is proposed, by training a conditional VAE to reconstruct the next state, where the current state and action are provided as context to the encoder and decoder. Contribution: One of the strong-points of the submission is the fact that it features several different, orthogonal contributions. I also think that the considered problem setting is relatively interesting. However, also when considering real applications such as robotics, I am not convinced that explicitly modeling the dynamic changes is necessary. Some existing imitation learning methods focus on a setting where the dynamics of the expert may differ from the agent, but the dynamics of agents do not change. This setting does not require dynamic-contextual policies and seems to be applicable to typical robot applications. Soundness: The different components of the proposed methods seem reasonable to me. They do not come with (and arguably do not require) new derivations but seem rather like pragmatic solutions for the encountered problem. The optimization problem (Eq.2) seems to be formulated slightly inaccurate, because the last term should in my opinion not depend on theta. If I understand correctly, the policy should not maximize the likelihood of the dynamics posterior. The contrastive regularization loss needs to be better motivated. A high KL in the first term may not necessary be bad, for example, if the confidence in the prediction of (s_0, a_0, s*_0) is lower compared to (s_1, a_1, s*_1). If a similar regularizer has been used in prior work, such work should be referenced. Otherwise, it needs to be motivated. Presentation/Clarity: The presentation of the work is arguably the main weakness of the paper. The submission does not seem polished. It contains a large amount of typos. Figure 2 is confusing and adds little compared to the text description. Also the structure could be improved. For example, the submission introduces the posterior loss and outlines the algorithm before describing the individual components. The paper uses some techniques such as conditional VAEs [1] or contextual policy search [2] are used but not described / referenced. Evaluation: I like that the different aspects of the proposed method are also evaluated individually. The ablations with respect to the adaptability and GRL are crucial. The evaluation of the performance could be improved. PPO and UP-True use the true reward function, so the only real competitor is a naive GAIL-baseline that uses randomized dynamics during training. I*m not aware of prior work that considers the exact same setting as the manuscript. However, one of the main arguments for inverse reinforcement learning is the claimed generalizability of a reward function as opposed to a policy. I see that learning a new policy after each change in the dynamics may be too costly in some settings. However, comparisons to methods such as AIRL that aim to learn reward functions that are robust to changes in the dynamics would be highly interesting. [1] Sohn Kihyuk, Honglak Lee, and Xinchen Yan. “Learning Structured Output Representation using Deep Conditional Generative Models.” Advances in Neural Information Processing Systems. 2015. [2] Marc Peter Deisenroth, Gerhard Neumann, and Jan Peters. A survey on policy search forrobotics.Foundations and Trends in Robotics, 2(1–2):328–373, 2013. Comparison: How about IRL, e.g. AIRL? What about state-only GAIL? Typos: Equations are not properly integrated into the sentences (missing punctuations) *domains, as oppose to one domain.* *Inspired by to GANs* *and generates a environment* Figure 5a (legend): *dyanmics* *that can generalized across* Algorithmbox: *A environment*, *and Generate environment* *is achieved through 1) allowing* *the policy is mainly concerned with the end-effector of the latent parameters* Question: What are the network architectures? According to line 10 of the algorithmbox only the current trajectory is used for updating the dynamics posterior. Why?