LEARNING INVARIANT REWARD FUNCTIONS
Year: 2021

Abstract:
Inverse reinforcement learning methods aim to retrieve the reward function of a Markov decision process based on a data set of expert demonstrations. The commonplace scarcity of such demonstrations potentially leads to the absorption of spurious correlations in the data by the learning model. Consequently, this adaptation often exhibits behavioural overfitting to the expert data set when trained on the obtained reward function. The invariant risk minimization principle provides us with a novel regularization approach for the maximum entropy inverse reinforcement learning problem. By applying this regularization to both exact and approximate formulations of the learning task, we recover superior reward functions across domains in the transfer setting and, as a result, induce better performing policies than the empirical risk minimization baseline.

heading: 1 INTRODUCTION
text: Recent advancements in reinforcement learning methods applied to real world scenarios have revealed a major limitation, namely how to design a problem specific reward function. This design challenge typically requires a cumbersome and error-prone process of handcrafting a heuristic function to account for all the intricacies of the task at hand. Eliciting the correct behaviour via the optimization of a reward function is of paramount importance both for training artificial agents, such as robots, as well as human agents, for instance, in the case of teaching skills in a simulated environment, but modelling rewards for behavior learning remains an open challenge in a large number of domains.
Inverse reinforcement learning (IRL) solves the problem of inferring the reward function of a Markov Decision Process (MDP) based on a dataset of temporal behaviors. Such trajectories obtained from an agent are assumed to demonstrate near-optimal performance in the respective MDP. The problem to recover the reward from the statistics of the expert trajectories is ill-posed, as there typically exist many reward functions which satisfy the optimization constraints (Ng et al., 1999). Maximum entropy inverse reinforcement learning resolves this ambiguity in a unique way. Typically, the amount of expert demonstrations used to infer the reward is limited, which poses the danger of overfitting to the noise in the expert demonstrations. In particular, high-dimensional, parameterized reward models without appropriate regularization tend to absorb spurious correlations present in the trajectory data which do not generalize outside of the training setup. While optimizing the empirical cumulative reward obtained by such models, such an adaptation causes the agent to learn behaviours which frequently fail at test time. This type of observational overfitting has been shown in the context of reinforcement learning (Song et al., 2019). An additional difficulty arises when there is a significant discrepancy between expert demonstrations originating from different experts. Adversarial imitation methods fail to recover accurate behaviour from diverse experts as shown in (Li et al., 2017), since such learning results in ambiguities when inferring reward functions. We require to learn reward functions that elicit succinct behaviours of agents. Such rewards should avoid adopting behavioural features of experts that might be irrelevant to the task or even detrimental as a solution strategy.
Invariant risk minimization (IRM) (Arjovsky et al., 2019) exploits the concept of invariant causal prediction (Peters et al., 2015; Heinze-Deml et al., 2017). This recently introduced approach studies the generalization problem of classification models through the lens of causality. The IRM method postulates that the conditional distribution of the class label must be stable across datasets, in order to avoid absorption of spurious correlations into the model predictions. This adaptation principle can be utilized in the context of reward function learning, where we aim to elicit behavioural policies
without exploiting spurious reward features, i.e. features which cause nontransferable behaviours. To achieve this goal, we make the assumption that variations in expert demonstrations are a product of causal interventions on the data generating process of trajectories. Our contributions are as follows:
• We formulate the assumption that the variations between experts performing optimally on the same task can be seen as interventional settings of the underlying trajectory distribution.
• We propose a regularization principle based on invariant risk minimization. This modelling choice allows us to learn reward functions that are invariant to spurious correlations present in the expert data.
• We demonstrate the efficacy of this approach in both tractable, finite state-spaces, which we refer to as the exact setting as well as large continuous state-spaces, which we denote as approximate setting.
The rest of the paper is structured as follows: we provide an overview of related work in Section 2 and introduce the problem setting from the perspective of IRL as well as causal considerations in Section 3. We then present our model for the exact and approximate settings in Section 4 and demonstrate the experimental results in Section 5. We conclude with a brief discussion and outlook in Section 6.

heading: 2 RELATED WORK
text: IRL and imitation learning Maximum entropy inverse reinforcement learning has been introduced in (Ziebart et al., 2008). An extension to functions parameterized by deep neural networks has been presented in (Wulfmeier et al., 2015). A formulation which considers large continuous state spaces where an approximation of the partition functions is required has initially been studied in (Finn et al., 2016) and expanded to an imitation setting in the work by (Ho & Ermon, 2016). A disentangled version of the adversarial imitation approach which allows to recover rewards has been presented in (Fu et al., 2017). The authors of (Zolna et al., 2019) propose a model which also tackles the issue of spurious correlations being absorbed from expert data. However, their focus is on visual features in a solely imitation learning setting as opposed to our approach, which recovers reward functions that perform favorably in a transfer setting.
Invariant representation learning Invariant risk minimization is a method proposed in (Arjovsky et al., 2019) which builds upon the invariant causal prediction (Peters et al., 2015; Heinze-Deml et al., 2017) principle and proposes a way to tackle the issue of spurious correlations impeding out-of-distribution generalization of classification algorithms. Various other formulations of the IRM principle have subsequently been proposed. IRM games (Ahuja et al., 2020) reformulates the IRM objective from a game theoretic perspective and achieves lower variance compared to the original IRM method. Invariant rationalization proposes an adversarial method for recovering invariant rationales in the context of natural language processing (Chang et al., 2020) by employing an auxiliary model which is trained to be agnostic to interventional settings.
Invariance and causality in RL The concept of invariance has been used in a number of works in the reinforcement learning domain. Invariant causal prediction has been utilized in (Zhang et al., 2020) to learn model invariant state abstractions in a multiple MDP setting with a shared latent space. Invariant policy optimization (Sonar et al., 2021) uses the IRM games (Ahuja et al., 2020) formulation to learn policies invariant to certain domain variations. The authors of (de Haan et al., 2019) tackle the problem of causal confusion in imitation learning by making use of causal structure of demonstrations. (Angelov et al., 2020) address the problem of user specification in robotic learning through the causal lens by augmenting the images with specific symbols. To the best of our knowledge, our algorithm is the first proposed method to use invariant causal prediction in the context of inverse reinforcement learning.

heading: 3 PROBLEM SETTING
text: MDP We consider environments modelled by a Markov decision processM = (S,A, T , P0, R), where S is the state space, A is the action space, T is the family of transition distributions on S indexed by S ×A with p(s′|s, a) describing the probability of transitioning to state s′ when taking
action a in state s, P0 is the initial state distribution, and R : S × A → R is the reward function. A policy π is a conditional probability distribution of actions a ∈ A given states s ∈ S with state descriptors or features ϕ(s).
Inverse reinforcement learning aims to estimate a suitable reward function rψ parameterized by weights ψ based on a dataset of expert trajectories DE = {τi}i≤K where τi = (s(i)1:T , a (i) 1:T ) is a sequence of states and actions of expert i of length T . To achieve this goal, the feature expectation statistics under the state occupancy measure of the expert, Eτ∼DE [ϕ(τ)], are matched with the statistics of the student state occupancy measure Eτ∼p(τ |ψ)[ϕ(τ)]. ϕ(τ) = ∑ s∈τ ϕ(s) denotes the sum of the state features ϕ(s). The trajectory distribution p(τ |ψ) is induced by the policy πrψ trained on the reward function estimate rψ. Given the ill-posed nature of the problem, a maximum entropy formulation of the IRL problem (Ziebart et al., 2008) is typically chosen in order to find a unique solution. The model describes the fact that the expert trajectories are sampled from the Gibbs distribution p(τ |ψ,ϕ) = 1Zϕ,ψ exp(rψ(τ)) = 1 Zϕ,ψ
exp(ψTϕ(τ)) which corresponds to the solution of the entropy maximization problem under feature matching and simplex constraints.
max ψ
H(p(τ |ψ)) s.t. { Eτ∼p(τ |ψ)[ϕ(τ)] = Eτ∼DE [ϕ(τ)]; ∫ T p(τ |ψ)dτ = 1; p(τ |ψ) > 0 } (1)
where Eτ∼DE [ϕ(τ)] is the sample average of the expert trajectories. In the feature matching case, the reward rψ(τ) is typically assumed to be linear in the state features rψ = ψTϕ(τ). We further consider the scenario where the state features ϕ(τ) are learned using the DeepMaxEnt model (Wulfmeier et al., 2015). In large state spaces, the partition function Z presents an estimation challenge and an importance sampling scheme is typically employed.
Adversarial IRL formulation It has been shown (Finn et al., 2016) that the trajectory distribution for the importance sampling scheme can be obtained by solving an adversarial minimax game implemented by a generative adversarial network (Goodfellow et al., 2014) architecture. Adversarial inverse RL (AIRL) methods yield a reward function by learning to distinguish between the transitions sampled from the dataset of expert trajectories τ ∼ DE and transitions continually sampled from the improving policy τ ∼ π(τ), which maximizes an expected cumulative reward objective based on the discriminator output. We adopt the discriminator structure from (Fu et al., 2017) which separates into the state-action dependent reward function gξ and the state dependent shaping term hφ.
Dξ,φ,θ(s, a, s ′) =
exp(gξ(s, a) + γhφ(s ′)− hφ(s))
exp(gξ(s, a) + γhφ(s′)− hφ(s)) + πθ(a|s) (2)
We also consider a state-only formulation of AIRL where the reward function gξ(s) only depends on the state s. The optimization objective of the discriminator is the binary cross-entropy loss on the discriminator output:
LBCE(ξ, φ, θ) = E(s,a)∼πθ log(Dξ,φ,θ(s, a))− E(s,a)∼πE log(1−Dξ,φ,θ(s, a)) (3)
The generator corresponds to the policy optimizing the expected cumulative reward given by the discriminator output and is trained using the proximal policy optimization (PPO) (Schulman et al., 2017) method.
As with many neural network approximators, the discriminator model absorbs spurious correlations and this learning effect poses a serious problem (Arjovsky et al., 2019). These correlations coincide with the binary label information encoding the optimality of the expert. In this work, we avoid this effect by applying the invariant risk minimization principle to the discriminator part of our model.
Invariant risk minimization (Arjovsky et al., 2019) extends the invariant causal prediction (Peters et al., 2015; Heinze-Deml et al., 2017) principle to nonlinear settings where the underlying SCM might not be explicitly given. The goal of IRM is to learn correlations which are invariant across training settings and ignore spurious correlations. The set of training settings Etr contains datasets De :=
{(x(i)e , y(i)e )}|Etr|i=1 sampled from interventional distributions P̃ (Xe, Ye) of the SCM of the underlying data generating process. The IRM method effectively learns an invariant data representation ϕ for the predictor w ◦ ϕ (where w is a linear classifier) which enables a stable conditional distribution P (Y e|Xe) by optimizing the objective stated in the definition below:
Definition 1 The data representation ϕ : X → H mapping the input to a hypothesis space, elicits an invariant predictor w ◦ ϕ across the set of interventional settings E if there is a classifier w : H → Y simultaneously optimal for all settings:
min ϕ:X→H,w:H→Y ∑ e∈Etr Le(w ◦ ϕ) s. t.: w ∈ argmin w̄:H→Y Le(w̄ ◦ ϕ) ∀e ∈ Etr (4)
The aim of this bi-level optimization problem is to minimize the prediction loss on the union of the training sets constrained on selecting the predictor which minimizes the losses on the individual training set e, Le. The tractable approximation of Eq. 4 uses the gradient norm penalty D(w = 1.0, ϕ, e) = ||∇w|w=1.0Le(w ◦ ϕ)||2 which quantifies the violation of the normal equations to measure the optimality of a fixed classifier (w = 1.0) at each setting e. This leads to the following unconstrained formulation of the problem:
min ϕ:X→Y ∑ e∈Etr Le(ϕ) + λ||∇w|w=1.0Le(w ◦ ϕ)||2 (5)
In the following section, we will show how to incorporate this optimization principle into the IRL setting.

heading: 4 MODEL
text: In this section, we present the invariance regularization for the maximum entropy IRL models in exact and approximate settings which allows us to eliminate spurious correlations present in the expert datasets. We assume that different experts stem from interventional settings of the trajectory SCM. We first present the types of interventions on the trajectory SCM that we consider in this work. In the second subsection, we show how to apply this principle in a feature matching scenario. In the final section, we propose a regularization method for the adversarial learning scenario, which penalizes the discriminator function in order to exclude spurious correlations when distinguishing between policy and expert transitions.

heading: 4.1 INTERVENTIONS
text: The principle of invariant risk minimization relies on the assumption that the difference in training settings arises as a result of an intervention on the data generating process described by a structural causal model. In the case of maximum entropy IRL, the generative model is the SCM of trajectories p(τ |O1:T ) conditioned on the optimality variable O1:T for timesteps t = {1, ..., T} factorizes into the initial state distribution p0(s1), the conditional distribution p(Ot|st,at) = exp(rψ(st,at)) of the binary optimality variable Ot at timestep t and the MDP transition dynamics p(st+1|st,at) (Levine, 2018) is defined by the following expression:
p(τ |O1:T ) ∝ p(τ,O1:T ) = p0(s1) T∏ t=1 p(Ot = 1|st,at)p(st+1|st,at) (6)
Assumption 1 Different expert demonstrations stem from interventional settings of the trajectory SCM.
The assumption is valid in a scenario where expert demonstrations were gathered on different dynamics. For instance, recordings of surgeons performing a procedure on different patients with variations of anatomy would satisfy this assumption.
In order to apply the IRM principle to the inverse reinforcement learning problem, we must obtain a set of settings E which all contain the relation of interest invariant with respect to these settings. For the trajectory model state above, the causal relation of interest in the context of reward function inference is the conditional P (Ot|st, at), which should be invariant under interventions to the transition SCM (Fig. 1a). We motivate this by the fact that despite the discrepancies in the demonstrations, all experts are assumed to perform the task in an optimal fashion. This implies that all experts optimize the same underlying reward that we would like to recover. By considering the conditional independence structure in the graphical model of the state-action transition in (Fig. 1a) we can observe the underlying assumption that the causal mechanism producing the optimality label Ot is invariant across variations of the setting index variable E provided that the latent variable W remains unobserved.
Spurious correlations in model of transitions We will now describe the two scenarios in which a non-causal information path corresponding to spurious correlations is formed in the transition SCM. In Fig. 1b we can observe the scenario where we do not condition our reward function reprentation r(s) on the action. By conditioning on the collider node st+1 and not observing the action node a, a path is formed between the setting index E and the optimality variable Ot, resulting in the violation of their conditional independence relationship. A second scenario can be observed in Fig. 1c. This scenario requires the assumption that the orientation of the edge from node Ot to node st+1 is temporally causal, meaning that the optimality of a state at time t is a causal parent of the next state. In this case, observing the collider node st+1 implies the following conditional independence relationship: E 6⊥⊥ Ot|st+1. Beyond these scenarios, one can further assume a more general partitioning of an arbitratry transition input (s, a, s′) into the causal transition feature components x(c) and x(nc) : (s, a, s′) = (x(c),x(nc)), illustrated in Fig 1d, whereby conditioning on the x(nc) collider introduces a spurious correlation path (Huszar, 2019).
The first scenario arises in the feature matching case where the trajectory feature representation ϕ(τ) learned using the DeepMaxEnt (Wulfmeier et al., 2015) model only has access to the subsequent state but not the action observations. The second scenario arises in the adversarial learning case where the discriminator Dξ,φ,θ(s, a, s′) 2 depends on the node s′, hence observing a collider.
Following the assumption of the IRM principle that we only make interventions on non-causal parents of the optimality variable Ot, possible interventions are state interventions caused by altering the initial state distribution or the dynamics of the MDP and the action interventions, caused by intervening on the agent policy.

heading: 4.2 MAXIMUM ENTROPY IRL REGULARIZATION
text: For the application of the IRM regularization, we first consider the maximum entropy feature matching scenario. For the Gibbs distribution p(τ |ψ,ϕ) = exp(ψTϕ(τ))/Zϕ,ψ over trajectories, we can write down the constrained optimization problem analogously to Eq. 4:
min ϕ,ψ ∑ e∈Etr ∑ τ∈De log p(τ |ψ,ϕ) s.t. ψ ∈ argmin ψ̄ ∑ τ∈De log p(τ |ψ̄, ϕ) (7)
In simple settings such as gridworlds, where the computation of the partition function is tractable, we propose the following regularization approach for the MaxEnt maximum likelihood objective:
min ϕ ∑ e∈Etr (∑ τ∈De log ( 1 Zψ,φ exp(ψTϕ(τ)) ) + λD(ψ,ϕ, e) ) (8)
In an analogous fashion to the invariant risk minimization approach, D(ψ,ϕ, e) is a distance function representing the violation of the constraints w.r.t. the optimal solution, which in our case corresponds to the parameters maximizing the likelihood of the Gibbs distribution. The gradient of the loglikelihood loss LMLE = ∑ τ∈De log p(τ |ψ,ϕ) w.r.t. to ψ is then the difference of the feature expectations EDE [ϕ(τ)]− Ep(τ |ψ)[ϕ(τ))] (Ziebart et al., 2008). The squared norm of this gradient constitutes the equivalent of the IRM penalty in the maximum entropy IRL case.
D(ψ,ϕ; e) = ||∇ψ|ψ=1.0LMLE(ψ,ϕ)||2 = ||EDE [ϕ(τ)]− Ep(τ |ψ)[ϕ(τ))]||2 (9)
This closed form of the gradient norm penalty can be utilized in the maximum causal entropy solver (Ziebart et al., 2010). We assume the state features to be the output of a neural network ϕθ(s) according to the DeepMaxEnt model (Wulfmeier et al., 2015). The gradient norm penalty 9 is then backpropagated to enforce invariant features ϕ.

heading: 4.3 ADVERSARIAL IRL REGULARIZATION
text: We now present the adversarial version of the invariant regularization in algorithm 1. The training procedure of alternating policy updates using a policy gradient method with discriminator updates is similar to (Fu et al., 2017). There are two main differences compared to the baseline algorithm. The first is the fact that we use multiple experts in a distinct fashion as opposed to pooling the demonstrations into one big dataset. The second is the regularization of the discriminator objective (Eq. 3) using the gradient norm penalty D(ξ, φ, ω; e) = ||∇ω|ω=1.0LBCE(ξ, φ, ω; e)||2 in a similar fashion to Eq. 4, where ω = 1.0 corresponds to a fixed scalar classifier. The regularized discriminator effectively iterates over tuples (DeE ,Dπ) ∀e ∈ Etr during adversarial training, where Dπ is dataset of transitions generated by policy π.
Algorithm 1: IRM regularized adversarial IRL Input: Expert trajectories DeE assumed to be obtained by intervening on p(τ) in settings e; Result: Reward rξ, trained student policy πθ Initialize policy πθ and discriminator Dξ,φ; for setting e in {1, ..., Etr} do
for step t in {1, ..., N} do Collect trajectory buffer Dπ = {τi}i≤|Dπ| by executing the policy πθ; Update Dξ,φ,θ(s, a) via binary logistic regression by maximizing;
L(ξ, φ, ω; e) = LBCE(ξ, φ, ω; e) + λ||∇ω|ω=1.0LBCE(ξ, φ, ω; e)||2
using expert-specific tuple (DeE ,Dπ) Update policy πθ w.r.t. rψ,φ = logDξ,φ(s, a, s′)− log(1−Dξ,φ) function of IRM
regularized discriminator using a policy gradient method (e.g. PPO); end
end
The training process yields two outputs: the trained reward component rξ of the discriminator model as well as the trained student policy. In our experiments we show that the policy obtained through the adversarial training procedure shows improved zero-shot generalization performance on environments with dynamics sampled outside of the training distribution compared to the baseline. The performance is measured using the ground truth reward. Furthermore, we show that rewards obtained using this procedure elicit better performing policies when trained from random initialization on environments with modified dynamics.

heading: 5 EXPERIMENTS
text: The experiments are designed to answer the following questions:
• Can the MaxEnt IRL setting benefit from applying the invariance principle? • Is the IRM principle helpful for adversarial reward learning to infer invariant discriminators? • Do causal assumptions improve robustness of reward functions?
To answer these questions, we evaluate our model in three settings. The first setting considers a grid-world scenario, where the partition function is tractable and the reward is recovered using our modification of the maximum entropy feature expectation matching algorithm. In the second setting, we test the invariance regularization in an adversarial setting on simulated robotic locomotion environments. Finally, we demonstrate generalization of the obtained reward functions by retraining policies on the reward functions obtained using the adversarial methods.
Throughout this section, we compare the empirical risk minimization (ERM) baseline where the trajectory datasets gathered through interventions are pooled together into one dataset to the IRM version of the algorithms where we regularize either the feature expectation matching algorithm (sec. 5.1) or the adversarial formulation by assuming interventions on the expert trajectories (sec. 5.2, 5.3) as described in section 4.

heading: 5.1 TRACTABLE SETTING: GRIDWORLD EXPERIMENTS USING FEATURE MATCHING
text: To illustrate the principle of invariant risk minimization in the tractable IRL setting as discussed in 4.2, we choose a simple gridworld environment illustrated in (Fig. 2a). The goal of the agent is to navigate from the bottom left to the top right corner. The gridworld has stochastic dynamics: the chance of uniformly transferring to a state around the target is pslip = 0.2.
Setup In order to construct dataset settings, we generate dataset of 3 groups of expert trajectories by training the policies using a value iteration method on modified versions of the MDP. The initial and final states of the trajectories are fixed. We introduce a selection bias into the IRL feature expectation matching problem by manipulating the trajectory dataset in a way that we have a different number of trajectories for each of the three paths chosen by the experts (Fig. 2a): 400 trajectories for 1st group, 25 trajectories for 2nd group, 3 trajectories for 3rd group.
Results In Fig 2, we can observe that both the unregularized MaxEnt IRL algorithm (ERM) (Fig. 2b) and L2-regularized MaxEnt IRL algorithm (ERM) (Fig. 2c) exhibit overfitting to the expert datasets and partially to recover a meaningful reward and respective policy where as the IRMregularized version recovers a shaped reward function which takes the different optimal paths into account in a correct manner. In particular, increasing the regularization strength λ improves the reward significantly. (Fig. 2d - Fig. 2e).

heading: 5.2 ADVERSARIAL SETTING: PYBULLET ROBOT ENVIRONMENTS
text: Setup In this section, the experiments are performed in PyBullet (Ellenberger, 2018–2019) gym (Brockman et al., 2016) environments, which are an open-source alternative to the MuJoCo physics
simulator. We generate the demonstration datasets by training policies using the PPO algorithm (Schulman et al., 2017) 1 and varying the dynamics of the environments. An environment is defined as an instance of the variable E, which has an impact on the prior distribution of transitions (and by extension, trajectories) p(s, a, s′) and p(τ) in a minibatch of rolled out on-policy and expert transitions. We have used 10 expert trajectories for every environment for the physical parameter modification experiments. For the goal intervention experiments, we have used a proportion of 1, 3 and 10 trajectories for the three goal directions respectively.
Results on physical parameter interventions We test the performance using the policy rollout method with respect to the ground truth reward. We compare the policies trained using the adversarial approach with and without the invariance penalty. The evaluation environment setting is chosen to be outside of the value range presented to the model at training time. We can observe in Fig. 3 that the IRM-regularized version of AIRL outperforms the baseline in a zero-shot generalization setting in both of the locomotion environments (CustomAntMuJoCo and CustomHalfCheetahMuJoCo) 2), where interventions were performed on the hind leg length (2x and 3x of original length) and friction coefficients respectively. The friction coefficients have been modified to have a value of (1.5, 0.1, 0.1) and (2.0, 0.6, 0.6) tangential, torsional and rolling coeffients respectively. In constrast, we can see that for the CustomReacherPyBulletEnv environment, where the demonstrations were recorded by policies acting on the environment with a varying gravity coefficient, both models perform similarly. We attribute this to the fact that the gravity plays little role in the context of the Reacher environment.
Results on goal interventions We have also performed a second set of experiments using the following trajectory modification strategy. By altering the target y-coordinate of the locomotion environments and effectively changing the goal direction, we have generated a set of three distinct trajectory distributions. We have further introduced a selection bias into the dataset by selecting a higher number of trajectories which favors the right (ytarget = −900) in addition to the trajectories with the default goal position. The target y-coordinate in the out-of-distribution testing environment has been chosen to be the hard left (ytarget = 900). We can see an even stronger improvement in performance in IRM-regularized models (fig. 4) when compared to the ERM baseline, which fails to learn an effective policy in the case of the Half Cheetah.
1stable-baselines3 implementation 2*MuJoCo refers to the environment (MDP) specification. The simulation engine is PyBullet.

heading: 5.3 REWARD TRANSFER
text: In this experimental scenario, we use trained reward functions extracted from the discriminator model in order to retrain a new policy from scratch. The training dynamics of the policy with respect to the frozen version of the reward parameters is different as opposed to the training dynamics during adversarial training due to the varying discriminator quality during training. In order to train the models, we use the same reference PPO implementation as for the generation of the expert trajectories in the previous setting. In this scenario, we additionally use the LunarLander-v2 environment as a benchmark with interventions on the bounce parameter for the interaction of the landing gear and the platform as well as the engine power multiplier.
In table 1, we can observe that the policies trained on the rewards obtained through regularized adversarial training show improved performance compared to the ERM baseline in terms of cumulative reward but fail to reach expert performance which corresponds to a policy trained using the ground truth reward. We consider this to be of limited importance since we typically do not have access to the ground truth reward function in an IRL setting.

heading: 6 CONCLUSION
text: We have presented a regularization objective for inverse reinforcement learning to recover robust reward functions which avoid to learn spurious correlations present in demonstration data sets. The robustness manifests itself as improved generalization performance in out-of-distribution settings both in the maximum entropy IRL case based on feature expectation matching as well as the adversarial setting. The student policies that optimize the adversarial discriminator signal regularized by our method, demonstrate improved zero-shot generalization capabilities as shown by our experiments.
An open challenge remains the question how to construct effective interventional settings, i.e., how to determine whether the interventional assumptions are fulfilled by the given expert datasets. While the presented model successfully improves upon the baseline, it remains an open question what is the correct degree of invariance the model should aim for while avoiding degenerate invariant solutions.
In future work, one could envision a setting where such interventions could be discovered in an adversarial manner, for instance, by leveraging procedurally generated environments.
In the process of adversarial training, the distribution over transitions P (s, a, s′) is subject to shifts due to an improving student policy πθ. The invariant regularization objective imposed on the discriminator and the exploitation of this effect by the student policy coupled with the adversarial training dynamics requires further analysis. Explicit use of the regularization term as a reward component for the student policy could be another avenue to explore.

heading: A DERIVATION OF THE PENALTY TERM FOR MAXIMUM ENTROPY IRL
text: The gradient of the log-likelihood loss LMLE = ∑ τ∈De log p(τ |ψ,ϕ) w.r.t. to ψ is computed as follows:
LMLE = ∑ τ∈De log p(τ |ψ,ϕ) = ∑ τ∈De log(exp(ψTϕ(τ)))− logZψ,ϕ = ∑ τ∈De ψTϕ(τ)− logZψ,ϕ
Differentiating w.r.t. ψ yields:
∂LMLE ∂ψ = EDe [ϕ(τ)]− 1
Zψ,ϕ
∫ exp(ψTϕ(τ))ϕ(τ)dτ
= EDE [ϕ(τ)]− Ep(τ |ψ)[ϕ(τ)]
The gradient penalty term from Eq. 9 with respect to the features ϕ is derived as follows:
∇ϕ ∥∥∇ψ|ψ=1.0Le (r (ψ,ϕ))∥∥2 = ∂||∂Le(r(ψ,ϕ)·)∂ψ |ψ=1.0||2
∂ϕ
We employ the chain rule:
∂Le (r (ψ,ϕ)) ∂ψ = ∂Le (r (ψ,ϕ)) ∂r · ∂ (r (ψ,ϕ)) ∂ψ = ∂Le (r (ψ,ϕ)) ∂r · ϕ
where the last equality holds because we assume a linear reward with respect to the features ϕ: r (ψ,ϕ) = ψTϕ. Also, in section 4.2 we showed that :
∂Le (r (ψ,ϕ))
∂r = EDE [ϕ(τ)]− Ep(τ |ψ)[ϕ(τ))
where EDE [ϕ(τ)] are the feature statistics of the expert and E [ρπ] are the feature statistics of the student policy. For the sake of simplicity we define:
C := EDE [ϕ(τ)]− Ep(τ |ψ)[ϕ(τ)) which is independent of ϕ: Then:
∂Le (r (ψ,ϕ))
∂ψ = Cϕ
We obtain the IRM penalty for the feature matching maximum entropy IRL case as the following term:
∇ϕ ∥∥∇ψ|ψ=1.0Le (r (ψ,ϕ))∥∥2 = ∂||∂Le(r(ψ,ϕ)·)∂ψ |ψ|ψ=1.0||2
∂ϕ = ∂ ‖Cϕ‖2 ∂ϕ
= ∂ [Cϕ]
T [Cϕ]
∂ϕ
= CTC ∂ϕTϕ
∂ϕ = 2 ‖C‖2 ϕ = 2||ρE − E [ρψ] ||2ϕ

heading: B STRUCTURAL CAUSAL MODELS
text: A structural causal model (SCM) is defined as a tuple G = (S, P (ε)), where P (ε) = ∏ i≤K P (εi) is a product distribution over exogenous latent variables εi and S = {f1, ..., fK} is a set of structural mechanisms where pa(i) denotes the set of parent nodes of variable xi:
xi := fi(pa(xi), εi) for i ∈ |S| (10)
G induces a directed acyclic graph (DAG) over the variables nodes xi and entails a joint observational distribution PG = ∏ i≤K p(xi|pa(xi)) over variables xi conditioned on the parents of xi for some probability distribution p(·|pa(xi)) describing the mechanism fi. Interventions on G constitute modifications of one or more structural mechanisms fi yielding interventional distributions P̃G. In this paper, we have considered interventional distributions of expert trajectories P̃Ge(τ).

heading: C MODEL ARCHITECTURE AND TRAINING DETAILS
text: For the gridworld experiments, we use a DeepMaxEnt (Wulfmeier et al., 2015) formulation of the IRL problem. The state features are parametrized by a 2-layer MLP with a 1-dimensional hidden layer. We use RMSProp as an optimizer with a learning rate of 1e− 3. For the adversarial learning experiments, we use an actor and critic network with two layers of 64 neurons ( MlpPolicy from stable-baselines3 (Raffin et al., 2019)) We use Adam (Kingma & Ba, 2014) as the optimizer with an initial learning rate of η = 3e− 4 and an L2 weight decay coefficient of λ = 1e − 4. We perform one update of the discriminator network for every 3 updates of the actor-critic networks. We use the tuned hyperparameters from rl-baselines-zoo-3 3 to train the policy networks both using the ground truth reward as well as during adversarial training and using the recovered reward for the transfer experiments.
D INTERVENTIONAL SETTINGS

Review summaries:
The authors' propose a regularization loss for inverse reinforcement learning (IRL) based on invariant risk minimisation (IRM). They argue that existing IRL methods suffer from overfitting due to the sparsity of expert trajectories, and that this is analogous to the dataset level overfitting tackled by IRM. This approach is validated by comparing against unregularized and L2 regularized maximum entropy IRL on a grid-world, and is used in a scalable adversarial IRL algorithm when dealing with more complex environments. They close by evaluating the reward function learnt by this regularized adversarial IRL in terms of its utility in training a new policy from scratch.

This work draws inspiration from the invariant risk minimization (IRM) framework to improve the generalization capability of inverse reinforcement learning (IRL) methods. The paper's contributions are deriving a new IRL objective from IRM. This new algorithm includes a penalty regularization term to the standard objective which helps the agent from overfitting to the peculiarities of the expert demonstrations. 

This paper investigates a specific technique for Inverse reinforcement learning to avoid the detrimental effect of "spurious correlations in the data by the learning model". The objective is to avoid "behavioural overfitting to the expert data set". The paper uses an invariant risk minimization principle as a regularization approach for the maximum entropy inverse RL problem.

The paper studies the Inverse Reinforcement Learning (IRL) problem, addressing the issue of avoiding overfitting when having access to a finite set of expert demonstrations. The paper proposes an approach based on Invariant Risk Minimization (IRM) as a regularization approach for maximum entropy IRL. After having presented the method, an experimental evaluation on both finite and continuous environments, showing the advantages of the presented approach.