Original: A UTO -E NCODING
Revised: A UTO -E NCODING

Original: I NVERSE R EINFORCEMENT
Revised: I NVERSE R EINFORCEMENT

Original: L EARNING
Revised: L EARNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Reinforcement learning (RL) provides a powerful framework for decisionmaking, but its application in practice often requires a carefully designed reward function.
Revised: Reinforcement learning (RL) provides a powerful framework for decisionmaking, but its application in practice often requires a carefully designed reward function.

Original: Inverse Reinforcement Learning (IRL) has shed light on automatic reward acquisition, but it is still difﬁcult to apply IRL to solve real-world tasks.
Revised: Inverse Reinforcement Learning (IRL) has shed light on automatic reward acquisition, but it is still difﬁcult to apply IRL to solve real-world tasks.

Original: In this work, we propose Auto-Encoding Inverse Reinforcement Learning (AEIRL), a robust and scalable IRL framework, which belongs to the adversarial imitation learning class.
Revised: In this work, we propose Auto-Encoding Inverse Reinforcement Learning (AEIRL), a robust and scalable IRL framework, which belongs to the adversarial imitation learning class.

Original: To recover reward functions from expert demonstrations, AEIRL utilizes the reconstruction error  as the learning signal, which provides more information for optimizing policies, compared to the binary logistic loss.
Revised: To recover reward functions from expert demonstrations, AEIRL utilizes the reconstruction error of an auto-encoder as the learning signal, which provides more information for optimizing policies, compared to the binary logistic loss.

Original: Subsequently, we use the derivedsafe objective functions to train the reward function and the RL agent.
Revised: Subsequently, we use the derived objective functions to train the reward function and the RL agent.

Original: In the experiments, we show that AEIRL performs superior in comparison with state-of-the-art methods in the MuJoCo environments.
Revised: Experiments show that AEIRL performs superior in comparison with state-of-the-art methods in the MuJoCo environments.

Original: More importantly, inthe more realistic settings, AEIRL shows much better robustnesscompared to other methods, when the expert demonstrations are noisy.
Revised: More importantly, in more realistic settings, AEIRL shows much better robustness when the expert demonstrations are noisy.

Revised: Speciﬁcally, our method achieves 16% relative improvement compared to the best baseline FAIRL on clean expert data and 38% relative improvement compared to the best baseline PWIL on noisy expert data both with the metric overall averaged scaled rewards.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Reinforcement learning (RL) provides a powerful framework for automating decision making.
Revised: Reinforcement learning (RL) provides a powerful framework for automating decision making.

Original: However, RL still requires signiﬁcantly engineered reward functions for good practical performance.
Revised: However, RL still requires signiﬁcantly engineered reward functions for good practical performance.

Original: To make RL more applicable in the real-world, it is important to learn a reward function from expert demonstrations.
Revised: To make RL more applicable in the real-world, it is important to learn a reward function from expert demonstrations.

Original: Imitation learning offers the instruments to learning policies directly from the data, without an explicit reward function.
Revised: Imitation learning offers the instruments to learn policies directly from the data, without an explicit reward function.

Original: Imitation learning enables the agents to learn to solve tasks from expert demonstrations, such as helicopter control (Abbeel et al., 2006; 2007; Ng et al., 2004;
Revised: Imitation learning enables the agents to learn to solve tasks from expert demonstrations, such as helicopter control (Abbeel et al., 2006; 2007; Ng et al., 2004;

Original: Coates et al., 2008; Abbeel et al., 2008a; 2010), robot navigation (Ratliff et al., 2006; Abbeel et al., 2008b; Ziebart et al., 2008; 2010), and building controls (Barrett & Linder, 2015).
Revised: Coates et al., 2008; Abbeel et al., 2008a; 2010), robot navigation (Ratliff et al., 2006; Abbeel et al., 2008b; Ziebart et al., 2008; 2010), and building controls (Barrett & Linder, 2015).

Original: The goal of imitation learning is to extract the expert policies from the expert demonstrations without the access to the reward signal from the environment.
Revised: The goal of imitation learning is to extract the expert policies from the expert demonstrations without the access to the reward signal from the environment.

Original: The algorithms in this ﬁeld can be divided into two broad categories: behavioral cloning (BC) and Inverse Reinforcement Learning (IRL).
Revised: The algorithms in this ﬁeld can be divided into two broad categories: behavioral cloning (BC) and Inverse Reinforcement Learning (IRL).

Original: BC formulates the learning task as a supervised learning problem, which learns a mapping from the states to the actions, using the expert trajectories.
Revised: BC formulates the learning task as a supervised learning problem, which learns a mapping from the states to the actions, using the expert trajectories.

Original: However, BC methods suffer from the problem of compounding errors, i.e., covariate shift, which only learn the actions of the expert but not reason about what the expert attempts to achieve.
Revised: However, BC methods suffer from the problem of compounding errors, i.e., covariate shift, which only learn the actions of the expert but not reason about what the expert attempts to achieve.

Original: In the contrast, IRL directly infers the underlying reward function from the data, and then teaches the RL agent to learn a policy to achieve the highest accumulated reward.
Revised: In the contrast, IRL directly infers the underlying reward function from the data, and then teaches the RL agent to learn a policy to achieve the highest accumulated reward.

Original: Empirical results show that IRL methods are more efﬁcient than BC methods in multi-step decision making tasks.
Revised: Empirical results show that IRL methods are more efﬁcient than BC methods in multi-step decision making tasks.

Original: However, how to make IRL stable and efﬁcient to use is still subject to research.
Revised: However, how to make IRL stable and efﬁcient to use is still subject to research.

Original: Current IRL methods, such as adversarial imitation learning approaches, model the reward function as a neural network to learn the mapping from the state-action pair to a scalar value, i.e., the reward (Ho & Ermon, 2016; Fu et al., 2017; Ghasemipour et al., 2020).
Revised: Current IRL methods, such as adversarial imitation learning approaches, model the reward function as a discriminator to learn the mapping from the state-action pair to a scalar value, i.e., the reward (Ho & Ermon, 2016; Fu et al., 2017; Ghasemipour et al., 2020).

Original: However, this formulation focuses on the  state-action featuresrather than the distribution of state-action visitations.
Revised: However, this formulation may be easily over-conﬁdent to focus on the minor differences between the state-action features of the expert and generated samples.

Original: We propose to formulate the reward function asthe reconstruction error of an auto-encoder, which learns to reconstruct the  expertdemonstrations. In the adversarial training framework, we design several loss functions and the corresponding reward functions for stable optimizing the auto-encoder and the policy.
Revised: To the best of our knowledge, this is the ﬁrst instance of formulating the reward function as an auto-encoder, which has the ability to learn full scale differences between expert and alternative policies.

Original: The derived formulations ensure the reward signal is informative and bounded, which are efﬁcientand safe for optimization.
Revised: The derived formulation ensures the reward signal is informative, which are efﬁcient for optimization.

Original: Additionally, The encoding-decoding process empowers the  capabilityto denoise and make the agent robust to the noisy expert, inthe more realistic settings.
Revised: Additionally, the encoding-decoding process empowers the denoising capability and make the agent robust to the noisy expert, in more realistic settings.

Original: In the experiments, we show that the proposed method AEIRL achieves the best overall performance on both clean and noisy settings.
Revised: In the experiments, we show that the proposed method AEIRL achieves the best overall performance on both clean and noisy expert demonstrations.

Original: Our contributions are two-folds:
Revised: Our contributions are three-fold:

Original: • We propose the Auto-Encoding Inverse Reinforcement Learning (AEIRL) architecture, which models the reward function as a surrogate function using the reconstruction error of the state-action features.
Revised: • We propose the Auto-Encoding Inverse Reinforcement Learning (AEIRL) architecture, which models the reward function as a surrogate function using the reconstruction error of an auto-encoder.

Original: The reconstruction error provides more informative learning signals, compared to the binary logistic loss.
Revised: The reconstruction error provides more informative learning signals, compared to the binary logistic loss.

Revised: • To show the contributing factors of our method, we conduct ablation studies based on different distribution divergences and alternative auto-encoders.

Revised: Experiments show that they achieve comparable results which indicates the encoding-decoding process is the major contributing factor for our method.

Original: • The experimental results on the MoJoCo tasks show that AEIRL outperforms state-of-theart imitation learning methods.
Revised: • The experimental results on the MoJoCo tasks show that our method outperforms state-ofthe-art imitation learning methods on both clean and noisy expert demonstrations.

Original: Furthermore, the empirical results show that  the learning processes of our methods are  more stable in general.
Revised: Empirical analysis show that our learned reward function can be more informative and robust. Furthermore, the learning processes of our methods are also more stable in general.

Original:  R ELATED W ORK AND B ACKGROUND
Revised:  R ELATED W ORK 2. I NVERSE R EINFORCEMENT L EARNING

Original: IRL infers the reward function using the expert demonstrations, which are assumed to be the observations of optimal behaviors (Ng & Russell, 2000). In the past two decades, the related works can be categorized into two directions. The ﬁrst direction is to enhance the representation ability of the reward functions, such as FIRL (Levine & Popvi ´ c, 2010), non-parametric IRL (Levine & Popvi ´ c, 2010), GP-IRL (Jin et al., 2015), and a class of methods that formulates the  reward functions using deep neural networks, including GCL, GAIL, AIRL, f -GAIL
Revised: Adversarial imitation learning such as GAIL, AIRL, DAC, f -GAIL, EAIRL, FAIRL (Ho & Ermon, 2016; Fu et al., 2017; Kostrikov et al., 2019; Zhang et al., 2020; Qureshi et al., 2019; Ghasemipour et al., 2020) formulates the learned reward function as a discriminator that learns to differentiate expert transitions from non-expert ones.

Original: (Finn et al., 2016b; Ho & Ermon, 2016; Fu et al., 2017; Zhang et al., 2020).
Revised: Among these methods, GAIL (Ho & Ermon, 2016) considers the Jensen-Shannon divergence, while AIRL (Fu et al., 2017) considers the Kullback-Leibler (KL) divergence.

Original: The second direction is to extend the previous methods to the more complex settings, such as the high dimensional state space environments (Cai et al., 2021), dynamic environments (Finn et al., 2016b), and continuous control environments (Ho & Ermon, 2016). In general, IRL is formulated as a bi-level optimization process meaning that iteratively training the reward function and optimizing the policy. A series of works show that Energy-Based Models (EBMs) play an important role in IRL. The basic assumption of EBMs is that the expert trajectories lie in a low energy space.
Revised: DAC (Kostrikov et al., 2019) extends GAIL to the off-policy setting and signiﬁcantly improves the sample-efﬁciency of adversarial imitation learning.

Original: Thus, the probability distribution of these trajectories can be modeled as a Boltzmann distribution: [Equation] where the energy is given by the negative reward function − r θ , and Z is the partition function normalizing the probability. Maximum entropy IRL (MaxEnt IRL)
Revised: Furthermore, f -divergence is utilized in f -GAIL

Original: (Ziebart et al., 2008; 2010) maximizes the following log-likelihood:
Revised: (Zhang et al., 2020), which is considered more sample-efﬁcient.

Original: [Equation] However, estimating the partition function Z is computationally expensive. MatEnt IRL is applicable in a ﬁnite MDP. Guided Cost Learning (GCL)
Revised: Recently,

Original: (Finn et al., 2016b) uses a mixture of policy distributions to estimate the importance weights of the demonstrations, in order to adapt the method tomore complex environments. Finn et al. (2016a) address the connection between GCL and GAIL (Ho & Ermon, 2016)
Revised: FAIRL utilizes the forward KL divergence (Ghasemipour et al., 2020) and achieves better performance than AIRL (Fu et al., 2017) comprehensively, but it is still not robust enough. However, these methods rely heavily on a carefully tuned discriminator, which might easily overﬁt to the minor differences between the expert and the generated samples.

Original: i.e., the equivalent discriminators and the generator loss functionals, and further propose GAN-GCL. In practice, using full trajectories in the reward function optimization (GANGCL) could result high variance estimates. Adversarial Inverse Reinforcement Learning (AIRL) (Fu et al., 2017) addresses the connection between the reward function r θ ( s, a ) and the discriminator D ( s, a ) of adversarial imitation learning.
Revised: In comparison, our auto-encoder based reward function helps to learn the full scale differences between the expert and generated samples, which provides more informative reward signals. The robustness of adversarial imitation learning is also questionable even with a little noise in observations (Stadie et al., 2017; Berseth & Pal, 2020), actions, transition models (Gangwani & Peng, 2020; Christiano et al., 2016), expert demonstrations (Brown et al., 2019;

Original: Mathematically: [Equation] where π ( s, a ) is the current policy we need to optimize.
Revised: Shiarlis et al., 2016;

Original: Minimizing the JS divergence between the occupancy measure of expert policy and the current policy leads to the loss functions of the original GAN method (Goodfellow et al., 2014). Although AIRL has limited improvements upon GAIL, AIRL can learn disentangled rewards and transfer better among different environments.
Revised: Jing et al., 2020) and their combinations (Kim et al., 2020).

Original: Following the direction of AIRL, EAIRL (Qureshi et al., 2019) and FAIRL (Ghasemipour et al., 2020) focus [Equation] [Figure] [Equation]
Revised: Previous robust IRL methods require the demonstrations to be annotated with conﬁdence scores (Wu et al., 2019; Brown et al., 2019; Grollman & Billard, 2012), when the expert data is noisy.

Original: [Figure] [Equation] [Figure] [Equation]
Revised: However, these annotations are rather expensive.

Original: [Figure] [Equation] [Figure] on more meaningful reward functions.
Revised: Compared to this, our auto-encoder based reward function helps to denoise the expert data through the encoding-decoding process.

Original: Energy-Based Imitation Learning (EBIL) (Liu et al., 2021) realizes IRL by two stages. First, estimate and use the expert energy as the surrogate reward function(Saremi et al., 2018). Secondly, utilize the reward to learn the policy.
Revised: Our method AEIRL is relatively succinct and robust to noisy expert demonstrations and does not require any annotated data. Another category of IRL uses an ofﬂine similarity function to estimate the rewards (Boularias et al., 2011; Klein et al., 2013; Piot et al., 2016).

Original: This method does not need to train the  policy in an inner loop ofoptimizing the reward function.
Revised: The idea of these methods is still inducing the expert policy by minimizing the distance between the state action distributions of the expert and sampled trajectories.

Original: Thus, it has a good convergence rate. Strictly batch imitation learning (EDM) (Jarrett et al., 2020) proposes ofﬂine imitation learning with energy-based distribution matching.
Revised: To the best of our knowledge, the most powerful method in this category is Primal Wasserstein Imitation Learning (PWIL) (Dadashi et al., 2021), which utilizes the upper bound of its primal form as the optimization objective.

Original: It doesn’t need any interactions with the environment, but it also suffers from the problem of compounding errors. Energy Based GAN model (EBGAN) (Zhao et al., 2016) takes the inspiration from EBMs, which formulates the discriminator as an auto-encoder and achieves good image generation performance. The generator tries to minimize the energy, i.e., the reconstruction error of dataexamples. The energy function is optimized via a margin loss.
Revised: The advantage of these methods is that they are relatively more robust compared to adversarial imitation learning methods, when the expert data is noisy.

Original: However,EBGAN is hard to train due to the unbounded energy value. Another category of IRL is using a ﬁxed similarity function to calculate the rewards.
Revised: However, the performance of these methods heavily depends on the similarity measurement, and therefore, it varies greatly on different tasks.

Original: To the best of ourknowledge, the most powerful method in this category is Primal Wasserstein Imitation Learning (PWIL) (Dadashi et al., 2021), which utilizes the upper bound of its primal form as the optimization objective.
Revised: Compared to PWIL, our method achieves superior performance. 2. A UTO -E NCODING BASED GAN S

Original: our method AEIRL takes inspirations from EBGAN  and utilizes an auto-encoder as the reward function and derives the efﬁcientand safe reconstruction error-based objective functions.
Revised: Auto-encoders have been successfully applied to improve the training stability and modes capturing in GANs. Auto-encoding based GANs can be classiﬁed into three categories: (1) utilizing an auto-encoder as the discriminator such as energy-based GANs (Zhao et al., 2016) and boundaryequilibrium GANs (Berthelot et al., 2017); (2) using a denoising auto-encoder to derive an auxiliary loss for the generator (Warde-Farley & Bengio, 2017); (3) combining variational auto-encoder and GANs to generate both vivid and diverse samples by balancing the objective of reconstructing the training data and confusing the discriminator (Larsen et al., 2016). Our method AEIRL takes inspirations from EBGAN (Zhao et al., 2016) and utilizes an auto-encoder as the reward function and derives the efﬁcient reconstruction error based surrogate reward signal and its corresponding objective functions.

Original: Compared to the state-of-the-art methods, our method AEIRL achieves superior performance.
Revised: A Markov decision process (MDP) is a tuple ( S, A, T, γ, P, r ) . In this tuple, S is a state space; A is an action space; T is a probability matrix for state transitions; γ ∈ (0 , 1] is a discount factor; P is a initial-state transition distribution; and r : S × A → R is a reward function.

Original: Furthermore, the robustness of adversarial imitation learning methods is questionable, even with a little noise in the  expert demonstrations.
Revised: Additionally, we also deﬁne a stochastic policy π , which is a mapping from states to probability distributions over actions. IRL infers the reward function using the expert demonstrations, which are assumed to be the observations of optimal behaviors (Ng & Russell, 2000). In general, IRL is formulated as a bi-level optimization process meaning that iteratively training the reward function and optimizing the policy. Assume that we are given an expert policy π E , IRL (Ziebart et al., 2008; 2010) ﬁts a reward function from a family of functions R with the optimization problem:

Original: Recently, learning from imperfect or noisy demonstrations become a popular topic in imitation learning (Sasaki & Yamashina, 2021).
Revised: [Equation] Moreover, the expert policy π E will only be provided by a set of expert demonstrations, so the expected reward of π E is estimated by these trajectories. IRL looks for a reward function r ( s, a ) that assigns high values to the expert policy and low values to other policies. Therefore, it allows the expert policy to be found with a certain reinforcement learning procedure:

Original: Previous methods require the demonstrations to be annotated with conﬁdence scores (Wu et al., 2019).
Revised: [Equation] The RL process will induce the expert policy via maximizing the expected cumulative rewards.

Original: However, these annotations are rather expensive.
Revised: Meanwhile the entropy term can be optionally added to the reinforcement learning objective to encourage the exploration in policy searching. Typically, IRL models the reward function as a discriminator which can be easily overﬁt to the expert data. Also its training stability is questionable once confronted with even a little noise in the expert data.

Original: In comparison to this, our method AEIRL is relatively robust to  noisy expert demonstrations and does not require annotated data.
Revised: Therefore, we propose auto-encoding IRL to utilize an auto-encoder based reward function which achieve strong performance both on clean and noisy expert demonstrations.

Original:  A UTO -E NCODING I NVERSE R EINFORCEMENT L EARNING 3. O UR M ETHOD
Revised:  A UTO -E NCODING I NVERSE R EINFORCEMENT L EARNING 4. M OTIVATION

Revised: The reward function in GAIL is a discriminator, which attempts to assign high values to the regions near the expert demonstrations and assign low values to the other regions.

Revised: However, this form of reward function could be easily overﬁtting to the expert data.

Revised: Consider the CartPole balancing task, the state of the environment is a feature consisting of position, angle, angle’s rate, and cart velocity, while the action is moving left or right.

Revised: Here, we assume that all the expert states’ velocity is 2 for example.

Revised: When the generated states’ velocity is 1 . 9 and other dimensions of state action pairs are

Original: Generally, our approach is to minimize the distance between the  state action distribution of  the policy we seek to train π and the state action distribution of the expert demonstrations D E .
Revised: [Figure] the same as the expert’s, the discriminator of GAIL would still give a low reward on these generated state action pairs. However it may actually perform very well on the goal of mimicking the expert’s behaviors. In other words, the reward function in GAIL on this example could easily overﬁt to the minor differences between the expert and the sampled data, while missing the underlying goal of the expert.

Revised: In our paper, we propose an auto-encoder based reward function for inverse reinforcement learning.

Revised: It utilizes the reconstruction error of the state action pairs to yield an informative reward signal.

Revised: The reconstruction error based reward signal signiﬁcantly retains the information of state action pairs, rather than focusing on the minor differences.

Revised: Such a reward signal wouldn’t lead to overconﬁdence in distinguishing the expert and the generated samples.

Revised: Recall the CartPole balancing example, the mean square error between states’ velocity 1 . 9 and 2 is very small.

Revised: And it could still feed a good reward to the agent under this situation.

Revised: Thus, the reconstruction error based reward signal focuses on the full scale differences between the expert and generated state action pairs rather than the minor parts.

Revised: This yields a much more informative reward signal.

Revised: Experiments in section 5.3, ﬁgure 5, shows a more informative reward signal learned on Walker2d task of our method.

Revised: Furthermore, the training process of the policy becomes smoother, as shown in ﬁgure 10 (See in Appendix A.4).

Revised: [Table]

Revised: The expert demonstrations usually contains noise bacause we sample these human trajectories with sensors and other devices in the real world.

Revised: Therefore, adversarial imitation learning, such as GAIL, might be easily affected to learn the noisy expert behaviors, which are not the real intentions of the expert.

Revised: The reward function under this situation could overﬁt to the noisy features.

Original: {+Recall the CartPole balancing example, we assume that the expert states’ velocity is 2 but it tends to be 
+}
Revised: {+Recall the CartPole balancing example, we assume that the expert states’ velocity is 2 but it tends to be 
+}

Revised: + δ due to the noisy sampling process.

Revised: When the learned policy is good enough as the sampled states’ velocity, the discriminator in GAIL still takes it as a bad policy, and therefore, it loses its efﬁcacy, when learning from noisy demonstrations.

Revised: Figure 2 depicts the denoising process of an autoencoder.

Revised: When the auto-encoder is trained to minimize the averaged squared errors, the vector points approximately towards the nearest points on the manifold, since the auto-encoder estimates the center of mass of the clean points (Goodfellow et al., 2016).

Revised: Therefore, the reconstruction error of the auto-encoder can eliminate some effects when learning from noisy expert demonstrations and make the reward signal more robust.

Original: We aim to optimize theEarth - Mover distance i.e., Wasserstein-1 distance between two distributions P r and P g :
Revised: Recall the CartPole example, the reconstruction error of the auto-encoder is δ 2 which doesn’t play an important role in the rewards for the agent. 4. M ETHOD Our approach is to minimize the distance between the state action distribution of the policy π θ and the expert demonstrations D E . The distance we used in our method is the Wasserstein distance (Arjovsky et al., 2017):

Original: [Equation]
Revised: [Equation]

Original: whereΠ( P r , P g ) denotes the set of all joint distributions γ ( x, y ) whose marginals are respectively
Revised: where the reward function network’s parameters are denoted as w and the policy network’s parameters are represented as θ .

Original: P r and P g .
Revised: We choose a class of reward functions from neural network functions for computational convenience. Minimizing this distance is actually inducing the expert policy from expert demonstrations.

Original: Intuitively, γ ( x, y ) indicates how much mass must be transported from x to y in order to transform the distributions P r into the distribution P g .
Revised: Therefore, the optimization of the policy π θ and recovering period of the reward function r w ( s, a ) forms a bi-level optimization problem, which can be formally deﬁned as: [Equation]

Original: Therefore, this distance is  the cost of the optimal transport plan.
Revised: This leads to an adversarial formulation for inverse RL. The outer level minimization with respect to the policy leads to a learned policy which is close to the expert. The inner level maximization recovers a reward function which attributes higher values to regions close to the expert data, and penalizes all other regions.

Original: In our work, we propose that the reward function isstructured with an auto-encoder which is deﬁned as:
Revised: In traditional adversarial imitation learning methods, such as WGAIL (Arjovsky et al., 2017), the reward function is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: where x represents the state-action pairs.
Revised: where D w ( s, a ) is the output of the discriminator.

Original: To train the auto-encoder adversarially with the reinforcement learning step, we use the Wasserstein1 distance in Equation 4:
Revised: In our method, we use an auto-encoder based surrogate reward function, which is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: where AE is the reconstruction error of  the auto-encoder as deﬁned in Equation 5.
Revised: where AE is the reconstruction error of an auto-encoder: [Equation] Here, x represents the state-action pairs.

Original: The loss function can be quickly minimized with just updating the last or last several layers of the auto-encoder while using unbound mean square error directly. So here, we use the reciprocal form to train the auto-encoder rather than the MSE directly which may leads to unstable training of the auto-encoder.
Revised: It is a mean square error between the sampled state action pairs and the reconstructed samples. This form of the reward signal uses the reconstruction error of an auto-encoder to score the state action pairs in trajectories. It is a monotonically decreasing function over the reconstruction error of the auto-encoder.

Original: Figure 1 depicts the architecture for AEIRL.
Revised: Low reconstruction errors ensure high reward values for state action pairs and vice versa. Section 4.1 tells that this form of reward signal focuses more on the full scale differences between the expert and generated samples and can help to denoise the expert demonstrations.

Original: The reinforcement learning component is trained to roll out trajectories, which are thesequences of state-action pairs. The auto-encoder takes either expert or generated state-action pairs, and estimates the reconstruction error AE ( x ) accordingly.
Revised: Training the auto-encoder is an adversarial process considering the objective 4, which is minimizing the reconstruction error for the expert samples and meanwhile maximizing this error for generated samples.

Original: The output of the auto-encoder goes through an objective functional in order to shape the reward function which is deﬁned as:
Revised: So the auto-encoder based reward function training objective is:

Original: [Equation]
Revised: [Equation]

Original: Considered that when the auto-encoder maps the samples to a latent space (might be a linear embedding with some noises) andreconstructs to the sample space, the noises can be removed by this reconstruction period.
Revised: The auto-encoder learns to maximize the full scale differences between the expert and the generated samples with the adversarial objective.

Original: So our method can achieve great performance even with some noisy expert.
Revised: It leads to a better feedback signal to the agent.

Original: Require: Initial parameters of policy, auto-encoder θ 0 , w 0 ; Expert trajectories D E . Ensure: i , a i ) ∼ π
Revised: Furthermore, the auto-encoder also retains more information and helps to denoise the expert data via the encoding-decoding process.

Original: θ
Revised: This can further help us to match the distribution between the expert and generated behaviors in a latent space.

Original: i and ( s E , a E ) ∼ D E with same batch size.
Revised: The variational auto-encoder (Kingma & Welling, 2014) is also an alternative form of reward function.

Original: Update w i to w i +1 by decreasing with the gradient:
Revised: When we use an variational auto-encoder based reward function, the objective function for training the variational auto-encoder in our method is:

Original: [Equation]
Revised: [Equation]

Original: Take a policy step from θ i to θ i +1 , using the TRPO update rule with the reward function 1 / (1 + AE w i ( s, a )) , and the objective function for TRPO is:
Revised: p

Original: [Equation]
Revised: [Equation]

Original: 5: end for Algorithm 1 depicts the procedure for training AEIRL.
Revised: [Equation]

Original: The ﬁrst step is to sample thestate-action pairs from expert demonstrationsD E and the trajectories rolled out by current policy with the same batch size. Then we update the auto-encoder by decreasing with the gradient with loss function 9.
Revised: where p model ( z )) is a ﬁxed Gaussian distribution. When minimizing this objective, the variational auto-encoder based reward function still assigns high values near the expert demonstrations and low values to other regions.

Original: The Third step is to update the policy assuming that the reward function is equation 7 which leads to a adversarial training paradigm. We runs these steps with N iterations until the policy converges. 3. A LTERNATIVE S URROGATE R EWARD S IGNALS AND A UTO -E NCODERS As mentioned in section 3.1, we design the reward function to avoid auto-encoder ”cheating” problem with bounded loss functions.
Revised: Meanwhile the latent distribution for the expert samples is widely distributed to induce more diverse expert policies.

Original: In detail, this problem is that once the loss function to train the auto-encoder is unbounded in adversarial training, the loss may be easily minimized with just updating the last layer or part of the nets. So we can use some other distribution distances to build a bounded reward signals without loss of generality. Here we propose another two reward signal formulations while following Jensen–Shannon distance to build the adversarial training paradigm.
Revised: Figure 1 depicts the architecture for AEIRL.

Original: One rewardformulation is: [Equation] where the corresponding loss function for updating the auto-encoder is: [Equation]
Revised: The auto-encoder based reward function takes either expert or the generated state action pairs, and estimates the reconstruction error based rewards accordingly.

Original: Another reward formulation is: [Equation] where the corresponding loss function for updating the auto-encoder is: [Equation]
Revised: The auto-encoder and the policy is iteratively optimized under this adversarial training paradigm.

Original: This two forms of reward signals achieve comparable performance on almost all the tasks shown in Appendix 8. While the ordinary auto-encoder is used in our paper, variational auto-encoder could also be an alternative with just restrict the latent space distribution to a ﬁxed Gaussian.  T HEORETICAL A NALYSIS
Revised: Algorithm 1 in Appendix A.1 depicts the pseudo code for training AEIRL.

Original: The reward function in our method can be represented as r ( x, ˜ x w ) , where w means the parameters for the auto-encoder and x represents the state action pair, ˜ x w is the reconstructed pair. Let π θ denote an agent’s stochastic policy with parameter θ . And ρ π ( x ) denotes the probability density of visiting some state action pairs x , which is also known as occupancy measure (Syed et al., 2008). The optimization of the policy π θ and recovering period of the reward function r ( x, ˜ x w ) forms a bi-level optimization problem, which can be formally deﬁned as:
Revised: Not limited to the utilized Wasserstein distance, other distribution divergences are also applicable to form the reward functions without loss of generality.

Original: [Equation] where W and Θ denote the parameter space of the auto-encoder and the policy respectively, H denote the entropy of policy. Obviously, the lower bound of the inner optimization problem is: [Equation]
Revised: Appendix A.4 shows other two forms of objectives and reward functions following Jensen-Shannon divergence.

Original: where we can optimize this lower bound instead of the original one empirically. In this section, we want to show some analysis on the question whether the reward function could converge or not theoretically. Assumption 1. There is no region where ρ π ∗ ( x ) = 0 where π ∗ is the induced optimal policy.
Revised: Experiments show that different forms of objectives achieve comparable results on almost all locomotion tasks.

Original: Theorem 1 (Main Results) . Suppose Assumption 1 holds, the reward function of our method could converge to a constant η in [0, 1] once the game for bi-level optimization problem reaches a Nash equilibrium. Proof. See Appendix A.1.
Revised: It indicates that an auto-encoder based reward function is the major contributing factor to improve the performance on mimicking the expert rather than the different distribution divergences and the corresponding forms of surrogate reward function.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We conduct the experiments on six locomotion tasks of varying dynamics and difﬁculty:
Revised: We conduct the experiments on six locomotion tasks with varying dynamics and difﬁculty:

Original: Swimmerv2 (Schulman et al., 2015a), Hopper-v2 (Levine & Koltun, 2013), Walker2d-v2 (Schulman et al., 2015a), HalfCheetah-v2 (Heess et al., 2015), Ant-v2 (Schulman et al., 2016), Humanoid-v2 (Tassa et al., 2012).
Revised: Swimmer-v2 (Schulman et al., 2015), Hopper-v2 (Levine & Koltun, 2013), Walker2d-v2 (Schulman et al., 2015), HalfCheetah-v2 (Heess et al., 2015), Ant-v2 (Schulman et al., 2016), and Humanoid-v (Tassa et al., 2012).

Original: The goal for all these tasks is to move forward as quickly as possible.
Revised: The goal for all these tasks is to move forward as quickly as possible.

Original: These tasks are more challenging than the basic tasks due to high degrees of freedom.

Original: In addition, a great amount of exploration is needed to learn to move forward without getting stuck at local optima.

Original: The baselines we choose for comparison includes: behavior cloning (BC) (Dhariwal et al., 2017), GAIL (Ho & Ermon, 2016), Wasserstein GAIL (WGAIL) (Arjovsky et al., 2017),  and PWIL (Dadashi et al., 2021).
Revised: The baselines we choose for comparison includes: behavior cloning (BC) (Dhariwal et al., 2017), GAIL (Ho & Ermon, 2016), Wasserstein GAIL (WGAIL) (Arjovsky et al., 2017), FAIRL (Ghasemipour et al., 2020), and PWIL (Dadashi et al., 2021).

Original: For  fair comparison, we use TRPO (Schulman et al., 2015b) as the policy search method for all the imitation learning algorithms which is implemented byOpenAI Baselines (Dhariwal et al., 2017). The expert demonstrations are sampled with PPO (Schulman et al., 2017) for 5 tasks except for Humanoid task. And the expert data for Humanoid is sampled with SAC (Haarnoja et al., 2018) to get higher expert performance. To our knowledge, PWIL achieves state-of-the-art performance, which outperforms DiscriminatorActor-Critic (Kostrikov et al., 2019) on almost all MuJoCo tasks (Dadashi et al., 2021), especially on Humanoid that claimed to be best performance so far.
Revised: For a fair comparison, we use TRPO (Schulman et al., 2015) as the policy search method for all the algorithms, which is implemented by

Original: We used all the algorithms to train policies of the same neural network architecture for all tasks: two hidden layers of 64 units each, with tanh nonlinearities in between. All networks were always initialized randomly at the start of each trial. For each task, we optimize BC with 100k iterations and gave GAIL, WGAIL, PWIL and our AEIRL exactly the same amount of environment interactions 10M for training. We ran all algorithms 5 times over different random seeds in all environments.
Revised: OpenAI Baselines (Dhariwal et al., 2017).

Original: Detailed experiments including  implementationsettings and more results on toy experiments and MuJoCo tasks can be seen in Appendix.
Revised: Experimental details, including the implementation and more results, are shown in Appendix A.2 to A.

Original: Note that the reward function tags used in experiments for our methods are shown in Table 4.

Original: [Table]
Revised: [Table]

Original: 5. L EARNING FROM E XPERT D EMONSTRATIONS
Revised: 5. L EARNING FROM E XPERT D EMONSTRATIONS

Revised: [Table]

Original:  Figure 2 depicts the results which are the training curve of original rewards for differentimitation learning algorithms.
Revised: To evaluate the performance of our proposed methods and other baselines, we run the experiments on the clean expert demonstrations. The metric we used is the ground truth rewards for different learned policies. Therefore, higher rewards indicate better mimicking expert behaviors. Figure 3 depicts the results, which are the training curve of ground truth rewards for different algorithms.

Original: The training curve shows the mean and standard deviation return of the evaluation policy over 10 rollouts and 5 seeds learning from non-noisy expert data.
Revised: The training curve shows the mean and standard deviation return of the evaluation policy over 10 rollouts and 5 seeds learning from non-noisy expert data.

Original: Table 2 shows the ﬁnal learned policy performance on MuJoCo tasks with raw episode return.
Revised: Table 1 shows the ﬁnal learned policy performance on MuJoCo tasks with ground truth episode return.

Revised: GAIL, WGAIL and FAIRL have comparable results, while WGAIL and FAIRL are little better than

Revised: GAIL comprehensively on these tasks but still not robust enough.

Revised: PWIL achieves higher performance compared with GAIL, WGAIL and FAIRL on Walker2d, Hopper and Humanoid tasks but

Original: [Table]
Revised: [Table]

Original: Obviously
Revised: performs worse than BC on Swimmer. Thus, the performance of PWIL varies greatly on different tasks. We take the random and expert trajectory as 0 and 1 respectively to compute the scaled ﬁnal rewards for different methods. The overall averaged scaled rewards for our method is about 0 .

Original: [Table]
Revised: 907 where the best baseline is 0 . 78 for FAIRL. It is a more than 16 .

Original:  our proposed AEIRL method achieves expert performane on these tasks.
Revised: 1% relative improvement. Clearly, our proposed AEIRL methods achieve expert performance.

Original: AEIRL outperforms other SoTA methods onalmost allthe locomotion tasks except for Hopper where PWIL achieves best performance on this task.
Revised: In detail, our method outperforms state-of-the-art baselines on all locomotion tasks except for Hopper where our method doesn’t outperform PWIL.

Original: PWIL achieves higher performance compared with GAIL and WGAIL on Walker2d and Hopper tasks but performs even worse than BC on Swimmer. GAIL and WGAIL get similar results while WGAIL could be a little better than GAIL on these tasks. Table 2 shows the best top 3 results marked with different colors.
Revised: Moreover, the variational auto-encoder variant of our method gets the best scores on Walker2d and Hopper.

Original: Three different loss functions for our methods are applied and all achieves strong results and  get best performance compared with other methods except for on Hopper.
Revised: In summary, our methods, including vanilla auto-encoder based and variational autoencoder based variants, get better performance than other baselines comprehensively.

Original: 5. L EARNING FROM N OISY D EMONSTRATIONS
Revised: 5. L EARNING FROM N OISY D EMONSTRATIONS

Original: [Table]
Revised: To show the robustness of our proposed AEIRL, we run our method and baselines on noisy expert demonstrations.

Original: [Equation] Figure 3 depicts theoriginal rewards for different methods on noisy expert data with Gaussian noise (0 , 0 . 3) on every task except for  Humanoid which with (0 , 0 . 01)
Revised: The noise added to the expert data is a Gaussian noise (0 , 0 . 3) on every dimension of state action pairs on all tasks except for Humanoid. Since the dimension of state action pairs in Humanoid is much higher than other tasks, we choose (0 , 0 . 01)

Original: Gaussian noise.
Revised: Gaussian noises added to the dataset in order to lower the effects of noise component.

Original:  Table 3 shows the ﬁnal learned policy performance on MuJoCo tasks with raw episode return.
Revised: Figure 7 (See in Appendix A.3) depicts training curves for the ground truth rewards for different methods on noisy expert data while Table 2 shows the ﬁnal learned policy performance on MuJoCo tasks with ground truth episode return.

Original: These results show that our method outperforms other state-of-the-art algorithmsalmost on all tasks except for Hopper  which PWIL wins.
Revised: These results show that our method outperforms other state-of-the-art algorithms on all tasks except for Hopper and Swimmer, on which PWIL and GAIL wins respectively.

Original: And for  the top 3 results, This deﬁnitely shows that our 
Revised: Table 1 and 2 show that the noise component affects the performance of mimicking the expert behaviors. But our method could help to denoise the expert demonstrations and provide more robust reward signals for inducing the expert policies. The overall scaled rewards for our AEIRL is 0 . 667 where the best baseline is 0 .for

Original: AEIRL can effectively denosie the expert data with the reconstruction period.
Revised: PWIL. Additionally, PWIL also has some capabilities to denoise the expert as it is deﬁned as an ofﬂine reward function. And other adversarial imitation learning methods, includes GAIL, WGAIL and FAIRL, are very sensitive to the noisy expert.

Original: 5. M ORE I NFORMATIVE R EWARD F UNCTION AND M ORE S TABLE T RAINING P ROCEDURE
Revised: 5. A NALYSIS ON R OBUSTNESS OF O UR M ETHOD

Original: We use t-SNE visulization to show the latent representations of the auto-encoder for our methods and the discriminator for baselines.
Revised: To compare the robustness of different reward signals, we use t-SNE visualization to show the latent representations of the auto-encoder of our method and the discriminator of the baselines.

Original: The auto-encoder or the discriminator is using two hidden layers of 100 units each, with tanh nonlinearities in between, and  the ﬁrst layer is using normalization and  the last layer is using sigmoid for GAIL and WGAIL.
Revised: More indistinguishable latent representations between policy samples and expert samples indicate better robustness of the reward function. Since PWIL considers an ofﬂine reward function without neural networks, we only compare our methods with GAIL, WGAIL and FAIRL. Figure 4 depicts that the latent representations of our method are more indistinguishable than other baselines both on clean and noisy expert data. It depicts the better robustness of our auto-encoder based reward function.

Original: So the latent representations is actually the output of  the second layer.
Revised: Different Objectives To analyze the major contributing factors of our method, we conduct the ablation studies based on different distribution divergences and its formulation of reward functions. Comparable performances would indicate that the major contributing factor is that we utilize the reconstruction error of an auto-encoder rather than the formulation of reward function and its corresponding objective. We formulate some other forms of surrogate reward functions based on JensenShannon divergence (Goodfellow et al., 2014) to justify this hypothesis.

Original: Figure 4 depicts that the latent representations are more indistinguishable between policy samples and expert samples for our methods than GAIL and WGAIL. And under the noisy expert, the latent representations of the samples becomes two clusters while our methods still learns the expert sample distributions.
Revised: Figure 8 (See in Appendix A.4) shows that three different formulations of surrogate reward signal achieve comparable results on almost all tasks which are all based on the vanilla auto-encoder. This shows the major contributing factor in our method is the auto-encoder based reward function which provides more informative learning signals, compared to the discriminator based reward function.

Revised: Alternative Auto-Encoders Our method is not limited to vanilla auto-encoders and we also take the variational auto-encoder as an alternative.

Revised: We run variational auto-encoder based variant on MuJoCo tasks under the same setting with vanilla auto-encoder.

Revised: Figure 9 (See in Appendix A.4) shows the training curves for variational auto-encoder based and vanilla auto-encoder based variants.

Revised: Overall, the scaled rewards of variational auto-encoder based variant is 0 . 845 and 0 .

Revised: 415 on clean and noisy expert respectively, which is comparable to the vanilla auto-encoder based method.

Revised: Auto-Encoder based Reward Signal is Denser A denser reward signal is beneﬁcial for policy searching (Ng et al., 1999;

Revised: Hu et al., 2020).

Revised: Section 4.1 tells that our auto-encoder based reward function focuses on the full scale differences between the expert and generated samples.

Revised: Thus, it provides more information to the agent compared to the discriminator based reward function.

Revised: To justify this hypothesis, we want to illustrate the generated rewards for different trajectories in the whole trajectory space.

Revised: Meanwhile more distinct rewards for different trajectories means the reward signal is denser and more informative.

Original: We train PPO on Walker2d to sample 50 different  levels of trajectories which can be seen as the trajectories in the whole trajectory space.
Revised: So we use the ﬁnal learned reward function of our method and the baselines to score different trajectories for analysis. We collect 50 levels of trajectories with different performances on Walker2d from random to the expert performance. This can be realized by saving trajectories every 100 iterations via training PPO on this task. Since the trajectories with different performances have different lengths, we normalize the generated rewards with the length of trajectories.

Revised: We also scale the generated rewards considering the expert samples are 1 and random policy samples are 0 .

Original:  Figure 5 shows that the ﬁnal recovered rewardsof the state action pairs for different levels of the  trajectories.
Revised: So the scaled normalized rewards for different levels of trajectories reﬂect the denseness of a reward signal. And higher degree of distinction indicates it is a denser reward signal. Figure 5 shows that the ﬁnal learned scaled and normalized rewards for different levels of the trajectories on Walker2d. The higher level in the abscissa means higher performance of its trajectories.

Original: And the curve depicts that  our method is more informative than GAIL and WGAIL with higher degree of distinction for normalized learned rewards.
Revised: Note that learned rewards under our formulation are readily distinct at different levels when compared to learned rewards under WGAIL. Our much steeper curves indicates that the auto-encoder derived reward function provides a denser and more informative learning signal when compared to the discriminator based reward functions. Stability of Our Method The variance of gradients is computed via different mini-batches of training examples. Lower variance of gradients means smoother updating of the network (Faghri et al., 2020).

Revised: Now that our auto-encoder based reward function is more informative and denser than the discriminator based ones, how about the training stability of our method?

Original: To show the training stability of our AEIRL, we record the norm of weights for policy network and the variance of policy gradients.
Revised: To show the training stability of different methods, we record the norm of weights of the policy network and the variance of policy gradients.

Original: Consider that the relative variance of policy gradients shows the training stability for different methods.
Revised: The relative variance of policy gradients shows the training stability for different methods.

Original: Figure 6 shows that our method get lower relative variance of gradients which means that the  network updating can besmoother and more stable.
Revised: Figure 10 (See in Appendix A.4) shows that our method gets lower relative variance of gradients which means that the policy network training process can be more stable.

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: In our paper, we present a simpliﬁed and robust adversarial imitation learning method based on AutoEncoding (AEIRL).
Revised: This paper presents a succinct and robust adversarial imitation learning method based on AutoEncoding (AEIRL).

Original: We structured the reward function with the reconstruction error of an autoencoder.
Revised: We utilize the reconstruction error of an auto-encoder as the surrogate reward function for reinforcement learning.

Revised: The advantages of our method can be summarized into two aspects: (1) the auto-encoder based reward function focuses on the full scale differences between the expert and generated samples, which won’t easily overﬁt to the expert data; (2) the auto-encoder denoises the expert data through the reconstruction process.

Original: The reward signal  is more robust and informative compared with  existing methods.
Revised: Thus, the learned reward signal in our method is more robust and informative compared with the existing methods.

Original: And further more, AEIRL could help denoise the expert demonstrations with the reconstruction period of training the reward function.

Original: Experiment results show that ourproposed methods achieve strong competitive performance on locomotion tasks even from imperfect (noisy) demonstrations.
Revised: Experimental results show that our methods achieve strong performances on locomotion tasks on both clean and noisy expert demonstrations.

Original: We hope that our work can further be applied into more realistic scenarios such as autonomous driving systems and so on.
Revised: In the future, we want to further investigate our method in more realistic scenarios, such as autonomous driving.

Original:  R EPRODUCIBILITY S TATEMENT
Revised:  R EPRODUCIBILITY S TATEMENT

Original: Code and data are attached in supplementary materials for reproducing the results, more details on experimental implementations are shown in Appendix A.2. A A PPENDIX A. P ROOF FOR T HEOREM Proof. The bi-level optimization problem for our method is discussed before:
Revised: Data and source code are attached in supplementary materials.

Original: [Equation]

Original: where W and Θ denote the parameter space of the auto-encoder and the policy respectively, H denote the entropy of policy.

Original: We deﬁne quantity V ( π, r ) as:

Original: [Equation]

Original: [Equation]

Original: and quantity Q ( π, r ) as:

Original: [Equation]

Original: and quantity U ( π, r ) as: [Equation] where quantity U is the lower bound of the quantity Q . In our experiments and proofs, we optimize
Revised: Experimental implementation details, including our methods and baselines, see in Appendix A.2.

Original: U as the reinforcement learning objective.

Original: Solving the bi-level optimization problem 14 is actually ﬁnding a Nash equilibrium between the two players in this game, policy π and reward function r .

Original: A Nash equilibrium of system 14 is a pair ( π ∗ , r ∗ ) which satisﬁes:

Original: [Equation]

Original: Due to that r ( x, ˜ x w ) = 1 / (1 + (cid:107) x − ˜ x w (cid:107) 22 ) , the reward is bounded in [0 , 1] .

Original: Now we prove that the reward function r would converge to a constant by proof of contradiction once the system reaches a Nash equilibrium.

Original: We assume that r is not constant, then there exists a constant C and a set S of non-zero measure such that ∀ x ∈ S , r ≤ C and ∀ x / ∈ S , r > C .

Original: Suppose Assumption 1 holds, then we have ρ π ∗ ( x ) > 0 on S . We can build a policy whose occupancy measure is ρ 0 such that ρ π 0 ( x ) < ρ π ∗ ( x ) over S . Then we compute [Equation]
Revised: A. I MPLEMENTATION

Original: Considered that ρ π ∗ ( x ) and ρ π 0 ( x ) are probability density functions, so we can hold that:

Original: [Equation]
Revised: [Figure]

Original: So by equation 19 and equation 20, we have: [Equation] So we can obtain that: [Equation]
Revised: ) (e) (f)

Original: Here we know the quantity ρ π ∗ ( x ) − ρ π 0 ( x ) ≥ 0 and r ∗ ( x, ˜ x w ) − C ≤ 0 for x ∈ S , and the quantity ρ π ∗ ( x ) − ρ π 0 ( x ) < 0 and r ∗ ( x, ˜ x w ) − C > 0 for x ∈ X \S .

Original: So for equation 22, i.e., U ( π ∗ , r ∗ ) − U ( π 0 , r ∗ ) < 0 , which violates equation 18.

Original: So the reward function of our method could converge to a constant η in [0, 1] when the game for bi-level optimization problem reaches a Nash equilibrium.

Original: This completes the proof.

Original: A. I MPLEMENTATION A.2. T ASKS

Original: [Figure]

Original: The goal for all these tasks is to move forward as quickly as possible.
Revised: The goal for all these tasks is to move forward as quickly as possible.

Original: These tasks are more challenging than the basic tasks due to high degrees of freedom.
Revised: These tasks are more challenging than the basic tasks due to high degrees of freedom.

Original: In addition, a great amount of exploration is needed to learn to move forward without getting stuck at local optima.
Revised: In addition, a great amount of exploration is needed to learn to move forward without getting stuck at local optima.

Original: Since we penalize for excessive controls as well as falling over, during the initial stage of learning, when the robot is not yet able to move forward for a sufﬁcient distance without falling, apparent local optima exist including staying at the origin or diving forward slowly.
Revised: Since we penalize for excessive controls as well as falling over, during the initial stage of learning, when the robot is not yet able to move forward for a sufﬁcient distance without falling, apparent local optima exist including staying at the origin or diving forward slowly.

Original: Figure 7 depicts locomotion tasks’ environments.
Revised: Figure 6 depicts locomotion tasks’ environments.

Original: A.2. B ASIC S ETTINGS
Revised: A.2. B ASIC S ETTINGS

Revised: The expert trajectories are sampled with PPO on Walker2d, Hopper, and Swimmer, HalfCheetah and Ant with 25 trajectories.

Revised: And we sample 320 trajectories on Humanoid with SAC to get higher expert performance since PPO is not good enough on Humanoid.

Original: On Walker2d, Hopper, and Swimmer, we train these algorithms directly, while on HalfCheetah, Ant and Humanoid we use BC to pre-train the policy with 10k iterations.
Revised: On Walker2d, Hopper, and Swimmer, we train these algorithms directly, while on HalfCheetah, Ant and Humanoid we use BC to pre-train the policy with 10k iterations.

Original: We test all the algorithms on these six locomotion tasks with environment seed 0 , 1 , 2 , 3 , 4 , and using 10 rollouts to estimates the original reward value to plot the training curves.
Revised: We test all the algorithms on these six locomotion tasks with environment seed 0 , 1 , 2 , 3 , 4 , and using 10 rollouts to estimates the ground truth reward value to plot the training curves.

Original: The maximum length for sampled trajectories is 1024.
Revised: The maximum length for sampled trajectories is 1024.

Original: The discounted factor is 0 .
Revised: The discounted factor is 0 .

Original: Each iteration we update 3 times of policies and 1 time of reward function (discriminator or auto-encoder).
Revised: Each iteration we update 3 times of policies and 1 time of reward function (discriminator or auto-encoder).

Original: While updating TRPO, we update the critic 5 times with learning rate 2e-4 every iteration.
Revised: While updating TRPO, we update the critic 5 times with learning rate 2e-4 every iteration.

Original: A.2. B ASELINES
Revised: A.2. B ASELINES

Original: We use open sourced code, OpenAI baselines (Dhariwal et al., 2017), to get the performance of BC and the training iterations is 100k for all tasks.
Revised: We use open sourced code, OpenAI baselines (Dhariwal et al., 2017), to get the performance of BC and the training iterations is 100k for all tasks.

Revised: GAIL (Ho & Ermon, 2016) borrows the GAN (Goodfellow et al., 2014) architecture to realize the inverse reinforcement learning process and it is applicable in high dimensional dynamic environments.

Original: We implement GAIL using open sourced code, OpenAI baselines (Dhariwal et al., 2017), with best tunned hyperparameters.
Revised: We implement GAIL using open source code, OpenAI baselines (Dhariwal et al., 2017), with best tunned hyperparameters.

Original: We choose TRPO as the reinforcement learning algorithm.
Revised: We choose TRPO as the reinforcement learning algorithm.

Revised: WGAIL (Arjovsky et al., 2017) utilizes the Wasserstein distance to train the policy and reward function which can be more stable.

Original: WGAIL is implemented based on GAIL with new loss function and reward signals as in (Arjovsky et al., 2017).
Revised: WGAIL is implemented based on GAIL with new loss function and reward signals as in (Arjovsky et al., 2017).

Original: We  choose TRPO as the reinforcement learning algorithm.
Revised: We also choose TRPO as the reinforcement learning algorithm.

Revised: FAIRL (Ghasemipour et al., 2020) utilizes the forward KL divergence as the objective which shows competitive performances on MuJoCo tasks.

Revised: FAIRL is implemented based on GAIL with new loss function and reward signals as in (Ghasemipour et al., 2020).

Revised: We choose TRPO as the reinforcement learning algorithm.

Revised: PWIL (Dadashi et al., 2021) is an adversarial imitation learning method based on an ofﬂine similarity reward function.

Revised: It is quite robust and easy to ﬁne-tune the hyper-parameters of its reward function (only two hyper-parameters).

Original: PWIL is implemented based on open sourced code (Dadashi et al., 2021) with TRPO as the reinforcement learning step for fair comparison (original PWIL code uses SAC as the RL algorithm).
Revised: PWIL is implemented based on open source code (Dadashi et al., 2021) with TRPO as the reinforcement learning step for fair comparison (original PWIL code uses SAC as the RL algorithm).

Original: A.2. O UR M ETHODS
Revised: A.2. O UR M ETHODS

Original: Policy Net: the same with baselines, two hidden layers with 64 units, with tanh nonlinearities in between.
Revised: Policy Net: the same with baselines, two hidden layers with 64 units, with tanh nonlinearities in between.

Original: The learning rate to update the policy is 0.01.
Revised: The learning rate to update the policy is 0.01.

Original: Auto-Encoder Net and VAE Net: the same number of layers with the architecture of discriminator in GAIL and WGAIL, 4 layers.
Revised: Auto-Encoder Net (AE & VAE): the same number of layers with the architecture of discriminator in GAIL and WGAIL, 4 layers.

Original: The two hideen layers are with 100 units, with tanh nonlinearities in between, the ﬁnal output layer is identity.
Revised: The two hideen layers are with 100 units, with tanh nonlinearities in between, the ﬁnal output layer is identity.

Original: And we normalize the state feature after the input layer and we use the raw state action features to compute the mean square error.
Revised: And we normalize the state feature after the input layer and we use the raw state action features to compute the mean square error.

Original: Learning rate for the auto-encoder is 3e-4.
Revised: Learning rate for the auto-encoder is 3e-4.

Original: We doesn’t add the entropy term to the reinforcement learning step but already get great performance.
Revised: We don’t add the entropy term to the reinforcement learning step and already get great performances.

Original: So  ﬁne-tuning the hyper-parameters might have better performance on these tasks.
Revised: So we believe ﬁne-tuning the hyper-parameters might have better performances on these tasks.

Original: A.2. T OY E XPERIMENTS
Revised: A. O MITTED E XPERIMENT R ESULTS

Original: We design a 1D walking task as toy experiments.
Revised: Figure 7 shows the training curves for different imitation learning algorithms on noisy expert demonstrations.

Original: The state space is discrete with 40 states and the action is 0 walking left one step or 1 walking right one step. The initial state is sampled by uniform distribution between state 10 to 29. The maximum length of trajectory is 10.
Revised: The noise is a Gaussian distribution noise which is (0 , 0 . 3) for Walker2d, Hopper, Swimmer, HalfCheetah, and Ant which (0 , 0 . 01) for Humanoid.

Original: We set the expert policy is always walking right and sample 100 trajectories. Every algorithms will be trained 3 times with random initialization of  the policy, reward function and critic network.
Revised: The noise component affects the performance of imitation learning differently. We can also observe that this affects becomes greater when the dimension for state action pairs becomes larger.

Original: The timesteps for training is set to 100k. The reward function is structured with 4 dense layers with relu nonlinerities in between. The policy net is structured with 4 dense layers with 20 units with tanh nonlinerities in between.
Revised: That’s why we choose a smaller noise component on Humanoid.

Original: We use PPO to train the policy as the reinforcement learning step.
Revised: In comparison, our method is more robust than other baselines.

Original: A. A BLATION S TUDIES
Revised: A. A BLATION S TUDIES

Original: Figure 8 depicts the comparison for different variants of our methods with different loss functions on non-noisy expert demonstrations.
Revised: In this section, we propose another two forms of reward functions following the Jensen–Shannon divergence (Goodfellow et al., 2014). The ﬁrst reward function form is: [Equation]

Original: Clearly, Our methods with different loss function and reward signal formulation have comparable results on Walker2d, Swimmer, HalfCheetah, Hopper, and Ant.
Revised: where the corresponding loss function for updating the auto-encoder is: [Equation] [Equation] In this formulation, we take the exponential of negative reconstruction error as a probability.

Original: But on Humanoid, the Wasserstein-1 distance achieves better performance than JS distance based loss functions.
Revised: Minimizing the Jensen-Shannon divergence between state action distribution of the expert and generated trajectories assigns high probabilities to the expert samples and low probabilities to other regions. In this situation, minimizing Jensen-Shannon divergence is equivalent to maximizing the cross entropy for these data examples. When we simplify the corresponding objective, we can obtain Eq. The second reward function form is:

Original: [Figure]
Revised: [Equation]

Original: Figure 9 depicts the training curve of loss function for  the auto-encoder which shows the convergence of the auto-encoder on non-noisy expert demonstrations. A. M ORE E XPERIMENTS ON N OISY E XPERT D ATA
Revised: where the corresponding loss function for updating the auto-encoder is:

Original: Figure 10 shows that the training curve for different methods on more noisy expert data. Clearly, our methods get higher performance than other baselines on almost all tasks except for Hopper which is not better than PWIL. With the greater noise, we can see that our methods can effectively denoise the expert data while GAIL and WGAIL cannot. Additionally, PWIL also has some capability to denoise the expert data with a ﬁxed reward function.
Revised: [Equation]

Original: Figure 11 shows that t-SNE visualization results for the latent space of reward functions on Walker2d and non-noisy and different noisy expertdata. The indistinguishable policy samples and  expert samples means better reward recovering.
Revised: In this formulation, we take the sigmoid output of the negative reconstruction error as a probability. Also minimizing the Jensen-Shannon divergence between state action distribution of the expert and generated trajectories assigns high probabilities to the expert samples and low probabilities to other regions.

Original: Clearly, when the noise becomes greater, GAIL and WGAIL fail to learn the reward distribution and two sets of samples become separate clusters. However, our methods still recover the reward function efﬁciently except for strong noisy expert. A. O NE S HOT I MITATION L EARNING ON N ON -N OISY E XPERT D ATA A. T OY E XPERIMENTS
Revised: And we can obtain this objective 13 with some simpliﬁcations.

Original: Note that the reward function tags used in experiments for our methods are shown in Table 
Revised: Note that the reward function tags used in experiments for our methods are shown in Table 3.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: These three forms of reward signals including our method achieve comparable performance on almost all the tasks except for Humanoid shown in Figure 8. This indicates that an auto-encoder based reward function is the major contributing factor to improve the performance on mimicking the expert rather than the different distribution divergences or the different forms of reward function.

Original: Figure 12 shows the KL divergence between the learned policy and the expert policy.
Revised: Figure 9 shows the training curves for vanilla auto-encoder based variant and varational auto-encoder based variant on non-noisy expert demonstrations.

Original: Our methods learn the policies quite close to the expert with the KL divergence close to 0. Figure 13 shows the heatmaps for recovered reward functions for the entire state action space.
Revised: The performances for these two variants are comparable on almost all tasks while vanilla auto-encoder based method is a little better than varational auto-encoder based variant on Humanoid. Figure 10 shows the training stability of our method compared to the baseline WGAIL.

Original: The reward function of our methods are more informative than other baselines as shown in the ﬁgure.
Revised: It depicts that our method can be more stable and the policy network updating is smoother.