Title: ADVERSARIAL IMITATION VIA VARIATIONAL INVERSE REINFORCEMENT LEARNING

Abstract: We consider a problem of learning the reward and policy from expert examples under unknown dynamics. Our proposed method builds on the framework of generative adversarial networks and introduces the empowerment-regularized maximum-entropy inverse reinforcement learning to learn near-optimal rewards and policies. Empowerment-based regularization prevents the policy from overfitting to expert demonstrations, which advantageously leads to more generalized behaviors that result in learning near-optimal rewards. Our method simultaneously learns empowerment through variational information maximization along with the reward and policy under the adversarial learning formulation. We evaluate our approach on various high-dimensional complex control tasks. We also test our learned rewards in challenging transfer learning problems where training and testing environments are made to be different from each other in terms of dynamics or structure. The results show that our proposed method not only learns nearoptimal rewards and policies that are matching expert behavior but also performs significantly better than state-of-the-art inverse reinforcement learning algorithms.

1. INTRODUCTION:

Reinforcement learning (RL) has emerged as a promising tool for solving complex decision-making and control tasks from predefined high-level reward functions (Sutton et al., 1998) . However, defining an optimizable reward function that inculcates the desired behavior can be challenging for many robotic applications, which include learning social-interaction skills (Qureshi et al., 2018; 2017) , dexterous manipulation (Finn et al., 2016b) , and autonomous driving (Kuderer et al., 2015) .

Inverse reinforcement learning (IRL) (Ng et al., 2000) addresses the problem of learning reward functions from expert demonstrations, and it is often considered as a branch of imitation learning (Argall et al., 2009) . The prior work in IRL includes maximum-margin (Abbeel & Ng, 2004; Ratliff et al., 2006) and maximum-entropy (Ziebart et al., 2008) formulations. Currently, maximum entropy (MaxEnt) IRL is a widely used approach towards IRL, and has been extended to use non-linear function approximators such as neural networks in scenarios with unknown dynamics by leveraging sampling-based techniques (Boularias et al., 2011; Finn et al., 2016b; Kalakrishnan et al., 2013) . However, designing the IRL algorithm is usually complicated as it requires, to some extent, hand engineering such as deciding domain-specific regularizers (Finn et al., 2016b) .

Rather than learning reward functions and solving the IRL problem, imitation learning (IL) learns a policy directly from expert demonstrations. Prior work addressed the IL problem through behavior cloning (BC), which learns a policy from expert trajectories using supervised learning (Pomerleau, 1991) . Although BC methods are simple solutions to IL, these methods require a large amount of data because of compounding errors induced by covariate shift (Ross et al., 2011) . To overcome BC limitations, a generative adversarial imitation learning (GAIL) algorithm (Ho & Ermon, 2016) was proposed. GAIL uses the formulation of Generative Adversarial Networks (GANs) (Goodfellow et al., 2014) , i.e., a generator-discriminator framework, where a generator is trained to generate expert-like trajectories while a discriminator is trained to distinguish between generated and expert trajectories. Although GAIL is highly effective and efficient framework, it does not recover transferable/portable reward functions along with the policies, thus narrowing its use cases to similar problem instances in similar environments. Reward function learning is ultimately preferable, if possible, over direct imitation learning as rewards are portable functions that represent the most basic and complete representation of agent intention, and can be re-optimized in new environments and new agents.

Reward learning is challenging as there can be many optimal policies explaining a set of demonstrations and many reward functions inducing an optimal policy (Ng et al., 2000; Ziebart et al., 2008) . Recently, an adversarial inverse reinforcement learning (AIRL) framework (Fu et al., 2017) , an extension of GAIL, was proposed that offers a solution to the former issue by exploiting the maximum entropy IRL method (Ziebart et al., 2008) whereas the latter issue is addressed through learning disentangled reward functions by modeling the reward as a function of state only instead of both state and action. However, AIRL fails to recover the ground truth reward when the ground truth reward is a function of both state and action. For example, the reward function in any locomotion or ambulation tasks contains a penalty term that discourages actions with large magnitudes. This need for action regularization is well known in optimal control literature and limits the use cases of a state-only reward function in most practical real-life applications. A more generalizable and useful approach would be to formulate reward as a function of both states and actions, which induces action-driven reward shaping that has been shown to play a vital role in quickly recovering the optimal policies (Ng et al., 1999) .

In this paper, we propose the empowerment-regularized adversarial inverse reinforcement learning (EAIRL) algorithm1 . Empowerment (Salge et al., 2014 ) is a mutual information-based theoretic measure, like state-or action-value functions, that assigns a value to a given state to quantify the extent to which an agent can influence its environment. Our method uses variational information maximization (Mohamed & Rezende, 2015) to learn empowerment in parallel to learning the reward and policy from expert data. Empowerment acts as a regularizer to policy updates to prevent overfitting the expert demonstrations, which in practice leads to learning robust rewards. Our experimentation shows that the proposed method recovers not only near-optimal policies but also recovers robust, transferable, disentangled, state-action based reward functions that are near-optimal. The results on reward learning also show that EAIRL outperforms several state-of-the-art IRL methods by recovering reward functions that leads to optimal, expert-matching behaviors. On policy learning, results demonstrate that policies learned through EAIRL perform comparably to GAIL and AIRL with non-disentangled (state-action) reward function but significantly outperform policies learned through AIRL with disentangled reward (state-only) and GAN interpretation of Guided Cost Learning (GAN-GCL) (Finn et al., 2016a) .

2. BACKGROUND:

We consider a Markov decision process (MDP) represented as a tuple (S, A, P, R, ρ 0 , γ) where S denotes the state-space, A denotes the action-space, P represents the transition probability distribution, i.e., P : S × A × S → [0, 1], R(s, a) corresponds to the reward function, ρ 0 is the initial state distribution ρ 0 : S → R, and γ ∈ (0, 1) is the discount factor. Let q(a|s, s ) be an inverse model that maps current state s ∈ S and next state s ∈ S to a distribution over actions A, i.e., q : S × S × A → [0, 1]. Let π be a stochastic policy that takes a state and outputs a distribution over actions such that π : S × A → [0, 1]. Let τ and τ E denote a set of trajectories, a sequence of state-action pairs (s 0 , a 0 , • • • s T , a T ), generated by a policy π and an expert policy π E , respectively, where T denotes the terminal time. Finally, let Φ(s) be a potential function that quantifies a utility of a given state s ∈ S, i.e., Φ : S → R. In our proposed work, we use an empowerment-based potential function Φ(•) to regularize policy update under MaxEnt-IRL framework. Therefore, the following sections provide a brief background on MaxEnt-IRL, adversarial reward and policy learning, and variational information-maximization approach to learn the empowerment.

2.1. MAXENT-IRL:

MaxEnt-IRL (Ziebart et al., 2008) models expert demonstrations as Boltzmann distribution using parametrized reward r ξ (τ ) as an energy function, i.e.,

EQUATION

where r ξ (τ ) = T t=0 r ξ (s t , a t ) is a commutative reward over given trajectory τ , parameterized by ξ, and Z is the partition function. In this framework, the demonstration trajectories are assumed to be sampled from an optimal policy π * , therefore, they get the highest likelihood whereas the suboptimal trajectories are less rewarding and hence, are generated with exponentially decaying probability. The main computational challenge in MaxEnt-IRL is to determine Z. The initial work in MaxEnt-IRL computed Z using dynamic programming (Ziebart et al., 2008) whereas modern approaches (Finn et al., 2016b; a; Fu et al., 2017) present importance sampling technique to approximate Z under unknown dynamics.

2.2. ADVERSARIAL INVERSE REINFORCEMENT LEARNING:

This section briefly describes Adversarial Inverse Reinforcement Learning (AIRL) (Fu et al., 2017) algorithm which forms a baseline of our proposed method. AIRL is the current state-of-the-art IRL method that builds on GAIL (Ho & Ermon, 2016) , maximum entropy IRL framework (Ziebart et al., 2008) and GAN-GCL, a GAN interpretation of Guided Cost Learning (Finn et al., 2016b; a) . GAIL is a model-free adversarial learning framework, inspired from GANs (Goodfellow et al., 2014) , where the policy π learns to imitate the expert policy behavior π E by minimizing the Jensen-Shannon divergence between the state-action distributions generated by π and the expert state-action distribution by π E through following objective

min π max D∈(0,1) S×A E π [log D(s, a)] + E π E [log(1 -D(s, a))] -λH(π) (2)

where D is the discriminator that performs the binary classification to distinguish between samples generated by π and π E , λ is a hyper-parameter, and H(π) is an entropy regularization term E π [log π]. Note that GAIL does not recover reward; however, Finn et al. (2016a) shows that the discriminator can be modeled as a reward function. Thus AIRL (Fu et al., 2017) presents a formal implementation of (Finn et al., 2016a) and extends GAIL to recover reward along with the policy by imposing a following structure on the discriminator:

EQUATION

where f ξ,ϕ (s, a, s ) = r ξ (s) + γh ϕ (s ) -h ϕ (s) comprises a disentangled reward term r ξ (s) with training parameters ξ, and a shaping term F = γh ϕ (s ) -h ϕ (s) with training parameters ϕ. The entire D ξ,ϕ (s, a, s ) is trained as a binary classifier to distinguish between expert demonstrations τ E and policy generated demonstrations τ . The policy is trained to maximize the discriminative reward r(s, a, s ) = log(D(s, a, s ) -log(1 -D(s, a, s ))). Note that the function F = γh ϕ (s ) -h ϕ (s) consists of free-parameters as no structure is imposed on h ϕ (•), and as mentioned in (Fu et al., 2017) , the reward function r ξ (•) and function F are tied upto a constant (γ -1)c, where c ∈ R; thus the impact of F , the shaping term, on the recovered reward r is quite limited and therefore, the benefits of reward shaping are not fully realized.

2.3. EMPOWERMENT AS MAXIMAL MUTUAL INFORMATION:

Mutual information (MI), an information-theoretic measure, quantifies the dependency between two random variables. In intrinsically-motivated reinforcement learning, a maximal of mutual information between a sequence of K actions a and the final state s reached after the execution of a, conditioned on current state s is often used as a measure of internal reward (Mohamed & Rezende, 2015) , known as Empowerment Φ(s), i.e.,

Φ(s) = max I(a, s |s) = max E p(s |a,s)w(a|s) log p(a, s |s) w(a|s)p(s |s) (4)

where p(s |a, s) is a K-step transition probability, w(a|s) is a distribution over a, and p(a, s |s) is a joint-distribution of K actions a and final state s2 . Intuitively, the empowerment Φ(s) of a state s quantifies an extent to which an agent can influence its future. Thus, maximizing empowerment induces an intrinsic motivation in the agent that enforces it to seek the states that have the highest number of future reachable states.

Empowerment, like value functions, is a potential function that has been previously used in reinforcement learning but its applications were limited to small-scale cases due to computational intractability of MI maximization in higher-dimensional problems. Recently, however, a scalable method (Mohamed & Rezende, 2015) was proposed that learns the empowerment through the moreefficient maximization of variational lower bound, which has been shown to be equivalent to maximizing MI (Agakov, 2004) . The lower bound was derived (for complete derivation see Appendix A.1) by representing MI in term of the difference in conditional entropies H(•) and utilizing the non-negativity property of KL-divergence, i.e.,

EQUATION

where H(a|s) = -E w(a|s) [log w(a|s)], H(a|s , s) = -E p(s |a,s)w(a|s) [log p(a|s , s)], q φ (•) is a variational distribution with parameters φ and w θ (•) is a distribution over actions with parameters θ. Finally, the lower bound in Eqn. 5 is maximized under the constraint H(a|s) < η (prevents divergence, see (Mohamed & Rezende, 2015) ) to compute empowerment as follow:

EQUATION

where β is η dependent temperature term.

Mohamed & Rezende (2015) also applied the principles of Expectation-Maximization (EM) (Agakov, 2004) to learn empowerment, i.e., alternatively maximizing Eqn. 6 with respect to w θ (a|s) and q φ (a|s , s). Given a set of training trajectories τ , the maximization of Eqn. 6 w.r.t q φ (•) is shown to be a supervised maximum log-likelihood problem whereas the maximization w.r.t w θ (•) is determined through the functional derivative ∂I/∂w = 0 under the constraint a w(a|s) = 1. The optimal w * that maximizes Eqn. 6 turns out to be 1 Z(s) exp(βE p(s |s,a) [log q φ (a|s, s )]), where Z(s) is a normalization term. Substituting w * in Eqn. 6 showed that the empowerment Φ(s) = 1 β log Z(s)

(for full derivation, see Appendix A.2).

Note that w * (a|s) is implicitly unnormalized as there is no direct mechanism for sampling actions or computing Z(s). Mohamed & Rezende (2015) introduced an approximation w * (a|s) ≈ log π(a|s) + Φ(s) where π(a|s) is a normalized distribution which leaves the scalar function Φ(s) to account for the normalization term log Z(s). Finally, the parameters of policy π and scalar function Φ are optimized by minimizing the discrepancy, l I (s, a, s ), between the two approximations (log π(a|s) + Φ(s)) and β log q φ (a|s , s)) through either absolute (p = 1) or squared error (p = 2), i.e., l I (s, a, s

) = β log q φ (a|s , s) -(log π θ (a|s) + Φ ϕ (s)) p (7)

3. EMPOWERED ADVERSARIAL INVERSE REINFORCEMENT LEARNING:

We present an inverse reinforcement learning algorithm that learns a robust, transferable reward function and policy from expert demonstrations. Our proposed method comprises (i) an inverse model q φ (a|s , s) that takes the current state s and the next state s to output a distribution over actions A that resulted in s to s transition, (ii) a reward r ξ (s, a), with parameters ξ, that is a function of both state and action, (iii) an empowerment-based potential function Φ ϕ (•) with parameters ϕ that determines the reward-shaping function F = γΦ ϕ (s ) -Φ ϕ (s) and also regularizes the policy update, and (iv) a policy model π θ (a|s) that outputs a distribution over actions given the current state s. All these models are trained simultaneously based on the objective functions described in the following sections to recover optimal policies and generalizable reward functions concurrently.

3.1 INVERSE MODEL q φ (a|s, s ) OPTIMIZATION As mentioned in Section 2.3, learning the inverse model q φ (a|s, s ) is a maximum log-likelihood supervised learning problem. Therefore, given a set of trajectories τ ∼ π, where a single trajectory is a sequence states and actions, i.e., τ i = {s 0 , a 0 , • • • , s T , a T } i , the inverse model q φ (a|s , s) is trained to minimize the mean-square error between its predicted action q(a|s , s) and the action a taken according to the generated trajectory τ , i.e.,

l q (s, a, s ) = (q φ (•|s, s ) -a) 2 (8) 3.2 EMPOWERMENT Φ ϕ (s) OPTIMIZATION

Empowerment will be expressed in terms of normalization function Z(s) of optimal w * (a|s), i.e.,

Φ ϕ (s) = 1 β log Z(s).

Therefore, the estimation of empowerment Φ ϕ (s) is approximated by minimizing the loss function l I (s, a, s ), presented in Eqn. 7, w.r.t parameters ϕ, and the inputs (s, a, s ) are sampled from the policy-generated trajectories τ .

3.3. REWARD FUNCTION r ξ (s, a):

To train the reward function, we first compute the discriminator as follow:

EQUATION

where r ξ (s, a) is the reward function to be learned with parameters ξ. We also maintain the target ϕ and learning ϕ parameters of the empowerment-based potential function. The target parameters ϕ are a replica of ϕ except that the target parameters ϕ are updated to learning parameters ϕ after every n training epochs. Note that keeping a stationary target Φ ϕ stabilizes the learning as also mentioned in (Mnih et al., 2015) . Finally, the discriminator/reward function parameters ξ are trained via binary logistic regression to discriminate between expert τ E and generated τ trajectories, i.e.,

EQUATION

3.4 POLICY OPTIMIZATION POLICY π θ (a|s)

We train our policy π θ (a|s) to maximize the discriminative reward r(s, a, s ) = log(D(s, a, s )log(1 -D(s, a, s ))) and to minimize the loss function l I (s, a, s ) = β log q φ (a|s, s ) -(log π θ (a|s) + Φ ϕ (s)) p which accounts for empowerment regularization. Hence, the overall policy training objective is:

EQUATION

) where policy parameters θ are updated using any policy optimization method such as TRPO (Schulman et al., 2015) or an approximated step such as PPO (Schulman et al., 2017) .

Algorithm 1 outlines the overall training procedure to train all function approximators simultaneously. Note that the expert samples τ E are seen by the discriminator only, whereas all other models are trained using the policy generated samples τ . Furthermore, the discriminating reward r(s, a, s ) boils down to the following expression (Appendix B.1): r(s, a, s ) = f (s, a, s ) -log π(a|s) where f (s, a, s ) = r ξ (s, a) + γΦ ϕ (s ) -Φ ϕ (s). Thus, an alternative way to express our policy training objective is E τ [log π θ (a|s)r π (s, a, s )], where r π (s, a, s ) = r(s, a, s ) -λ I l I (s, a, s ), Algorithm 1: Empowerment-based Adversarial Inverse Reinforcement Learning Initialize parameters of policy π θ , and inverse model q φ Initialize parameters of target Φ ϕ and training Φ ϕ empowerment, and reward r ξ functions Obtain expert demonstrations τ E by running expert policy π E for i ← 0 to N do Collect trajectories τ by executing π θ Update φ i to φ i+1 with the gradient E τ [ φi l q (s, a, s )] Update ϕ i to ϕ i+1 with the gradient E τ [ ϕi l I (s, a, s )] Update ξ i to ξ i+1 with the gradient:

E τ [ ξi log D ξi,ϕi+1 (s, a, s )] + E τ E [ ξi (1 -log D ξi,ϕi+1 (s, a, s ))]

Update θ i to θ i+1 using natural gradient update rule (i.e., TRPO/PPO) with the gradient:

E τ θi log π θi (a|s)r ξi+1 (s, a, s ) + λ I E τ θi l I (s, a, s )

After every n epochs sync ϕ with ϕ which would undoubtedly yield the same results as Eqn. 11, i.e., maximize the discriminative reward and minimize the loss l I . The analysis of this alternative expression is given in Appendix B to highlight that our policy update rule is equivalent to MaxEnt-IRL policy objective (Finn et al., 2016a) except that it also maximizes the empowerment, i.e.,

EQUATION

where, λ and γ are hyperparameters, and Ĥ(•) is the entropy-regularization term depending on π(•) and q(•). Hence, our policy is regularized by the empowerment which induces generalized behavior rather than locally overfitting to the limited expert demonstrations.

4. RESULTS:

Our proposed method, EAIRL, learns both reward and policy from expert demonstrations. Thus, for comparison, we evaluate our method against both state-of-the-art policy and reward learning techniques on several control tasks in OpenAI Gym. In case of policy learning, we compare our method against GAIL, GAN-GCL, AIRL with state-only reward, denoted as AIRL(s), and an augmented version of AIRL we implemented for the purposes of comparison that has state-action reward, denoted as AIRL(s, a). In reward learning, we only compare our method against AIRL(s) and AIRL(s, a) as GAIL does not recover rewards, and GAN-GCL is shown to exhibit inferior performance than AIRL (Fu et al., 2017) . Furthermore, in the comparisons, we also include the expert It can be seen that our method performs significantly better than AIRL (Fu et al., 2017) and exhibits expert-like performance in all five randomly-seeded trials which imply that our method learns near-optimal, transferable reward functions.

performances which represents a policy learned by optimizing a ground-truth reward using TRPO (Schulman et al., 2015) . The performance of different methods are evaluated in term of mean and standard deviation of total rewards accumulated (denoted as score) by an agent during the trial, and for each experiment, we run five randomly-seeded trials. To evaluate the learned rewards, we consider a transfer learning problem in which the testing environments are made to be different from the training environments. More precisely, the rewards learned via IRL in the training environments are used to re-optimize a new policy in the testing environment using standard RL. We consider two test cases shown in the Fig. 1 .

In the first test case, as shown in Fig. 1 (a), we modify the agent itself during testing. We trained a reward function to make a standard quadruped ant to run forward. During testing, we disabled the front two legs (indicated in red) of the ant (crippled-ant), and the learned reward is used to reoptimize the policy to make a crippled-ant move forward. Note that the crippled-ant cannot move sideways (Appendix C.1). Therefore, the agent has to change the gait to run forward. In the second test case, shown in Fig. 2 compares the policy performance scores over five different trials of EAIRL, AIRL(s) and AIRL(s, a) in the aforementioned transfer learning tasks. The expert score is shown as a horizontal line to indicate the standard set by an expert policy. Table 1 summarizes the means and standard deviations of the scores over five trials. It can be seen that our method recovers near-optimal reward functions as the policy scores almost reach the expert scores in all five trials even after transfering to unseen testing environments. Furthermore, our method performs significantly better than both Next, we considered the performance of the learned policy specifically for an imitation learning problem in various control tasks.The tasks, shown in Fig. 3 , include (i) making a 2D halfcheetah robot to run forward, (ii) making a 3D quadruped robot (ant) to move forward, (iii) making a 2D swimmer to swim, and (iv) keeping a friction less pendulum to stand vertically up. For each algorithm, we provided 20 expert demonstrations generated by a policy trained on a ground-truth reward using TRPO (Schulman et al., 2015) . Table 2 presents the means and standard deviations of policy learning performance scores, over the five different trials. It can be seen that EAIRL, AIRL(s, a) and GAIL demonstrate similar performance and successfully learn to imitate the expert policy, whereas AIRL(s) and GAN-GCL fails to recover a policy. 

5. DISCUSSION:

This section highlights the importance of empowerment-regularized MaxEnt-IRL and modeling rewards as a function of both state and action rather than restricting to state-only formulation on learning rewards and policies from expert demonstrations.

In the scalable MaxEnt-IRL framework (Finn et al., 2016a; Fu et al., 2017) , the normalization term is approximated by importance sampling where the importance-sampler/policy is trained to minimize the KL-divergence from the distribution over expert trajectories. However, merely minimizing the divergence between expert demonstrations and policy-generated samples leads to localized policy behavior which hinders learning generalized reward functions. In our proposed work, we regularize the policy update with empowerment i.e., we update our policy to reduce the divergence from expert data distribution as well as to maximize the empowerment (Eqn.12). The proposed regularization prevents premature convergence to local behavior which leads to robust state-action based rewards learning. Furthermore, empowerment quantifies the extent to which an agent can control/influence its environment in the given state. Thus the agent takes an action a on observing a state s such that it has maximum control/influence over the environment upon ending up in the future state s .

Our experimentation also shows the importance of modeling discriminator/reward functions as a function of both state and action in reward and policy learning under GANs framework. The re-ward learning results show that state-only rewards (AIRL(s)) does not recover the action dependent terms of the ground-truth reward function that penalizes high torques. Therefore, the agent shows aggressive behavior and sometimes flips over after few steps (see the accompanying video), which is also the reason that crippled-ant trained with AIRL's disentangled reward function reaches only the half-way to expert scores as shown in Table 1 . Therefore, the reward formulation as a function of both states and actions is crucial to learning action-dependent terms required in most real-world applications, including any autonomous driving, robot locomotion or manipulation task where large torque magnitudes are discouraged or are dangerous. The policy learning results further validate the importance of the state-action reward formulation. Table 2 shows that methods with state-action reward/discriminator formulation can successfully recover expert-like policies. Hence, our empirical results show that it is crucial to model reward/discriminator as a function of state-action as otherwise, adversarial imitation learning fails to learn ground-truth rewards and expert-like policies from expert data.

6. CONCLUSIONS AND FUTURE WORK:

We present an approach to adversarial reward and policy learning from expert demonstrations by regularizing the maximum-entropy inverse reinforcement learning through empowerment. Our method learns the empowerment through variational information maximization in parallel to learning the reward and policy. We show that our policy is trained to imitate the expert behavior as well to maximize the empowerment of the agent over the environment. The proposed regularization prevents premature convergence to local behavior and leads to a generalized policy that in turn guides the reward-learning process to recover near-optimal reward. We show that our method successfully learns near-optimal rewards, policies, and performs significantly better than state-of-the-art IRL methods in both imitation learning and challenging transfer learning problems. The learned rewards are shown to be transferable to environments that are dynamically or structurally different from training environments.

In our future work, we plan to extend our method to learn rewards and policies from diverse human/expert demonstrations as the proposed method assumes that a single expert generates the training data. Another exciting direction would be to build an algorithm that learns from sub-optimal demonstrations that contains both optimal and non-optimal behaviors.

APPENDICES A VARIATIONAL EMPOWERMENT:

For completeness, we present a derivation of presenting mutual information (MI) as variational lower bound and maximization of lower bound to learn empowerment.

A ≥ -E w(a|s) log w(a|s) + E p(s |a,s)w(a|s) [log q(a|s , s)]

A.2 VARIATIONAL INFORMATION MAXIMIZATION:

The empowerment is a maximal of MI and it can be formalized as follow by exploiting the variational lower bound formulation (for details see (Mohamed & Rezende, 2015) ).

Φ(s) = max:

w,q E p(s |a,s)w(a|s) [-1 β log w(a|s) + log q(a|s , s)] (13)

As mentioned in section 2.3, given a training trajectories, the maximization of Eqn. 13 w.r.t inverse model q(a|s , s) is a supervised maximum log-likelihood problem. The maximization of Eqn. 13 w.r.t w(a|s) is derived through a functional derivative ∂I w,q /∂w = 0 under the constraint a w(a|s) = 1. For simplicity, we consider discrete state and action spaces, and the derivation is as follow: By using the constraint a w(a|s) = 1, it can be shown that the optimal solution w * (a|s) = 1 Z(s) exp(u(s, a)), where u(s, a) = βE p(s |a,s) [log q(a|s , s)] and Z(s) = a u(s, a). This solution maximizes the lower bound since ∂ 2 I w (s)/∂w 2 =a 1 w(a|s) < 0.

Îw (s) = E p(

B EMPOWERMENT-REGULARIZED MAXENT-IRL FORMULATION.

In this section we derive the Empowerment-regularized formulation of maximum entropy IRL. Let τ be a trajectory sampled from expert demonstrations D and p ξ (τ ) ∝ p(s 0 )Π T -1 t=0 p(s t+1 |s t , a t ) exp r ξ (st,at) be a distribution over τ . As mentioned in Section 2, the IRL objective is to maximize the likelihood:

max ξ J(ξ) = max ξ E D [log p ξ (τ )]

Furthermore, as derived in (Fu et al., 2017) , the gradient of above equation w.r.t ξ can be written as:

max ξ J(ξ) = E D [ T t=0 ∂ ∂ξ r ξ (s t , a t )] -E p ξ [ T t=0 ∂ ∂ξ r ξ (s t , a t )] = T t=0 E D [ ∂ ∂ξ r ξ (s t , a t )] -E p ξ,t [ ∂ ∂ξ r ξ (s t , a t )]

where r ξ (•) is a parametrized reward to be learned, and p ξ,t = s t =t,a t =t p ξ (τ ) denotes marginalization of state-action at time t. Since, it is unfeasible to draw samples from p ξ , Finn et al. (2016a) proposed to train an importance sampling distribution µ(τ ) whose varience is reduced by defining µ(τ ) as a mixture of polices, i.e., µ(a|s) = 1 2 (π(a|s) + p(a|s)), where p is a rough density estimate over demonstrations. Thus the above gradient becomes:

EQUATION

We train our importance-sampler/policy π to maximize the empowerment Φ(•) for generalization and to reduce divergence from true distribution by minimizing D KL (π(τ ) p ξ (τ )). Since, π(τ ) = p(s 0 )Π T -1 t=0 p(s t+1 |s t , a t )π(s t , a t ), the matching terms of π(τ ) and p ξ (τ ) cancel out, resulting into entropy-regularized policy update. Furthermore, as we also include the empowerment Φ(•) in the policy update to be maximized, hence the overall objective becomes:

EQUATION

Our discriminator is trained to minimize cross entropy loss as mention in Eqn. 10, and for the proposed structure of our discriminator Eqn. 9, it can be shown that the discriminator's gradient w.r.t its parameters turns out to be equal to Equation 14 (for more details, see (Fu et al., 2017) ). On the other hand, our policy training objective is

EQUATION

In the next section, we show that the above policy training objective is equivalent to Equation 15.

B.1 POLICY OBJECTIVE:

We train our policy to maximize the discriminative reward r(s, a, s ) = log(D(s, a, s ) -log(1 -D(s, a, s ))) and minimize the information-theoretic loss function l I (s, a, s ). The discriminative reward r(s, a, s ) simplifies to:

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

= log e f (s,a,s ) e f (s,a,s ) + π(a|s)

-log π(a|s) e f (s,a,s ) + π(a|s) = f (s, a, s ) -log π(a|s) where f (s, a, s ) = r(s, a) + γΦ(s ) -Φ(s). The entropy-regularization is usually scaled by the hyperparameter, let say λ h ∈ R, thus r(s, a, s ) = f (s, a, s ) -λ h log π(a|s). Hence, assuming single-sample (s, a, s ), absolute-error for l I (s, a, s ) = | log q φ (a|s, s ) -(log π(a|s) + Φ(s))|, and l i > 0, the policy is trained to maximize following: r π (s, a, s ) = f (s, a, s ) -λ h log π(a|s) -l I (s, a, s ) = r(s, a) + γΦ(s ) -Φ(s) -λ h log π(a|s) -log q(a|s, s ) + log π(a|s) + Φ(s) = r(s, a) + γΦ(s ) -λ h log π(a|s) -log q(a|s, s ) + log π(a|s) Note that, the potential function Φ(s) cancels out and we scale the leftover terms of l I with a hyperparameter λ I . Hence, the above equation becomes: r π (s, a, s ) = r(s, a, s ) + γΦ(s ) + (λ I -λ h ) log π(a|s) -λ I log q(a|s, s )

We combine the log terms together as: r π (s, a, s ) = r(s, a)

EQUATION

)

where λ is a hyperparameter, and Ĥ(•) is an entropy regularization term depending on q(a|s, s ) and π(a|s). Therefore, it can be seen that the Eqn. 17 is equivalent/approximation to Eqn. 15.

C TRANSFER LEARNING PROBLEMS:

C.1 ANT ENVIRONMENT

The following figures show the difference between the path profiles of standard and crippled Ant. It can be seen that the standard Ant can move sideways whereas the crippled ant has to rotate in order to move forward. all presented models and the inverse model q φ (•) of EAIRL are presented by two-layer RELU network with 32 units in each layer, where the network's output parametrizes the Gaussian distribution, i.e., we assume a Gaussian policy.

D.2 HYPERPARAMETERS:

For all experiments, we use the temperature term β = 1. We evaluated both mean-squared and absolute error forms of l I (s, a, s ) and found that both lead to similar performance in reward and policy learning. We set entropy regularization weight to 0.1 and 0.001 for reward and policy learning, respectively. The hyperparameter λ I was set to 1.0 for reward learning and 0.001 for policy learning. The target parameters of the empowerment-based potential function Φ ϕ (•) were updated every 5 and 2 epochs during reward and policy learning respectively. Although reward learning hyperparameters are also applicable to policy learning, we decrease the magnitude of entropy and information regularizers during policy learning to speed up the policy convergence to optimal values. Furthermore, we set the batch size to 2000-and 20000-steps per TRPO update for the pendulum and remaining environments, respectively. For the methods (Fu et al., 2017; Ho & Ermon, 2016) presented for comparison, we use their suggested hyperparameters. We also use policy samples from previous 20 iterations as negative data to train the discriminator of all IRL methods presented in this paper to prevent the parametrized reward functions from overfitting the current policy samples.

:

Supplementary material is available at https://sites.google.com/view/eairl

In our proposed work, we consider only immediate step transitions i.e., K = 1, hence variables s, a and s will be represented in non-bold notations.

Figures and Tables:

Figure 1: Transfer learning problems. Fig. (a) represents a problem where agent dynamics are modified during testing, i.e., a reward learned on a quadruped-ant (left) is transferred to a crippledant (right). Fig (b) represents a problem where environment structure is modified during testing, i.e., a reward learned on a maze with left-passage is transferred to a maze with right-passage to the goal (green).

Figure2: The performance of policies obtained from maximizing the learned rewards in the transfer learning problems. It can be seen that our method performs significantly better than AIRL(Fu et al., 2017) and exhibits expert-like performance in all five randomly-seeded trials which imply that our method learns near-optimal, transferable reward functions.

Fig 1(b), we change the environment structure. The agent learns to navigate a 2D point-mass to the goal region in a simple maze. We re-position the maze central-wall during testing so that the agent has to take a different path, compared to the training environment, to reach the target (Appendix C.2).

Figure 3: Benchmark control tasks for imitation learning

s |a,s)w(a|s) [-1 β log w(a|s) + log q(a|s , s)] β) -log w(a|s) + βE p(s |a,s) [log q(a|s , s)]} = 0 w(a|s) = e λ-β e βE p(s |a,s) [log q(a|s ,s)]

Figure 4: The top and bottom rows show the gait of standard and crippled ant, respectively.

Figure 5: The top and bottom rows show the path followed by a 2D point-mass agent (yellow) to reach the target (green) in training and testing environment, respectively.

The evaluation of reward learning on transfer learning tasks. Mean scores (higher the better) with standard deviation are presented over 5 trials.

Table content: <table><tr><td>Algorithm</td><td colspan="2">States-Only Pointmass-Maze</td><td>Crippled-Ant</td></tr><tr><td>Expert</td><td>N/A</td><td>-4.98 ± 0.29</td><td>432.66 ± 14.38</td></tr><tr><td>AIRL</td><td>Yes</td><td>-8.07 ± 0.50</td><td>175.51 ± 27.31</td></tr><tr><td>AIRL</td><td>No</td><td>-19.28 ± 2.03</td><td>46.12 ± 14.37</td></tr><tr><td>EAIRL(Ours)</td><td>No</td><td colspan="2">-7.01 ± 0.61 348.43 ± 43.17</td></tr><tr><td colspan="4">4.1 REWARD LEARNING PERFORMANCE (TRANSFER LEARNING EXPERIMENTS)</td></tr></table>

The evaluation of imitation learning on benchmark control tasks. Mean scores (higher the better) with standard deviation are presented over 5 trials for each method.

Table content: <table><tr><td>Methods</td><td>HalfCheetah</td><td>Ant</td><td colspan="2">Environments Swimmer</td><td>Pendulum</td></tr><tr><td>Expert</td><td>2139.83 ± 30.22</td><td colspan="2">935.12 ± 10.94</td><td>76.21 ± 1.79</td><td>-100.11 ± 1.32</td></tr><tr><td>GAIL</td><td>1880.05 ± 15.72</td><td colspan="2">738.72 ± 9.49</td><td>50.21 ± 0.26</td><td>-116.01 ± 5.45</td></tr><tr><td>GCL</td><td>-189.90 ± 44.42</td><td colspan="2">16.74 ± 36.59</td><td>15.75 ± 7.32</td><td>-578.18 ± 72.84</td></tr><tr><td>AIRL(s,a)</td><td>1826.26 ± 19.64</td><td colspan="2">645.90 ± 41.75</td><td>49.52 ± 0.48</td><td>-118.13 ± 11.33</td></tr><tr><td>AIRL(s)</td><td>121.10 ± 42.31</td><td colspan="2">271.31 ± 9.35</td><td>33.21 ± 2.40</td><td>-134.82 ± 10.89</td></tr><tr><td>EAIRL</td><td colspan="5">1870.10 ± 17.86 641.12 ± 25.92 49.55 ± 0.29 -116.26 ± 8.313</td></tr></table>

.1 VARIATIONAL INFORMATION LOWER BOUND As mentioned in section 2.3, the variational lower bound representation of MI is computed by defining MI as a difference in conditional entropies, and the derivation is formalized as follow.

Table content: <table><tr><td>I w,q (s) = H(a|s) -H(a|s , s)</td><td/></tr><tr><td colspan="2">= H(a|s) + E p(s |a,s)w(a|s) [log p(a|s , s)]</td></tr><tr><td>= H(a|s) + E p(s |a,s)w(a|s) [log</td><td>p(a|s , s)q(a|s , s) q(a|s , s) ]</td></tr><tr><td colspan="2">= H(a|s) + E p(s |a,s)w(a|s) [log q(a|s , s)] + E p(s |a,s)w(a|s) [log</td><td>p(a|s , s) q(a|s , s) ]</td></tr><tr><td colspan="3">= H(a|s) + E p(s |a,s)w(a|s) [log q(a|s , s)] + KL[p(a|s , s)||q(a|s , s)]</td></tr><tr><td colspan="2">≥ H(a|s) + E p(s |a,s)w(a|s) [log q(a|s , s)]</td></tr></table>

Peer review (iclr2019):
 The paper proposes a method for inverse reinforcement learning based on AIRL. It*s main contribution is that the shaping function is not learned while training the discriminator, but separately as an approximation of the empowerment (maximum mutual information). This shaping term aims to learn disentangled rewards without being restricted to learning state-only reward functions, which is a major restriction of AIRL. The main weakness of the paper is, that it does not justify or motivate the main deviations compared to AIRL. The new objective for updating the policy is especially problematic because it does no longer correspond to the RL objective but includes an additional term that biases the policy towards actions that increase its empowerment. Although both terms of the update can be derived independently from an IRL and Empowerment perspective respectively, optimizing the sum was not derived from a common problem formulation. By combining these objectives, the learned reward function may lead to policies that fail to match the expert demonstration without such bias. This does not imply that the approach is not sound per se, however, simply presenting such update without any discussion is insufficient--especially given that it constitutes the main novelty of the approach. I think the paper would be much stronger if the update was derived from an empowerment-regularized IRL formulation. And even then, the implications of such bias/regularization would need to be properly discussed and evaluated, in particular with respect to the trade-off lambda, which--again--is hardly mentioned in the submission. I*m also not sure if the story of the paper works out; when we simply want to use empowerment as shaping term, why not use two separate policies for computing the empowerment and reward function respectively. Is the bias in the policy update maybe more important than the shaping term in the discriminator update for learning disentangled rewards? Keeping these issues aside, I actually like the paper. It tackles the main drawback of AIRL and the idea seems quite nice. Having a reward function that does not actively induce actions that can be explained by empowerment, may not always be appropriate, but often enough it may be a sensible approach to get more generalizable reward functions. The paper is also well written with few typos. The parts that are discussed are clear and the experimental results seem fine as well (although more experiments on the reward transfer would be nice). Minor notes: I think there is a sign error in the policy update Typo in the theorem, grantee should be guarantee Question: Please confirm that the reward transfer was learned with a standard RL formulation. Does the learned policy change, when we use the empowerment objective as well? Update (22.11) I think that the revised version is much better than the original submission because it now correctly attributes the improved generalization to an inductive bias in the policy update. However, the submission still seems borderline to me. - The proposed method uses the empowerment both for regularization as well as for reward shaping, but it is not clear whether the latter improves generalization. If the reward shaping was not necessary, it would be cleaner to use empowerment only for regularization. If the reward shaping is beneficial, this should be shown in an ablative experiment. - The benefit of using empowerment (whether for reward shaping or for regularization) should be discussed. Empowerment for generalization is currently hardly motivated. - The derivation could be a bit more rigorous. As the presentation is now much more sound, I slightly increased my rating.