Original: L EARNING R OBUST R EWARDS WITH A DVERSARIAL I NVERSE R EINFORCEMENT L EARNING
Revised: L EARNING R OBUST R EWARDS WITH A DVERSARIAL I NVERSE R EINFORCEMENT 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 provides a powerful and general framework for decision making and control, but its application in practice is often hindered by the need for extensive feature and reward engineering.
Revised: Reinforcement learning provides a powerful and general framework for decision making and control, but its application in practice is often hindered by the need for extensive feature and reward engineering.

Original: Deep reinforcement learning methods can remove the need for explicit engineering of policy or value features, but still require a manually speciﬁed reward function.
Revised: Deep reinforcement learning methods can remove the need for explicit engineering of policy or value features, but still require a manually speciﬁed reward function.

Original: Inverse reinforcement learning holds the promise of automatic reward acquisition, but has proven exceptionally difﬁcult to apply to large, high-dimensional problems with unknown dynamics.
Revised: Inverse reinforcement learning holds the promise of automatic reward acquisition, but has proven exceptionally difﬁcult to apply to large, high-dimensional problems with unknown dynamics.

Original: In this work, we propose AIRL, a practical and scalable inverse reinforcement learning algorithm based on an adversarial reward learning formulation.
Revised: In this work, we propose AIRL, a practical and scalable inverse reinforcement learning algorithm based on an adversarial reward learning formulation.

Original: We demonstrate that AIRL is able to recover reward functions that are robust to changes in dynamics, enabling us to learn policies even under signiﬁcant variation in the environment seen during training.
Revised: We demonstrate that AIRL is able to recover reward functions that are robust to changes in dynamics, enabling us to learn policies even under signiﬁcant variation in the environment seen during training.

Original: Our experiments show that AIRL greatly outperforms prior methods in these transfer settings.
Revised: Our experiments show that AIRL greatly outperforms prior methods in these transfer settings.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: While reinforcement learning (RL) provides a powerful framework for automating decision making and control, signiﬁcant engineering of elements such as features and reward functions has typically been required for good practical performance.
Revised: While reinforcement learning (RL) provides a powerful framework for automating decision making and control, signiﬁcant engineering of elements such as features and reward functions has typically been required for good practical performance.

Original: In recent years, deep reinforcement learning has alleviated the need for feature engineering for policies and value functions, and has shown promising results on a range of complex tasks, from vision-based robotic control (Levine et al., 2016) to video games such as Atari (Mnih et al., 2015) and
Revised: In recent years, deep reinforcement learning has alleviated the need for feature engineering for policies and value functions, and has shown promising results on a range of complex tasks, from vision-based robotic control (Levine et al., 2016) to video games such as Atari (Mnih et al., 2015) and

Original: Minecraft (Oh et al., 2016).
Revised: Minecraft (Oh et al., 2016).

Original: However, reward engineering remains a signiﬁcant barrier to applying reinforcement learning in practice.
Revised: However, reward engineering remains a signiﬁcant barrier to applying reinforcement learning in practice.

Original: In some domains, this may be difﬁcult to specify (for example, encouraging “socially acceptable” behavior), and in others, a na¨ıvely speciﬁed reward function can produce unintended behavior (Amodei et al., 2016).
Revised: In some domains, this may be difﬁcult to specify (for example, encouraging “socially acceptable” behavior), and in others, a na¨ıvely speciﬁed reward function can produce unintended behavior (Amodei et al., 2016).

Original: Moreover, deep RL algorithms are often sensitive to factors such as reward sparsity and magnitude, making well performing reward functions particularly difﬁcult to engineer.
Revised: Moreover, deep RL algorithms are often sensitive to factors such as reward sparsity and magnitude, making well performing reward functions particularly difﬁcult to engineer.

Original: Inverse reinforcement learning (IRL) (Russell, 1998; Ng & Russell, 2000) refers to the problem of inferring an expert’s reward function from demonstrations, which is a potential method for solving the problem of reward engineering.
Revised: Inverse reinforcement learning (IRL) (Russell, 1998; Ng & Russell, 2000) refers to the problem of inferring an expert’s reward function from demonstrations, which is a potential method for solving the problem of reward engineering.

Original: However, inverse reinforcement learning methods have generally been less efﬁcient than direct methods for learning from demonstration such as imitation learning (Ho & Ermon, 2016), and methods using powerful function approximators such as neural networks have required tricks such as domain-speciﬁc regularization and operate inefﬁciently over whole trajectories (Finn et al., 2016b).
Revised: However, inverse reinforcement learning methods have generally been less efﬁcient than direct methods for learning from demonstration such as imitation learning (Ho & Ermon, 2016), and methods using powerful function approximators such as neural networks have required tricks such as domain-speciﬁc regularization and operate inefﬁciently over whole trajectories (Finn et al., 2016b).

Original: There are many scenarios where IRL may be preferred over direct imitation learning, such as re-optimizing a reward in novel environments (Finn et al., 2017) or to infer an agent’s intentions, but IRL methods have not been shown to scale to the same complexity of tasks as direct imitation learning.
Revised: There are many scenarios where IRL may be preferred over direct imitation learning, such as re-optimizing a reward in novel environments (Finn et al., 2017) or to infer an agent’s intentions, but IRL methods have not been shown to scale to the same complexity of tasks as direct imitation learning.

Original: However, adversarial IRL methods (Finn et al., 2016b;a) hold promise for tackling difﬁcult tasks due to the ability to adapt training samples to improve learning efﬁciency.
Revised: However, adversarial IRL methods (Finn et al., 2016b;a) hold promise for tackling difﬁcult tasks due to the ability to adapt training samples to improve learning efﬁciency.

Original: Part of the challenge is that IRL is an ill-deﬁned problem, since there are many optimal policies that can explain a set of demonstrations, and many rewards that can explain an optimal policy (Ng et al., 1999).
Revised: Part of the challenge is that IRL is an ill-deﬁned problem, since there are many optimal policies that can explain a set of demonstrations, and many rewards that can explain an optimal policy (Ng et al., 1999).

Original: The maximum entropy (MaxEnt) IRL framework introduced by Ziebart et al.
Revised: The maximum entropy (MaxEnt) IRL framework introduced by Ziebart et al.

Original: (2008) handles the former ambiguity, but the latter ambiguity means that IRL algorithms have difﬁculty distinguishing the true reward functions from those shaped by the environment dynamics.
Revised: (2008) handles the former ambiguity, but the latter ambiguity means that IRL algorithms have difﬁculty distinguishing the true reward functions from those shaped by the environment dynamics.

Original: While shaped rewards can increase learning speed in the original training environment, when the reward is deployedin at test-time on environments with varying dynamics, it may no longer produce optimal behavior, as we discuss in Sec. 5.
Revised: While shaped rewards can increase learning speed in the original training environment, when the reward is deployed at test-time on environments with varying dynamics, it may no longer produce optimal behavior, as we discuss in Sec. 5.

Original: To address this issue, we discuss how to modify IRL algorithms to learn rewards that are invariant to changing dynamics, which we refer to as disentangled rewards .
Revised: To address this issue, we discuss how to modify IRL algorithms to learn rewards that are invariant to changing dynamics, which we refer to as disentangled rewards .

Original: In this paper, we propose adversarial inverse reinforcement learning (AIRL), an inverse reinforcement learning algorithm based on adversarial learning.
Revised: In this paper, we propose adversarial inverse reinforcement learning (AIRL), an inverse reinforcement learning algorithm based on adversarial learning.

Original: Our algorithm provides for simultaneous learning of the reward function and value function, which enables us to both make use of the efﬁcient adversarial formulation and recover a generalizable and portable reward function, in contrast to prior works that either do not recover a reward functions (Ho & Ermon, 2016), or operates at the level of entire trajectories, making it difﬁcult to apply to more complex problem settings (Finn et al., 2016b;a).
Revised: Our algorithm provides for simultaneous learning of the reward function and value function, which enables us to both make use of the efﬁcient adversarial formulation and recover a generalizable and portable reward function, in contrast to prior works that either do not recover a reward functions (Ho & Ermon, 2016), or operates at the level of entire trajectories, making it difﬁcult to apply to more complex problem settings (Finn et al., 2016b;a).

Original: Our experimental evaluation demonstrates that AIRL outperforms prior IRL methods (Finn et al., 2016b) on continuous, high-dimensional tasks with unknown dynamics by a wide margin.
Revised: Our experimental evaluation demonstrates that AIRL outperforms prior IRL methods (Finn et al., 2016b) on continuous, high-dimensional tasks with unknown dynamics by a wide margin.

Original: When compared to GAIL (Ho & Ermon, 2016), which does not attempt to directly recover rewards, our method achieves comparable results on tasks that do not require transfer.
Revised: When compared to GAIL (Ho & Ermon, 2016), which does not attempt to directly recover rewards, our method achieves comparable results on tasks that do not require transfer.

Original: However, on tasks where there is considerable variability in the environment from the demonstration setting,
Revised: However, on tasks where there is considerable variability in the environment from the demonstration setting,

Original: GAIL and other IRL methods fail to generalize.
Revised: GAIL and other IRL methods fail to generalize.

Original: In these settings, our approach, which can effectivelyrecover disentangle the goals of the expert from the dynamics of the environment, achieves superior results.
Revised: In these settings, our approach, which can effectively disentangle the goals of the expert from the dynamics of the environment, achieves superior results.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Inverse reinforcement learning (IRL) is a form of imitation learning and learning from demonstration (Argall et al., 2009).
Revised: Inverse reinforcement learning (IRL) is a form of imitation learning and learning from demonstration (Argall et al., 2009).

Original: Imitation learning methods seek to learn policies from expert demonstrations, and IRL methods accomplish this by ﬁrst inferring the expert’s reward function.
Revised: Imitation learning methods seek to learn policies from expert demonstrations, and IRL methods accomplish this by ﬁrst inferring the expert’s reward function.

Original: Previous IRL approaches have included maximum margin approaches (Abbeel & Ng, 2004; Ratliff et al., 2006), and probabilistic approaches such as Ziebart et al.
Revised: Previous IRL approaches have included maximum margin approaches (Abbeel & Ng, 2004; Ratliff et al., 2006), and probabilistic approaches such as Ziebart et al.

Original: ; Boularias et al.
Revised: ; Boularias et al.

Original: In this work, we work under the MaxEnt IRL framework of Ziebart et al.
Revised: In this work, we work under the maximum causal IRL framework of Ziebart (2010).

Original: Some advantages of MaxEnt IRL are that itallows for sub-optimality in demonstrations, and removes ambiguity between demonstrations and the expert policy. Moreover, it allows us to cast the reward learning problem as a maximum likelihood problem, allowing us to connect IRL to generative model training.
Revised: Some advantages of this framework are that it removes ambiguity between demonstrations and the expert policy, and allows us to cast the reward learning problem as a maximum likelihood problem, connecting IRL to generative model training.

Original: Our proposed method most closely resembles the algorithms proposed by Uchibe (2017);
Revised: Our proposed method most closely resembles the algorithms proposed by Uchibe (2017);

Original: Ho & Ermon (2016);
Revised: Ho & Ermon (2016);

Original: Finn et al.
Revised: Finn et al.

Original: (2016a).
Revised: (2016a).

Original: Generative adversarial imitation learning (GAIL) (Ho & Ermon, 2016) differs from our work in that it is not an IRL algorithm that seeks to recover reward functions.
Revised: Generative adversarial imitation learning (GAIL) (Ho & Ermon, 2016) differs from our work in that it is not an IRL algorithm that seeks to recover reward functions.

Original: The critic or discriminator of GAIL is unsuitable as a reward since, at optimality, it outputs 0.5 uniformly across all states and actions.
Revised: The critic or discriminator of GAIL is unsuitable as a reward since, at optimality, it outputs 0.5 uniformly across all states and actions.

Original: Instead, GAIL aims only to recover the expert’s policy, which is a less portable representation for transfer.
Revised: Instead, GAIL aims only to recover the expert’s policy, which is a less portable representation for transfer.

Original: Uchibe (2017) does not interleave policy optimization with reward learning within an adversarial framework, whereas prior work has shown that interleaving policy optimization and reward learning is crucial forgood performanceon tasks that are continuous or high-dimensional (Finn et al., 2016b)
Revised: Uchibe (2017) does not interleave policy optimization with reward learning within an adversarial framework. Improving a policy within an adversarial framework corresponds to training an amortized sampler for an energy-based model, and prior work has shown this is crucial for performance (Finn et al., 2016b).

Original: Wulfmeier et al.
Revised: Wulfmeier et al.

Original: (2015) also consider learning cost function with neural networks, but only evaluate on simple domains where analytically solving the problem with value iteration is tractable.
Revised: (2015) also consider learning cost functions with neural networks, but only evaluate on simple domains where analytically solving the problem with value iteration is tractable.

Original: Previous methods which aim to learn nonlinear cost function have used boosting (Ratliff et al., 2007) and gaussian processes (Levine et al., 2011), but still suffer from the feature engineering problem.
Revised: Previous methods which aim to learn nonlinear cost functions have used boosting (Ratliff et al., 2007) and Gaussian processes (Levine et al., 2011), but still suffer from the feature engineering problem.

Original: Our IRL algorithm builds on the adversarial IRL framework proposed by Finn et al.
Revised: Our IRL algorithm builds on the adversarial IRL framework proposed by Finn et al.

Original: (2016a), with the discriminator corresponding to an odds ratio between the policy and exponentiated reward distribution.
Revised: (2016a), with the discriminator corresponding to an odds ratio between the policy and exponentiated reward distribution.

Original: The discussion in Finn et al.
Revised: The discussion in Finn et al.

Original: (2016a) is theoretical, and to our knowledge no prior work has reported a practical implementation of this method.
Revised: (2016a) is theoretical, and to our knowledge no prior work has reported a practical implementation of this method.

Original: Our experiments show that direct implementation of the proposed algorithm is ineffective, due to high variance from operating over entire trajectories.
Revised: Our experiments show that direct implementation of the proposed algorithm is ineffective, due to high variance from operating over entire trajectories.

Original: While it is straightforward to extend the algorithm to single state-action pairs, as we discuss in Section 4, a simple unrestricted form of the discriminator is susceptible to the reward ambiguity described in (Ng et al., 1999), making learning the portable reward functions difﬁcult.
Revised: While it is straightforward to extend the algorithm to single state-action pairs, as we discuss in Section 4, a simple unrestricted form of the discriminator is susceptible to the reward ambiguity described in (Ng et al., 1999), making learning the portable reward functions difﬁcult.

Original: As illustrated in our experiments, this greatly limits the generalization capability of the method: the learned reward functions are not robust the environment changes, and it is difﬁcult to use the algorithm for the purpose of inferring the intentions of agents.
Revised: As illustrated in our experiments, this greatly limits the generalization capability of the method: the learned reward functions are not robust to environment changes, and it is difﬁcult to use the algorithm for the purpose of inferring the intentions of agents.

Original: We discuss how to overcome this issue in
Revised: We discuss how to overcome this issue in

Original: Section 5.
Revised: Section 5.

Revised: Amin et al.

Revised: (2017) consider learning reward functions which generalize to new tasks given multiple training tasks.

Revised: Our work instead focuses on how to achieve generalization within the standard IRL formulation.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Our inverse reinforcement learning method builds on the maximum  entropy IRL framework (Ziebart et al., 2008), which considers an entropy-regularized Markov decision process (MDP) (Ziebart, 2010), deﬁned by the tuple ( S , A , T , R, γ, ρ 0 ) .
Revised: Our inverse reinforcement learning method builds on the maximum causal entropy IRL framework (Ziebart, 2010), which considers an entropy-regularized Markov decision process (MDP), deﬁned by the tuple ( S , A , T , r, γ, ρ 0 ) .

Original: S ,
Revised: S ,

Original: A are the state and action spaces, respectively, γ ∈ (0 , 1) is the discount factor.
Revised: A are the state and action spaces, respectively, γ ∈ (0 , 1) is the discount factor.

Original: The dynamics or transition distribution T ( s (cid:48) | a, s ) , the initial state distribution ρ 0 ( s ) , and the reward function R ( s, a ) are unknown in the standard reinforcement learning setup and can only be queried through interaction with the MDP.
Revised: The dynamics or transition distribution T ( s (cid:48) | a, s ) , the initial state distribution ρ 0 ( s ) , and the reward function r ( s, a ) are unknown in the standard reinforcement learning setup and can only be queried through interaction with the MDP.

Original: The goal of (forward) reinforcement learning is to ﬁnd the optimal policy π ∗ that maximizes the expected entropy-regularized discounted reward, under π , T , and ρ 0 :
Revised: The goal of (forward) reinforcement learning is to ﬁnd the optimal policy π ∗ that maximizes the expected entropy-regularized discounted reward, under π , T , and ρ 0 :

Original: [Equation]
Revised: [Equation]

Original: where[-R ( τ ) = (cid:80) Tt =
 γ t R ( s t , a t ) denotes the discounted sum of rewards,-] τ = ( s 0 , a 0 , ..., s T , a T ) denotes a trajectory, and [-π ( τ ) = ρ 0 ( s 0 ) (cid:81) Tt =
-]
Revised: where[-R ( τ ) = (cid:80) Tt =
 γ t R ( s t , a t ) denotes the discounted sum of rewards,-] τ = ( s 0 , a 0 , ...s T , a T ) denotes a sequence of states and [-π ( τ ) = ρ 0 ( s 0 ) (cid:81) Tt =
-]actions induced by the policy and dynamics.

Original: π ( a t | s t ) T ( s t +1 | s t , a t ) .

Original: It can be shown that the trajectory distribution induced by the optimal policy π ∗ ( a | s ) takes the form π ∗ ( τ ) ∝ exp { R ( τ ) }
Revised: It can be shown that the trajectory distribution induced by the optimal policy π ∗ ( a | s ) takes the form π ∗ ( a | s ) ∝ exp { Q ∗ soft ( s t , a t ) }

Original: (Ziebart, 2010).
Revised: (Ziebart, 2010; Haarnoja et al., 2017), where Q

Original:  Inverse reinforcement learning(IRL) instead seeks infer the reward function ˆ R given a set of demonstrations D
Revised: [Equation] [Equation] ] denotes the soft Qfunction. Inverse reinforcement learning instead seeks infer the reward function r ( s, a ) given a set of demonstrations D

Original: = { τ 1 , ..., τ N } .
Revised: = { τ 1 , ..., τ N } .

Original: Under maximum entropy IRL, we assume the demonstrations are drawn from an optimal policy 
Revised: In IRL, we assume the demonstrations are drawn from an optimal policy π ∗ ( a | s ) .

Original: τ i ∼ exp { R ( τ ) } , and we can frame the IRL problem as  maximum likelihoodestimation problem:
Revised: We can interpret the IRL problem as solving the maximum likelihood problem:

Original: [Equation]
Revised: [Equation]

Original: where p θ ( τ ) ∝ exp { R{+p ( s 0 ) (cid:81) Tt =
 p ( s t +1 | s t , a t ) e γ t r+} θ ( τ ) } takes the form of an energy-based model . The main challenge oftraining an energy-based model is computing the partition function
Revised: Where p θ ( τ ) ∝ {+p ( s 0 ) (cid:81) Tt =
 p ( s t +1 | s t , a t ) e γ t r+} θ ( s t ,a t ) parametrizes the reward function r θ ( s, a ) but ﬁxes the dynamics and initial state distribution to that of the MDP.

Original: Z θ = (cid:82) τ exp { R{+) ∝ e (cid:80) Tt =
 γ t r+} θ ( τ ) } , which is intractable in large or continuous MDPs. Alternative approaches, such as sampling (Finn et al., 2016b), must be used to estimate the partition function in such cases.
Revised: Note that under deterministic dynamics, this simpliﬁes to an energy-based model where for feasible trajectories, p θ ( τ {+) ∝ e (cid:80) Tt =
 γ t r+} θ ( s t ,a t ) (Ziebart et al., 2008).

Original: Finn et al.
Revised: Finn et al.

Original: (2016a) propose to castthe optimization in Eqn.
Revised: (2016a) propose to cast optimization of Eqn.

Original: 1 as a GAN (Goodfellow et al., 2014) optimization problem, where the discriminator takes on a particular form ( f θ ( τ ) is a learned function; π ( τ ) is precomputed and its value “ﬁlled in”):
Revised: 1 as a GAN (Goodfellow et al., 2014) optimization problem. They operate in a trajectory-centric formulation, where the discriminator takes on a particular form ( f θ ( τ ) is a learned function; π ( τ ) is precomputed and its value “ﬁlled in”):

Original: [Equation]
Revised: [Equation]

Original: and the policy π is trained to maximize R ( τ ) =
Revised: and the policy π is trained to maximize R ( τ ) =

Original: log(1 − D ( τ ))
Revised: log(1 − D ( τ ))

Original: − log D ( τ ) .
Revised: − log D ( τ ) .

Original: Updating the discriminator can be viewed as updating the reward function, and updating the policy can be viewed as improving the sampling distribution used to estimate the partition function.
Revised: Updating the discriminator can be viewed as updating the reward function, and updating the policy can be viewed as improving the sampling distribution used to estimate the partition function.

Original: If trained to optimality, it can be shown that an optimal reward function can be extracted from the optimal discriminator as f ∗ ( τ )
Revised: If trained to optimality, it can be shown that an optimal reward function can be extracted from the optimal discriminator as f ∗ ( τ )

Original: =
Revised: =

Original: R ∗ ( τ )+ const, and π recovers the optimal policy.
Revised: R ∗ ( τ )+ const, and π recovers the optimal policy.

Original: We refer to this formulation as generative adversarial network guided cost learning (GAN-GCL) to discriminate it from guided cost learning (GCL)
Revised: We refer to this formulation as generative adversarial network guided cost learning (GAN-GCL) to discriminate it from guided cost learning (GCL)

Original: (Finn et al., 2016a).
Revised: (Finn et al., 2016a).

Original: This formulation shares similarities with GAIL (Ho & Ermon, 2016), but GAIL does not place special structure on the discriminator, so the reward cannot be recovered.
Revised: This formulation shares similarities with GAIL (Ho & Ermon, 2016), but GAIL does not place special structure on the discriminator, so the reward cannot be recovered.

Original:  A DVERSARIAL I NVERSE R EINFORCEMENT L EARNING (AIRL)
Revised:  A DVERSARIAL I NVERSE R EINFORCEMENT L EARNING (AIRL)

Original: In practice, using full trajectories as proposed by GAN-GCL can result in high variance estimates as compared to using single state, action pairs, and our experimental results show that this results in very poor learning.
Revised: In practice, using full trajectories as proposed by GAN-GCL can result in high variance estimates as compared to using single state, action pairs, and our experimental results show that this results in very poor learning.

Original: We could instead propose a straightforward conversion of Eqn. 2 into the single state and action case, where:
Revised: We could instead propose a straightforward conversion of Eqn. 2 into the single state and action case, where:

Original: [Equation]
Revised: [Equation]

Original: As in the trajectory-centric case, we can show that, at optimality, f ∗ ( s, a ) = log π ∗ ( a | s ) =
Revised: As in the trajectory-centric case, we can show that, at optimality, f ∗ ( s, a ) = log π ∗ ( a | s ) =

Original: A ∗ ( s, a ) , the advantage function of the optimal policy.
Revised: A ∗ ( s, a ) , the advantage function of the optimal policy.

Original: We justify this, as well as a proof that this algorithm solves the IRL problem in Appendix B .
Revised: We justify this, as well as a proof that this algorithm solves the IRL problem in Appendix A .

Original: This change results in an efﬁcient algorithm for imitation learning.
Revised: This change results in an efﬁcient algorithm for imitation learning.

Original: However, it is less desirable for the purpose of reward learning.
Revised: However, it is less desirable for the purpose of reward learning.

Original: While the advantage is a valid optimal reward function, it is a heavily entangled reward, as it supervises each action based on the action of the optimal policy for the training MDP.
Revised: While the advantage is a valid optimal reward function, it is a heavily entangled reward, as it supervises each action based on the action of the optimal policy for the training MDP.

Original: Based on the analysis in the following Sec. 5, we cannot guarantee that this reward will be robust to changes in environment dynamics.
Revised: Based on the analysis in the following Sec. 5, we cannot guarantee that this reward will be robust to changes in environment dynamics.

Original: In our experiments we demonstrate several cases where this reward simply encourages mimicking the expert policy π ∗ , and fails to produce desirable behavior even when changes to the environment are made.
Revised: In our experiments we demonstrate several cases where this reward simply encourages mimicking the expert policy π ∗ , and fails to produce desirable behavior even when changes to the environment are made.

Original:  T HE R EWARD A MBIGUITY P ROBLEM
Revised:  T HE R EWARD A MBIGUITY P ROBLEM

Original: We now discuss why IRL methods can fail to learn robust reward functions.
Revised: We now discuss why IRL methods can fail to learn robust reward functions.

Original: First, we review the concept of reward shaping.
Revised: First, we review the concept of reward shaping.

Original: Ng et al.
Revised: Ng et al.

Original: (1999) describe a class of reward transformations that preserve the optimal policy.
Revised: (1999) describe a class of reward transformations that preserve the optimal policy.

Original: Their main theoretical result is that under the following reward transformation,
Revised: Their main theoretical result is that under the following reward transformation,

Original: [Equation]
Revised: [Equation]

Original: the optimal policy remains unchanged, for any function Φ : S → R .
Revised: the optimal policy remains unchanged, for any function Φ : S → R .

Original: Moreover, without prior knowledge of the dynamics, this is the only class of reward transformations that exhibits policy invariance.
Revised: Moreover, without prior knowledge of the dynamics, this is the only class of reward transformations that exhibits policy invariance.

Original: Because IRL methods only infer rewards from demonstrations given from an optimal agent, they cannot in general disambiguate between reward functions within this class of transformations, unless the class of learnable reward functions is restricted.
Revised: Because IRL methods only infer rewards from demonstrations given from an optimal agent, they cannot in general disambiguate between reward functions within this class of transformations, unless the class of learnable reward functions is restricted.

Original: We argue that shaped reward functions may not be robust to changes in dynamics.
Revised: We argue that shaped reward functions may not be robust to changes in dynamics.

Original: We formalize this notion by studying policy invariance in two MDPs M, M (cid:48) which share the same reward and differ only in the dynamics, denoted as T and T (cid:48) , respectively.
Revised: We formalize this notion by studying policy invariance in two MDPs M, M (cid:48) which share the same reward and differ only in the dynamics, denoted as T and T (cid:48) , respectively.

Original: Suppose an IRL algorithm recovers a shaped, policy invariant reward
Revised: Suppose an IRL algorithm recovers a shaped, policy invariant reward

Original: ˆ r ( s, a, s (cid:48) )
Revised: ˆ r ( s, a, s (cid:48) )

Original: under MDP M
Revised: under MDP M

Original: where Φ (cid:54) = 

Revised: where Φ (cid:54) =

Original: Then, there exists MDP pairs M, M (cid:48) where changing the transition model from T to T (cid:48) breaks policy invariance on MDP M (cid:48) .
Revised: Then, there exists MDP pairs M, M (cid:48) where changing the transition model from T to T (cid:48) breaks policy invariance on MDP M (cid:48) .

Original: As a simple example, consider deterministic dynamics
Revised: As a simple example, consider deterministic dynamics

Original: T ( s, a ) → s (cid:48) and state-action rewards ˆ r ( s, a ) =
Revised: T ( s, a ) → s (cid:48) and state-action rewards ˆ r ( s, a ) =

Original: r ( s, a ) + γ Φ( T ( s, a )) − Φ( s ) .
Revised: r ( s, a ) + γ Φ( T ( s, a )) − Φ( s ) .

Original: It is easy to see that changing the dynamics T to T (cid:48) such that T (cid:48)
Revised: It is easy to see that changing the dynamics T to T (cid:48) such that T (cid:48)

Original: ( s, a ) (cid:54) =
Revised: ( s, a ) (cid:54) =

Original: T ( s, a ) means that ˆ r ( s, a ) no longer lies in the equivalence class of Eqn.
Revised: T ( s, a ) means that ˆ r ( s, a ) no longer lies in the equivalence class of Eqn.

Original: 3 for M (cid:48) .
Revised: 3 for M (cid:48) .

Original: 5. D ISENTANGLING R EWARDS FROM D YNAMICS
Revised: 5. D ISENTANGLING R EWARDS FROM D YNAMICS

Original: To remove unwanted reward shaping with arbitrary reward function classes, the learned reward function can only depend on the current state s . Action-dependent terms may be used so long as they are independent of the next state s (cid:48) , but this is rarely the case for any MDP. This is both a sufﬁcient and necessary condition.
Revised: First, let the notation Q ∗ r,T ( s, a ) denote the optimal Q-function with respect to a reward function r and dynamics T , and π ∗ r,T ( a | s ) denote the same for policies.

Original: We present a short proof sketch, assuming ergodicity in the
Revised: We ﬁrst deﬁne our notion of a ”disentangled” reward.

Original: MDP, as follows: • Sufﬁciency : Because the learned reward can only depend on the current state, we must have Φ( s (cid:48) ) 
Revised: Deﬁnition 5.1 (Disentangled Rewards) . A reward function r (cid:48) ( s, a, s (cid:48) ) is (perfectly) disentangled with respect to a ground-truth reward r ( s, a, s (cid:48) ) and a set of dynamics T such that under all dynamics T ∈ T , the optimal policy is the same: π ∗ r (cid:48) ,T ( a | s )

Original: =
Revised: =

Original: c ∀ s (cid:48) , where c is a constant, meaning the shaping term γ Φ( s (cid:48) ) − Φ( s ) =
Revised: π ∗ r,T ( a | s ) We could also expand this deﬁnition to include a notion of suboptimality. However, we leave this direction to future work.

Original: ( γ − 1) c is constant.
Revised: Under maximum causal entropy RL, the following condition is equivalent to two optimal policies being equal, since Q-functions and policies are equivalent representations (up to arbitrary functions of state f ( s ) ): [Equation] To remove unwanted reward shaping with arbitrary reward function classes, the learned reward function can only depend on the current state s . We require that the dynamics satisfy a decomposability condition where functions over current states f ( s ) and next states g ( s (cid:48) ) can be isolated from their sum f ( s ) + g ( s (cid:48) ) .

Original: Therefore, with state-only rewards, rewards canonly be modiﬁed up to an additive constant, meaning no shaping is possible in any MDP.
Revised: This can be satisﬁed for example by adding self transitions at each state to an ergodic MDP, or any of the environments used in our experiments. The exact deﬁnition of the condition, as well as proof of the following statements are included in Appendix B. Theorem 5.1. Let r ( s ) be a ground-truth reward, and T be a dynamics model satisfying the decomposability condition.

Original: • Necessity : If learned rewards can be a function of s (cid:48) 
Revised: Suppose IRL recovers a state-only reward r (cid:48) ( s ) such that it produces an optimal policy in T :

Original: (or a that are not independent of s (cid:48) ), then there exists MDPs where reward shaping is possible (for example, the deterministic environment given in the previous section).
Revised: [Equation] Then, r (cid:48) ( s ) is disentangled with respect to all dynamics. Theorem 5.2. If a reward function r (cid:48) ( s, a, s

Original: This is the converse of what we wish to show, which is that forreward shaping to not be possible in all MDPs, the rewards must be stateonly.
Revised: [Equation] ) is disentangled for all dynamics functions, then it must be state-only. i.e. If for all dynamics T , [Equation]

Original:  In the traditional IRL setup, where we learn the reward in a single MDP, our analysis motivates learning reward functions that are solely functions of state.
Revised: Then r (cid:48) is only a function of state. In the traditional IRL setup, where we learn the reward in a single MDP, our analysis motivates learning reward functions that are solely functions of state.

Original: If the ground truth reward is also only a function of state, this allows us to recover the true reward up to a constant for any MDP using the
Revised: If the ground truth reward is also only a function of state, this allows us to recover the true reward up to a constant.

Original: MaxEnt IRL framework.

Original:  L EARNING D ISENTANGLED R EWARDS WITH AIRL
Revised:  L EARNING D ISENTANGLED R EWARDS WITH AIRL

Original: In the method presented in Section 4, we cannot learn a state-only reward function, r θ ( s ) , meaning that we cannot guarantee that learned rewards will not be shaped.
Revised: In the method presented in Section 4, we cannot learn a state-only reward function, r θ ( s ) , meaning that we cannot guarantee that learned rewards will not be shaped.

Original: In order to decouple the reward function from the advantage, we propose to modify the discriminator of Sec.
Revised: In order to decouple the reward function from the advantage, we propose to modify the discriminator of Sec.

Original: 4 with the form:
Revised: 4 with the form:

Original: [Equation]
Revised: [Equation]

Original: where f θ,φ is restricted to a reward approximator g θ and a shaping term h φ as
Revised: where f θ,φ is restricted to a reward approximator g θ and a shaping term h φ as

Original: [Equation]
Revised: [Equation]

Original: The additional shaping term helps mitigate the effects of unwanted shaping on our reward approximator g θ (and as we will show, in some cases it can account for all shaping effects).
Revised: The additional shaping term helps mitigate the effects of unwanted shaping on our reward approximator g θ (and as we will show, in some cases it can account for all shaping effects).

Original: The entire training procedure is detailed in Algorithm 1.
Revised: The entire training procedure is detailed in Algorithm 1.

Original: Our algorithm resembles GAIL (Ho & Ermon, 2016) and
Revised: Our algorithm resembles GAIL (Ho & Ermon, 2016) and

Original: GAN-GCL
Revised: GAN-GCL

Original: (Finn et al., 2016a), where we alternate between training a discriminator to classify expert data from policy samples, and update the policy to confuse the discriminator.
Revised: (Finn et al., 2016a), where we alternate between training a discriminator to classify expert data from policy samples, and update the policy to confuse the discriminator.

Original: The advantage of this approach is that we can now parametrize g θ ( s ) as solely a function of the state, allowing us to extract rewards that are disentangled from the dynamics of the environment in which they were trained.
Revised: The advantage of this approach is that we can now parametrize g θ ( s ) as solely a function of the state, allowing us to extract rewards that are disentangled from the dynamics of the environment in which they were trained.

Original: In fact, under this restricted case, we can show the following under deterministic environments with a state-only ground truth reward (proof in Appendix C):
Revised: In fact, under this restricted case, we can show the following under deterministic environments with a state-only ground truth reward (proof in Appendix C):

Original: [Equation]
Revised: [Equation]

Original: where r ∗ is the true reward function.
Revised: where r ∗ is the true reward function.

Original: Since f ∗ must recover to the advantage as shown in Sec. 4, h recovers the optimal value function V ∗ , which serves as the reward shaping term.
Revised: Since f ∗ must recover to the advantage as shown in Sec. 4, h recovers the optimal value function V ∗ , which serves as the reward shaping term.

Original: To be consistent with Sec. 4, an alternative way to interpret the form of Eqn. 4 is to view f θ,φ as the advantage under deterministic dynamics
Revised: To be consistent with Sec. 4, an alternative way to interpret the form of Eqn. 4 is to view f θ,φ as the advantage under deterministic dynamics

Original: [Equation]
Revised: [Equation]

Original: In stochastic environments, we can instead view f ( s, a, s (cid:48) ) as a single-sample estimate of A ∗ ( s, a ) .
Revised: In stochastic environments, we can instead view f ( s, a, s (cid:48) ) as a single-sample estimate of A ∗ ( s, a ) .

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In our experiments, we aim to answer two questions:
Revised: In our experiments, we aim to answer two questions:

Original: Can AIRL learn disentangled rewards that are robust to changes in environment dynamics?
Revised: Can AIRL learn disentangled rewards that are robust to changes in environment dynamics?

Original: Is AIRL efﬁcient and scalable to high-dimensional continuous control tasks?
Revised: Is AIRL efﬁcient and scalable to high-dimensional continuous control tasks?

Original: To answer 1, we evaluate AIRL in transfer learning scenarios, where a reward is learned in a training environment, and optimized in a test environment with signiﬁcantly different dynamics.
Revised: To answer 1, we evaluate AIRL in transfer learning scenarios, where a reward is learned in a training environment, and optimized in a test environment with signiﬁcantly different dynamics.

Original: We show that rewards learned with our algorithm under the constraint presented in Section 5 still produce optimal or near-optimal behavior, while na¨ıve methods that do not consider reward shaping fail.
Revised: We show that rewards learned with our algorithm under the constraint presented in Section 5 still produce optimal or near-optimal behavior, while na¨ıve methods that do not consider reward shaping fail.

Original: We also show that in small MDPs, we can recover the exact ground truth reward function.
Revised: We also show that in small MDPs, we can recover the exact ground truth reward function.

Original: To answer 2, we compare AIRL as an imitation learning algorithm against GAIL (Ho & Ermon, 2016) and the GAN-based GCL algorithm proposed by does notimplement or evaluate GAN-GCL and, to our knowledge, we present the ﬁrst empirical evaluation of this algorithm.
Revised: To answer 2, we compare AIRL as an imitation learning algorithm against GAIL (Ho & Ermon, 2016) and the GAN-based GCL algorithm proposed by Finn et al. (2016a), which we refer to as GAN-GCL, on standard benchmark tasks that do not evaluate transfer. Note that Finn et al.

Revised: (2016a) does not implement or evaluate GAN-GCL and, to our knowledge, we present the ﬁrst empirical evaluation of this algorithm.

Original: We ﬁnd that AIRLperforms on par with GAIL in a traditional imitation learning setup while vastly outperforming it in transfer learning setups, and outperforms GAN-GCL in both settings.
Revised: We ﬁnd that AIRL performs on par with GAIL in a traditional imitation learning setup while vastly outperforming it in transfer learning setups, and outperforms GAN-GCL in both settings.

Original: It is worth noting that, except for (Finn et al., 2016b), our method is the only IRL algorithm that we are aware of that scales to high dimensional tasks with unknown dynamics, and although GAIL (Ho & Ermon, 2016) resembles an IRL algorithm in structure, it does not recover disentangled reward functions, making it unable to re-optimize the learned reward under changes in the environment, as we illustrate below.
Revised: It is worth noting that, except for (Finn et al., 2016b), our method is the only IRL algorithm that we are aware of that scales to high dimensional tasks with unknown dynamics, and although GAIL (Ho & Ermon, 2016) resembles an IRL algorithm in structure, it does not recover disentangled reward functions, making it unable to re-optimize the learned reward under changes in the environment, as we illustrate below.

Original: For our continuous control tasks, we use trust region policy optimization (Schulman et al., 2015) as our policy optimization algorithm across all evaluated methods, and in the tabular MDP task, we use soft value iteration.
Revised: For our continuous control tasks, we use trust region policy optimization (Schulman et al., 2015) as our policy optimization algorithm across all evaluated methods, and in the tabular MDP task, we use soft value iteration.

Original: We obtain expert demonstrations by training an expert policy on the ground truth reward, but hide the ground truth reward from the IRL algorithm.
Revised: We obtain expert demonstrations by training an expert policy on the ground truth reward, but hide the ground truth reward from the IRL algorithm.

Revised: In this way, we simulate a scenario where we wish to use RL to solve a task but wish to refrain from manual reward engineering and instead seek to learn a reward function from demonstrations.

Original: Our code and additional supplementary material including videos will be available at https://sites.google.com/view/ adversarial-irl , and hyper-parameter and architecture choices are detailed in Appendix D.
Revised: Our code and additional supplementary material including videos will be available at https://sites.google.com/view/ adversarial-irl , and hyper-parameter and architecture choices are detailed in Appendix D.

Original: 7. R ECOVERING TRUE REWARDS IN TABULAR MDP S
Revised: 7. R ECOVERING TRUE REWARDS IN TABULAR MDP S

Original: We ﬁrst consider MaxEnt IRL in a toy task with randomly generated MDPs.
Revised: We ﬁrst consider MaxEnt IRL in a toy task with randomly generated MDPs.

Original: The MDPs havestates, 4 actions, randomly drawn transition matrices, and a reward function that always gives a reward of 1 . 0 when taking an action from state 0.
Revised: The MDPs havestates, 4 actions, randomly drawn transition matrices, and a reward function that always gives a reward of 1 . 0 when taking an action from state 0.

Original: The initial state is always state 1.
Revised: The initial state is always state 1.

Original: The optimal reward, learned reward with a state-only reward function, and learned reward using a state-action reward function are shown in Fig.
Revised: The optimal reward, learned reward with a state-only reward function, and learned reward using a state-action reward function are shown in Fig.

Original: We subtract a constant offset from all reward functions so that they share the same mean for visualization - this does not inﬂuence the optimal policy.
Revised: We subtract a constant offset from all reward functions so that they share the same mean for visualization - this does not inﬂuence the optimal policy.

Original: AIRL with a state-only reward function is able to recover the ground truth reward, but AIRL with a state-action reward instead recovers a shaped advantage function.
Revised: AIRL with a state-only reward function is able to recover the ground truth reward, but AIRL with a state-action reward instead recovers a shaped advantage function.

Original: We also show that in the transfer learning setup, under a new transition matrix T (cid:48) , the optimal policy under the state-only reward achieves optimal performance (it is identical to the ground truth reward) whereas the state-action reward only improves marginally over uniform random policy.
Revised: We also show that in the transfer learning setup, under a new transition matrix T (cid:48) , the optimal policy under the state-only reward achieves optimal performance (it is identical to the ground truth reward) whereas the state-action reward only improves marginally over uniform random policy.

Original: The learning curve for this experiment is shown in Fig 2.
Revised: The learning curve for this experiment is shown in Fig 2.

Original: 7. D ISENTANGLING R EWARDS IN C ONTINUOUS C ONTROL T ASKS
Revised: 7. D ISENTANGLING R EWARDS IN C ONTINUOUS C ONTROL T ASKS

Original: To evaluate whether our method can learn disentangled rewards in higher dimensional environments, we perform transfer learning experiments on continuous control tasks.
Revised: To evaluate whether our method can learn disentangled rewards in higher dimensional environments, we perform transfer learning experiments on continuous control tasks.

Original: In each task, a reward is learned via IRL on the training environment, and the reward is used to reoptimize a new policy on a test environment.
Revised: In each task, a reward is learned via IRL on the training environment, and the reward is used to reoptimize a new policy on a test environment.

Original: We train two IRL algorithms, AIRL and GAN-GCL, with state-only and stateaction rewards.
Revised: We train two IRL algorithms, AIRL and GAN-GCL, with state-only and stateaction rewards.

Original: We also include results for directly transferring the policy learned with GAIL, and an oracle result that involves optimizing the ground truth reward function with TRPO.
Revised: We also include results for directly transferring the policy learned with GAIL, and an oracle result that involves optimizing the ground truth reward function with TRPO.

Original: Numerical results for these environment transfer experiments are given in Table 1.
Revised: Numerical results for these environment transfer experiments are given in Table 1.

Original: The ﬁrst task involves a 2D pointmass navigating to a goal position in a small maze when the position of the walls are changed between train and test time.
Revised: The ﬁrst task involves a 2D point mass navigating to a goal position in a small maze when the position of the walls are changed between train and test time.

Original: At test time, the agent cannot simply mimic the actions learned during training, and instead must successfully infer that the goal in the maze is to reach the target.
Revised: At test time, the agent cannot simply mimic the actions learned during training, and instead must successfully infer that the goal in the maze is to reach the target.

Original: The task is shown in Fig.
Revised: The task is shown in Fig.

Original: Only AIRL trained with state-only rewards is able to consistently navigate to the goal when the maze is modiﬁed.
Revised: Only AIRL trained with state-only rewards is able to consistently navigate to the goal when the maze is modiﬁed.

Original: Direct policy transfer and state-action IRL methods learn rewards which encourage the agent to take the same path taken in the training environment, which is blocked in the test environment.
Revised: Direct policy transfer and state-action IRL methods learn rewards which encourage the agent to take the same path taken in the training environment, which is blocked in the test environment.

Original: We plot the learned reward in Fig.
Revised: We plot the learned reward in Fig.

Original: In our second task, we modify the agent itself.
Revised: In our second task, we modify the agent itself.

Original: We train a quadrupedal “ant” agent to run forwards, and at test time we disable and shrink two of the front legs of the ant such that it must signiﬁcantly change its gait (shown in Fig. 4).
Revised: We train a quadrupedal “ant” agent to run forwards, and at test time we disable and shrink two of the front legs of the ant such that it must signiﬁcantly change its gait.

Original: We ﬁnd that AIRL is able to learn reward functions that encourage the ant to move forwards, acquiring a modiﬁed gait that involves orienting itself to face the forward direction and crawling with its two hind legs.
Revised: We ﬁnd that AIRL is able to learn reward functions that encourage the ant to move forwards, acquiring a modiﬁed gait that involves orienting itself to face the forward direction and crawling with its two hind legs.

Original: Alternative methods, including transferring a policy learned by GAIL (which achieves near-optimal performance with the unmodiﬁed agent), fail to move forward at all.
Revised: Alternative methods, including transferring a policy learned by GAIL (which achieves near-optimal performance with the unmodiﬁed agent), fail to move forward at all.

Original: We show the qualitative difference in behavior in Fig.
Revised: We show the qualitative difference in behavior in Fig.

Original: We have demonstrated that AIRL can learn disentangled rewards that can accommodate signiﬁcant domain shift even in high-dimensional environments where it is difﬁcult to exactly extract the true reward.
Revised: We have demonstrated that AIRL can learn disentangled rewards that can accommodate signiﬁcant domain shift even in high-dimensional environments where it is difﬁcult to exactly extract the true reward.

Original: GAN-GCL can presumably learn disentangled rewards, but we ﬁnd that the trajectorycentric formulation does not perform well even in learning rewards in the original task, let alone transferring to a new domain.
Revised: GAN-GCL can presumably learn disentangled rewards, but we ﬁnd that the trajectorycentric formulation does not perform well even in learning rewards in the original task, let alone transferring to a new domain.

Original: GAIL learns successfully in the training domain, but does not acquire a representation that is suitable for transfer to test domains.
Revised: GAIL learns successfully in the training domain, but does not acquire a representation that is suitable for transfer to test domains.

Original: [Table]
Revised: [Table]

Original: 7. B ENCHMARK T ASKS FOR I MITATION L EARNING
Revised: B ENCHMARK T ASKS FOR I MITATION L EARNING

Original: Finally
Revised: Finally

Original: [Table]
Revised: [Table]

Original: we evaluate AIRL as an imitation learning algorithm against the GAN-GCL and the stateof-the-art GAIL on several benchmark tasks.
Revised: we evaluate AIRL as an imitation learning algorithm against the GAN-GCL and the stateof-the-art GAIL on several benchmark tasks.

Original: Each algorithm is presented with 50 expert demonstra-Note that there is little reward shaping, which enables the rewardto transfer well. tions, collected from a policy trained with TRPO on the ground truth reward function.
Revised: Each algorithm is presented with 50 expert demonstra- reward transfer well behavior tions, collected from a policy trained with TRPO on the ground truth reward function.

Original: For AIRL, we use an unrestricted state-action reward function as we are not concerned with reward transfer.
Revised: For AIRL, we use an unrestricted state-action reward function as we are not concerned with reward transfer.

Original: Numerical results are presented in Table 2, and sample images from these tasks are shown in Fig. These experiments do not test transfer, and in a sense can be regarded as “testing on the training set
Revised: Numerical results are presented in Table 2.These experiments do not test transfer and in a sense can be regarded as “testing on the training set,” but they match the settings reported in prior work (Ho &

Original: ” but they match the settings reported in prior work (Ho & Ermon, 2016)
Revised: Ermon, 2016).

Original: We ﬁnd that the performance difference between AIRL and GAIL is negligible, even though AIRL is a true IRL algorithm that recovers reward functions, while GAIL does not.
Revised: We ﬁnd that the performance difference between AIRL and GAIL is negligible, even though AIRL is a true IRL algorithm that recovers reward functions, while GAIL does not.

Original: Both methods achieve close to the best possible result on each task and there is little room for improvement.
Revised: Both methods achieve close to the best possible result on each task and there is little room for improvement.

Original: This result goes against the belief that IRL algorithms are indirect, and less efﬁcient that direct imitation learning algorithms (Ho & Ermon, 2016).
Revised: This result goes against the belief that IRL algorithms are indirect, and less efﬁcient that direct imitation learning algorithms (Ho & Ermon, 2016).

Original: The GAN-GCL method is ineffective on all but the simplest
Revised: The GAN-GCL method is ineffective on all but the simplest

Original: Pendulum task when trained with the same number of samples as AIRL and GAIL.
Revised: Pendulum task when trained with the same number of samples as AIRL and GAIL.

Original: We ﬁnd that a discriminator trained over trajectories easily overﬁts and provides poor learning signal for the policy.
Revised: We ﬁnd that a discriminator trained over trajectories easily overﬁts and provides poor learning signal for the policy.

Original: Our results illustrate that AIRL achieves the same performance as GAIL on benchmark imitation tasks that do not require any generalization.
Revised: Our results illustrate that AIRL achieves the same performance as GAIL on benchmark imitation tasks that do not require any generalization.

Original: On tasks that require transfer and generalization, illustrated in the previous section, AIRL outperforms GAIL by a wide margin, since our method is able to recover disentangled rewards that transfer effectively in the presence of domain shift.
Revised: On tasks that require transfer and generalization, illustrated in the previous section, AIRL outperforms GAIL by a wide margin, since our method is able to recover disentangled rewards that transfer effectively in the presence of domain shift.

Original: [Table]
Revised: [Table]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We presented AIRL, a practical and scalable IRL algorithm that can learn disentangled rewards and greatly outperforms both prior imitation learning and IRL algorithms.
Revised: We presented AIRL, a practical and scalable IRL algorithm that can learn disentangled rewards and greatly outperforms both prior imitation learning and IRL algorithms.

Original: We show that rewards learned with AIRL transfer effectively under variation in the underlying domain, in contrast to unmodiﬁed
Revised: We show that rewards learned with AIRL transfer effectively under variation in the underlying domain, in contrast to unmodiﬁed

Original: IRL methods which tend to recover brittle rewards that do not generalize well and GAIL, which does not recover reward functions at all.
Revised: IRL methods which tend to recover brittle rewards that do not generalize well and GAIL, which does not recover reward functions at all.

Original: In small MDPs where the optimal policy and reward are unambiguous, we also show that we can exactly recover the ground-truth rewards up to a constant.
Revised: In small MDPs where the optimal policy and reward are unambiguous, we also show that we can exactly recover the ground-truth rewards up to a constant.

Original: A PPENDICES A G UIDED C OST L EARNING
Revised: A PPENDICES A J USTIFICATION OF AIRL

Original: We justify AIRL by showing equivalence with guided cost learning (GCL). Thus, we ﬁrst give a brief overview of GCL Finn et al.
Revised: In this section, we show that the objective of AIRL matches that of solving the maximum causal entropy IRL problem. We use a similar method as Finn et al.

Original: (2016b).
Revised: (2016a), which shows the justiﬁcation of GAN-GCL for the trajectory-centric formulation. For simplicity we derive everything in the undiscounted case. A. S ETUP

Original: As explained in Section 3, the goal ofMaxEnt IRL Ziebart et al.
Revised: As mentioned in Section 3, the goal of IRL can be seen as training a generative model over trajectories as:

Original: (2008) is to train a generative model over demonstrations max θ E D [log p θ ( τ )] .
Revised: [Equation] Where the distribution p θ ( τ ) is parametrized as p θ ( τ ) ∝ p ( s 0 ) (cid:81) T

Original: We assume p θ takes the general form of an energy-based model, pθ ( τ ) ∝ exp { r θ ( τ )} . 
Revised: [Equation] 1 t =0 p ( s t +1 | s t , a t ) e r θ ( s t ,a t ) . We can compute the gradient with respect to θ as follows: [Equation]

Original: Using importance sampling from a policy π to estimate the partition function Z = (cid:82) τ exp { r θ ( τ ) } , the objective is:
Revised: Let p θ,t ( s t , a t ) = (cid:82) s t (cid:48)(cid:54) =

Revised: t ,a t (cid:48)(cid:54) =

Revised: t p θ ( τ ) denote the state-action marginal at time t .

Revised: Rewriting the above equation, we have:

Original: [Equation]
Revised: [Equation]

Revised: As it is difﬁcult to draw samples from p θ , we instead train a separate importance sampling distribution µ ( τ ) .

Original: Following Finn et al.
Revised: For the choice of this distribution, we follow Finn et al.

Original: (2016a), it is useful to replace π with a mixture distribution µ ( τ ) =
Revised: (2016a) and use a mixture policy µ ( a | s ) =

Original: 12 π ( τ ) + 12 ˆ p ( τ ) , where ˆ p ( τ ) is a rough density estimate over the demonstrations.
Revised: 12 π ( a | s ) + 12 ˆ p ( a | s ) , where ˆ p ( a | s ) is a rough density estimate trained on the demonstrations.

Original: This is justiﬁed as reducing the variance of the importance sampling estimate when the policy π ( τ ) has poor coverage over the demonstrations in the early stages of training. 
Revised: This is justiﬁed as reducing the variance of the importance sampling estimate when the policy π ( a | s ) has poor coverage over the demonstrations in the early stages of training. Thus, our new gradient is:

Original: [Equation]
Revised: [Equation]

Original: GCL additionally adapts the importance sampler  to reduce variance. The sampler minimizes the
Revised: We additionally wish to adapt the importance sampler π to reduce variance, by minimizing D KL ( π ( τ ) ||

Original: KL-divergence between π ( τ ) and e r θ ( τ ) /Z by optimizing the objective:
Revised: p θ ( τ )) . The policy trajectory distribution factorizes as π ( τ ) = p

Original: [Equation]
Revised: [Equation]

Original: The optimal policy that minimizes the KL is thus π ( τ ) ∝
Revised: T

Original: [Equation]
Revised: [Equation]

Original: ) , meaning if r θ ( τ ) is optimal, then the optimal π θ ( τ ) = π ∗ ( τ ) .
Revised: = [Equation] p [Equation]

Original: The exact form of the learned reward function is r θ ( τ ) =
Revised: s [Equation]

Original: [-(cid:80) Tt =
-]
Revised: [-(cid:80) Tt =
-]s t [Equation]

Original: γ t r θ (s t , a t  ) , where r θ is an arbitrary parametrized function, such as a neural network. B S INGLE STATE - ACTION DISCRIMINATORS
Revised: a t ) π ( a t | s t ) . The dynamics and initial state terms inside π ( τ ) and p θ ( τ ) cancel, leaving the entropy-regularized policy objective: [Equation]

Original: We replace the cost learning objective with training a discriminator of the following form:
Revised: In AIRL, we replace the cost learning objective with training a discriminator of the following form:

Original: [Equation]
Revised: [Equation]

Original: The objective of the discriminator is to minimize cross-entropy loss between expert demonstrations and generated samples:
Revised: The objective of the discriminator is to minimize cross-entropy loss between expert demonstrations and generated samples:

Original: [Equation]
Revised: [Equation]

Revised: θ ( s t , a t ))]

Original: We replace the policy optimization objective with the following reward:
Revised: We replace the policy optimization objective with the following reward:

Original: [Equation]
Revised: [Equation]

Original: B. E QUIVALENCE TO GCL
Revised: A. D ISCRIMINATOR O BJECTIVE

Original: We can now show equivalence with GCL in the undiscounted setting, using a similar method as Finn et al.
Revised: First, we show that training the gradient of the discriminator objective is the same as Eqn. 5.

Original: (2016a).

Original: In the discounted setting, an approximate equivalence can be made which lacks discounting on the entropy term of the policy objective, by replacing f θ ( s t , a t ) with γ t f θ ( s t , a t ) in the discriminator.

Original: We ﬁrst show that the cross-entropy discriminator loss matches the loss of the cost of GCL up to a constant.

Original: First, we write the negative loss to turn the minimization problem into maximization. We use µ to denote a mixture between the dataset and policysamples, and ˆ µ to denote a mixture between exp { f θ ( s, a ) } and policy samples.
Revised: We write the negative loss to turn the minimization problem into maximization, and use µ to denote a mixture between the dataset and policy samples.

Original: [Equation]
Revised: [Equation]

Original: Taking the derivative w.r.t.
Revised: Taking the derivative w [Equation] r

Revised: [Equation]

Original:  θ ,
Revised: t. θ ,

Original: [Equation]
Revised: [Equation]

Revised: Multiplying the top and bottom of the fraction in the second expectation by the state marginal π ( s t )

Revised: = (cid:82) a π t ( s t , a t ) , and grouping terms we get:

Revised: [Equation]

Revised: Where we have written ˆ p θ,t ( s t , a t )

Original:  This expressionexactly matchesthe derivative of Eqn. 6.
Revised: = exp { f θ ( s t , a t ) } π t ( s t ) , and ˆ µ to denote a mixture between ˆ p θ ( s, a ) and policy samples. This expression matches Eqn. 5, with f θ ( s, a ) serving as the reward function, when π maximizes the policy objective so that ˆ p θ ( s, a ) = p θ ( s, a ) .

Original: Next, we show that the policy objective matches that of the sampler of GCL.
Revised: Next, we show that the policy objective matches that of the sampler of Eqn. 6.

Original: The objective of the policy is to maximize with respect to the reward ˆ r t ( s, a ) .
Revised: The objective of the policy is to maximize with respect to the reward ˆ r t ( s, a ) .

Original: First, note that:
Revised: First, note that:

Original: [Equation]
Revised: [Equation]

Original: Thus, when ˆ r ( s, a ) is summed over entire trajectories, we obtain the  objective
Revised: Thus, when ˆ r ( s, a ) is summed over entire trajectories, we obtain the entropy-regularized policy objective

Original: [Equation]
Revised: [Equation]

Original: [-This exactly matches the objective of the GCL sampler, where the reward is deﬁned as r θ ( τ ) = (cid:80) Tt =
-] f θ ( s t , a t ) .
Revised: [-This exactly matches the objective of the GCL sampler, where the reward is deﬁned as r θ ( τ ) = (cid:80) Tt =
-]Where f θ serves as the reward function.

Original: B. f θ ( s, a ) RECOVERS THE ADVANTAGE
Revised: A. f θ ( s, a ) RECOVERS THE ADVANTAGE

Original: The global minimum of the discriminator objective is achieved when π = π E , where π denotes the learned policy (the ”generator” of a GAN) and π E denotes the policy under which demonstrations were collected (Goodfellow et al., 2014).
Revised: The global minimum of the discriminator objective is achieved when π = π E , where π denotes the learned policy (the ”generator” of a GAN) and π E denotes the policy under which demonstrations were collected (Goodfellow et al., 2014).

Original: At this point, the output of the discriminator is 12 for all values of s, a , meaning we have exp { f θ ( s, a ) } =
Revised: At this point, the output of the discriminator is 12 for all values of s, a , meaning we have exp { f θ ( s, a ) } =

Original: π E ( a | s ) , or f ∗ ( s, a ) = log π E ( a | s )
Revised: π E ( a | s ) , or f ∗ ( s, a ) = log π E ( a | s )

Original: =
Revised: =

Original: A ∗ ( s, a ) .
Revised: A ∗ ( s, a ) .

Revised: B S TATE -O NLY I NVERSE R EINFORCEMENT L EARNING

Revised: In this section we include proofs for Theorems 5.1 and 5.2, and the condition on the dynamics necessary for them to hold.

Original: C AIRL RECOVERS REWARDS UP TO CONSTANTS
Revised: Deﬁnition B.1 (Decomposability Condition) . Two states s 1 , s 2 are deﬁned as ”1-step linked” under a dynamics or transition distribution T ( s (cid:48) | a, s ) if there exists a state s that can reach s 1 and s 2 with positive probability in one time step. Also, we deﬁne that this relationship can transfer through transitivity: if s 1 and s 2 are linked, and s 2 and s 3 are linked, then we also consider s 1 and s 3 to be linked. A transition distribution T satisﬁes the decomposability condition if all states in the MDP are linked with all other states.

Revised: The key reason for needing this condition is that it allows us to decompose the functions state dependent f ( s ) and next state dependent g ( s (cid:48) ) from their sum f

Revised: ( s ) + g ( s (cid:48) ) , as stated below:

Revised: Lemma B.1.

Revised: Suppose the dynamics for an MDP satisfy the decomposability condition.

Revised: Then, for functions a ( s ) , b ( s ) , c ( s ) , d ( s ) , if for all s, s (cid:48) :

Revised: [Equation]

Revised: Then for for all s ,

Revised: [Equation]

Revised: Proof.

Revised: Rearranging, we have:

Revised: [Equation]

Revised: Let us rewrite f ( s )

Revised: = a ( s ) − c ( s ) .

Revised: This means we have f ( s ) =

Revised: b ( s (cid:48) )

Revised: − d ( s (cid:48) ) for some function only dependent on s .

Revised: In order for this to be representable, the term b ( s (cid:48) ) − d ( s (cid:48) ) must be equal for all successor states s (cid:48) from s .

Revised: Under the decomposability condition, all successor states must therefore be equal in this manner through transitivity, meaning we have b ( s (cid:48) )

Revised: − d ( s (cid:48) ) must be constant with respect to s .

Revised: Therefore, a ( s ) =

Revised: c ( s ) + const.

Revised: We can then substitute this expression back in to the original equation to derive b ( s ) =

Revised: d ( s ) + const.

Revised: We consider the case when the ground truth reward is state-only.

Revised: We now show that if the learned reward is also state-only, then we guarantee learning disentangled rewards, and vice-versa (sufﬁciency and necessity).

Revised: Theorem 5.1.

Revised: Let r ( s ) be a ground-truth reward, and T be a dynamics model satisfying the decomposability condition.

Revised: Suppose IRL recovers a state-only reward r (cid:48)

Revised: ( s ) such that it produces an optimal policy in T :

Revised: [Equation]

Revised: Then, r (cid:48) ( s ) is disentangled with respect to all dynamics.

Revised: Proof.

Revised: We show that r (cid:48) ( s ) must equal the ground-truth reward up to constants (modifying rewards by constants does not change the optimal policy).

Revised: Let r (cid:48)

Revised: ( s ) =

Revised: r ( s ) + φ ( s ) for some arbitrary function of state φ ( s ) .

Revised: We have:

Revised: [Equation]

Revised: From here

Revised: [Equation]

Revised: we see that:

Revised: [Equation]

Revised: Meaning we must have for all s, a :

Revised: [Equation]

Revised: This places the requirement that all successor states from s must have the same potential f ( s ) .

Revised: Under the decomposability condition, every state in the MDP can be linked with such an equality statement, meaning that f ( s ) is constant.

Revised: Thus, r (cid:48) ( s ) =

Revised: r ( s ) + const.

Revised: Theorem 5.2.

Revised: If a reward function r (cid:48) ( s, a, s (cid:48) ) is disentangled for all dynamics functions, then it must be state-only.

Revised: i.e. If for all dynamics T ,

Revised: [Equation]

Revised: Then r (cid:48) is only a function of state.

Revised: Proof.

Revised: We show the converse, namely that if r (cid:48) ( s, a, s (cid:48) ) can depend on a or s (cid:48) , then there exists a dynamics model T such that the optimal policy is changed, i.e. Q ∗ r,T ( s, a ) (cid:54)

Revised: [Equation]

Revised: f ( s ) ∀ s, a .

Revised: Consider the following 3-state MDP with deterministic dynamics and starting state S :

Revised: [Equation]

Revised: We denote the action with a small letter, i.e. taking the action a from S brings the agent to state A , receiving a reward of 0 .

Revised: For simplicity, assume the discount factor γ = 1 .

Revised: The optimal policy here takes the a action, returns to s , and repeat for inﬁnite positive reward.

Revised: An action-dependent reward which induces the same optimal policy would be to move the reward from the action returning to s to the action going to a or s : r (cid:48)

Revised: ( s, a ) =

Revised: State

Revised: Action

Revised: Reward

Revised: S a +

Revised: S

Revised: [Equation]

Revised: This corresponds to the shaping potential φ ( S ) = 0 , φ ( A ) = 1 , φ ( B ) = − 1 .

Revised: Now suppose we modify the dynamics such that action a leads to B and action b leads to A :

Revised: [Equation]

Revised: Optimizing r (cid:48) on this new MDP results in a different policy than optimizing r , as the agent visits B, resulting in inﬁnite negative reward.

Original: In this section, we prove that AIRL can recover the ground truth reward up to constants if the ground truth is only a function of state r ( s ) .
Revised: In this section, we prove that AIRL can recover the ground truth reward up to constants if the ground truth is only a function of state r ( s ) .

Original: For simplicity, we consider deterministic environments, so that s (cid:48) is uniquely deﬁned by s, a , and we restrict AIRL’s reward estimator g to only be a function of state. 
Revised: For simplicity, we consider deterministic environments, so that s (cid:48) is uniquely deﬁned by s, a , and we restrict AIRL’s reward estimator g to only be a function of state. Theorem C.1. Suppose we use AIRL with a discriminator of the form [Equation]

Original: AIRL uses a speciﬁc form of the discriminator used in Appendix B. Crucially, westill have that the advantage is recovered by f ∗ :
Revised: φ [Equation] We also assume we have deterministic dynamics, in addition to the decomposability condition on the dynamics from Thm 5.1. Then if AIRL recovers the optimal f ∗ ( s, a, s (cid:48) ) , we have

Original: [Equation]
Revised: [Equation]

Original: We can now expand A  ( s, a ) and symbol match:
Revised: Proof. From Appendix A.4, we have f ∗ ( s, a, s (cid:48) ) = A ∗ ( s, a ) , so f ∗ ( s, a, s (cid:48) ) =

Revised: Q ∗ ( s, a ) − V ∗ ( s ) =

Revised: r ( s

Revised: [Equation]

Revised: +

Revised: [Equation]

Revised: ) −

Original: [Equation]
Revised: [Equation]

Revised: s

Original: [Equation]
Revised: [Equation]

Original: We must have h ∗ ( s (cid:48) )
Revised: Substituting the form of f , we have for all s, s (cid:48) : [Equation] s [Equation]

Revised: Applying Lemma B.1 with a ( s )

Original: =
Revised: =

Original: V ( s (cid:48) ) + c for some arbitrary constant c because it is the only item that depends on the next state s (cid:48) . Under an ergodicity assumption, the spaces for s (cid:48) and s are identical, so γh ∗ ( s(cid:48) ) − h ∗ ( s ) =
Revised: g ∗ ( s ) − h ∗ ( s ) , b ( s (cid:48) )

Original: ( γ − 1) c , and g ∗ ( s )

Original: =
Revised: =

Original: r ( s  ) − ( γ − 1) c .
Revised: γh ∗ ( s (cid:48) ) , c ( s ) =

Original: Thus, we can recover both the reward and value function up to (tied) constants.
Revised: r ( s ) − V ∗ ( s ) , and d ( s (cid:48) ) = γV ∗ ( s (cid:48) ) we have the result.

Original: D E XPERIMENT D ETAILS
Revised: D E XPERIMENT D ETAILS

Original: In this section we detail hyperparameters and training procedures used for our experiments.
Revised: In this section we detail hyperparameters and training procedures used for our experiments.

Original: These hyperparameters were selected via a grid search.
Revised: These hyperparameters were selected via a grid search.

Original: D. N ETWORK A RCHITECTURES
Revised: D. N ETWORK A RCHITECTURES

Original: For the tabular MDP environment, we also use a tabular representation for function approximators.
Revised: For the tabular MDP environment, we also use a tabular representation for function approximators.

Original: For continuous control experiments, we use a two-layer ReLU network with 32 units for the discriminator of GAIL and GAN-GCL.
Revised: For continuous control experiments, we use a two-layer ReLU network with 32 units for the discriminator of GAIL and GAN-GCL.

Original: For AIRL, we use a linear function approximator for the reward term g and a 2-layer ReLU network for the shaping term h .
Revised: For AIRL, we use a linear function approximator for the reward term g and a 2-layer ReLU network for the shaping term h .

Original: For the policy, we use a two-layer (32 units)
Revised: For the policy, we use a two-layer (32 units)

Original: ReLU gaussian policy.
Revised: ReLU gaussian policy.

Original: D. O THER H YPERPARAMETERS
Revised: D. O THER H YPERPARAMETERS

Original: Entropy regularization : We use an entropy regularizer weight of 0 . 1 for Ant, Swimmer, and
Revised: Entropy regularization : We use an entropy regularizer weight of 0 . 1 for Ant, Swimmer, and

Original: HalfCheetah across all methods.
Revised: HalfCheetah across all methods.

Original: We use an entropy regularizer weight of 1 . 0 on the pointmass environment.
Revised: We use an entropy regularizer weight of 1 . 0 on the point mass environment.

Original: TRPO Batch Size : For Ant, Swimmer and HalfCheetah environments, we use a batch size ofsteps per TRPO update.
Revised: TRPO Batch Size : For Ant, Swimmer and HalfCheetah environments, we use a batch size ofsteps per TRPO update.

Original: For pendulum, we use a batch size of 2000.
Revised: For pendulum, we use a batch size of 2000.

Original: D. O THER T RAINING D ETAILS
Revised: D. O THER T RAINING D ETAILS

Original: IRL methods commonly learn rewards which explain behavior locally for the current policy, because the reward can ”forget” the signal that it gave to an earlier policy.
Revised: IRL methods commonly learn rewards which explain behavior locally for the current policy, because the reward can ”forget” the signal that it gave to an earlier policy.

Original: This makes rewards difﬁcult to optimize from scratch.
Revised: This makes rewards difﬁcult to optimize from scratch.

Original: To somewhat migitate this effect, we mix policy samples from the previous 20 iterations of training as negatives when training the discriminator.
Revised: To somewhat migitate this effect, we mix policy samples from the previous 20 iterations of training as negatives when training the discriminator.

Original: We use this strategy for both
Revised: We use this strategy for both

Original: AIRL and GAN-GCL.
Revised: AIRL and GAN-GCL.