Learning Invariant Reward Functions through Trajectory Interventions
Authors: Ivan Ovinnikov
Eugene Bykovets
Joachim M. Buhmann
Venue: ICLR 2022 Submitted (2022)

Abstract: Inverse reinforcement learning methods aim to retrieve the reward function of
a Markov decision process based on a dataset of expert demonstrations. The
commonplace scarcity of such demonstrations potentially leads to the absorption of
spurious correlations in the data by the learning model, which as a result, exhibits
behavioural overfitting to the expert dataset when trained on the obtained reward
function. We study the generalization properties of the maximum entropy method
for solving the inverse reinforcement learning problem for both exact and approximate
formulations and demonstrate that by applying an instantiation of the invariant
risk minimization principle, we can recover reward functions which induce better
performing policies across domains in the transfer setting.

Under review as a conference paper at ICLR 2022
LEARNING INVARIANT REWARD THROUGH TRAJECTORY INTERVENTIONS
FUNCTIONS
Anonymous authors Paper under double-blind review
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 pro- vides 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, recover superior reward functions across domains in the transfer setting and, as a result, induce better performing policies than the empirical risk minimization baseline.
INTRODUCTION
Recent advancements in reinforcement learning methods applied real world scenarios have revealed a major limitation, namely how to design a problem specific reward function. This design challenge typically requires cumbersome and error-prone process of handerafting heuristic function to account for all the intricacies of the task at hand. Eliciting the correct behaviour via the optimization of reward function is of paramount importance both for training artificial agents, such robots, as well as human agents, for instance, in the case of teaching skills in 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 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- 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 the noise in the expert demonstrations. In particular, high- nigh-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 significant discrepancy between expert demonstrations originating from different experts. Adversarial imitation methods fail 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 solution strategy.
Invariant risk minimization (IRM) (Arjovsky et al., 2019) exploits the concept of invariant causal prediction (Ptters et al., 2015 Heinze-Deml et 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 aim to elicit behavioural policies

---

Under review as a conference paper at ICLR 2022
without exploiting spurious reward features, i.. features which cause nontransferable behaviours. To achieve this goal, we make the assumption that variations in expert demonstrations are product of causal interventions 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 regularization principle based 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 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 causal considerations in Section 3. We then present our model for the exact and approximate settings in Section and demonstrate the experimental results in Section 5. We conclude with brief discussion and outlook in Section 6.
2 RELATED WORK
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 an imitation setting in the work by (Ho & Ermon 2016). A disentangled version of the ad adversarial imitation approach which allows to recover rewards has been presented in (Fu et al., 2017). The authors of (Zolna 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 transfer setting.
Invariant representation learning Invariant risk minimization is method proposed in (Arjovsky et al. 2019) which builds upon the invariant causal prediction (Peters et al. [2055] Heinze- Deml ct al. 2017) principle and proposes a way to tackle the issue of spurious correlations impeding out-of-distribuue generalization of classification algorithms. Various other formulations of the IRM principle have subsequently been proposed. IRM games (Ahuja et al.) (2200) reformulates the IRM objective from a game theoretic perspective and achieves lower variance compared the original IRM method. Invariant rationalization proposes an adversarial method for recovering invariant rationales in the context of natural language processing (Chang 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 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. 1 uses the IRM games (Ahuja et al. 2020 formulation to learn policies invariant to certain domain variations. The authors of (de Haan et al. (019)) 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, algorithm is the first proposed method to use invariant causal prediction in the context of inverse reinforcement learning.
PROBLEM SETTING
MDP
We consider environments modelled by a Markov decision process M Au, R). where 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, describing the probability of transitioning state s' when taking
2

---

Under review as a conference paper at ICLR 2022
action in state Po is the initial state distribution, and R 8 A IR is the reward function. A policy ı is a conditional probability distribution of actions a E given states s with state descriptors or features (s)).
Inverse reinforcement learning aims to estimate a suitable reward function Te parameterized by
weights v based on dataset of expert trajectories DE = {7x brik where т; = (3 a is sequence of states and actions of expert of length T. To achieve this goal, the feature expectation statisties under the state occupancy measure of the expert, E,, [o(]))(()]] are matched with the (f(T) = Eeer (ss) denotes the sum of the state features p(s). The trajectory distribution p(T((b is induced by the policy Tre trained on the reward function estimate Tw. Given the ill- ill-posed nature the problem, maximum entropy formulation of the IRL. problem (Ziebart et al.) 2008) is typically chosen in order to find unique solution. The model describes the fact that the expert trajectories are sampled from the Gibbs distribution ( sp) x exp(r (T)) 7.. exp v² (p()) which corresponds to the solution of the entropy maximization problem under feature matching and simplex constraints.
statistics of the student state occupancy measure (Br=p(rly)]<((7)
max Hp((Tl)) s.t. {Er (p)()))))pp(()))))))))))))))))) (ELDD)lP(()) J p((Tw))
p(T())> (1)
"
where (E>-D/(((()) is the sample average of the expert trajectories. In the feature matching case, the reward ry (T) is typically assumed to be linear in the state features I' = vi (TT). We further consider the scenario where the state features (p(r) are learned using the DeepMaxEnt model (Wulfmeier ct al. 2013). 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 generative adversarial network (Goodfellow et al. (2144) architecture. Adversarial inverse RL (AIRL) methods yield 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 т ~ (T), which maximizes an expected cumulative reward objective based on the discriminator output. We adopt the discriminator structure from (Fu et al. (207)) which separates into the state-action dependent reward function gE and the state dependent shaping term hg.
exploc(s.a) + (s) hs exp(ge a) + yho(s') ho(s)) + (a(s))
(s,
We also consider state-only formulation of AIRL where the reward function (8) only depends on the state 8. The optimization objective of the discriminator the binary cross-entropy loss on the discriminator output:
LECRlE, 0,0) E(s,a)~< log( De ((, (B (s,a)>= log(1 (s, a))
(3)
The generator corresponds to the policy optimizing the expected cumulative reward given by the discriminator output and trained using the proximal policy optimization (PPO) (Schulman et 2017) method.
As with many neural network approximators, the discriminator model absorbs spurious correlations and this learning effect poses 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- Deme 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 &,r contains datasets D,
3

---

Under review as a conference paper at ICLR 2022
((re) We (2) ) sampled from interventional distributions P(X,, Y.) of the SCM of the underlying data generating process. The IRM method effectively learns an invariant data representation şp for the predictor w o yp (where is linear classifier) which enables stable conditional distribution P) by optimizing the objective stated in the definition below:
Definition 1 The data representation &p X + H mapping the input to hypothesis space, elicits an invariant predictor w o &p across the set of interventional settings if there is a classifier w - y simultaneously optimal for all settings:
min Pe H
( L. w E argmin (1 [())p Ve E Etr w
(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 the individual training set é, L The tractable approximation of Eq. uses the gradient norm penalty D(w 1.0, (. e) V we 10/ ( =2]l which quantifies the violation of the normal equations to measure the optimality of fixed classifier (w = 1.0) at each setting e. This leads to the following unconstrained formulation of the problem:
min Ec(()) + A gi Y
+
fnj fo (w up)l
(5)
In the following section, will show how to incorporate this optimization principle into the IRL setting.
4 MODEL
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 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.
4.1 INTERVENTIONS
The principle of invariant risk minimization relies on the assumption that the difference in training settings arises result of an intervention the data generating process described by structural causal model. In the case of maximum entropy IRI., the generative model the SCM of trajectories p(r (1:r) conditioned the optimality variable Or.r for timesteps (1,..., T} factorizes into the initial state distribution (S1). the conditional distribution p(O,|sr, az) = exp(r (Sr, ax)) of the binary optimality variable O, at timestep and the MDP transition dynamics p( [Sc, as) (Levine) 2018 is defined by the following expression:
p(T(O1=) O(T, Po(()) Il p(O, 1|s,. a, (1)p(s(+11ss...
(6)
Assumption Different expert demonstrations stem from interventional settings of the trajectory SCM.
The assumption is valid in scenario where expert demonstrations were gathered different dynamics. For instance, recordings of surgeons performing a procedure on different patients with variations of anatomy would satisfy this assumption.
4

---

In order apply the IRM principle to the inverse reinforcement learning problem, we must obtain set of settings & 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 Under review as a conference paper at ICLR 2022
[6]
16
64
Figure I: (a) Probabilistic graphical model of transition under influence of the index variable E and latent variable W. The stable conditional is highlighted in blue. (b) Spurious correlations assuming state-only formulation. (c) Spurious correlations assuming wrong edge orientation Or + 8141- (d) General setting where Og depends on causal x(o) and non-causal ((n)) features of the transition.
is the conditional P[O,|st at), which should be invariant under interventions to the transition SCM (Fig. la). We motivate this by the fact that despite the discrepancies in the demonstrations, all experts are assumed to perform the task in 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. Ia) we can observe the underlying assumption that the causal mechanism producing the optimality label O 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 non-causal information path corresponding to spurious correlations is formed in the transition SCM. In Fig. 1b can observe the scenario where we do not condition on the action. By conditioning on the collider node 8 and not observing the action node a, a path is formed between the setting index E and the optimality variable O,, resulting in the violation of their conditional independence relationship. A second scenario can be observed in Fig. Ic. This scenario requires the assumption that the orientation of the edge from node O to node 8t41 is temporally causal, meaning that the optimality of state at time is a causal parent of the next state. In this case, observing the collider node $141 implies the following conditional independence relationship: E JL O.|st-1. Beyond these scenarios, one can further assume more general partitioning of an arbitratry transition input (s, 8') into the causal transition feature components x()) and x ((n)).
reward function reprentation r(s)
(s, s') (x((), x(me)), illustrated in Fig l whereby conditioning on the ((n)) collider introduces a spurious correlation path (Huszar 2019).
The first scenario arises in the feature matching case where the trajectory feature representation ((r) learned using the DeepMaxEnt (Wulfmeier et al,, (015) model only has access to the subsequent state but not the action observations. The second scenario arises in the adversarial learning where the discriminator (s, a, s') depends on the node s', hence observing a collider.
Following the assumption of the IRM principle that only make interventions non-causal parents of the optimality variable Ot, possible interventions are state interventions caused by altering the initial state distribution the dynamics of the MDP and the action interventions, caused by intervening on the agent policy.
4.2 MAXIMUM ENTROPY IRL REGULARIZATION
For the application of the IRM regularization, first consider the maximum entropy feature matching scenario. For the Gibbs distribution p(T sF) = exp("" (()((((("()))))))))))))))))))()) over trajectories, we can write down the constrained optimization problem analogously to Eq.
min P.
Σ
Σ
log logp(r|1b, s.t. vi E argmin
Σ TED.
logg(r|/,5p)
REEI-TED
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:
5

---

Under review as a conference paper at ICLR 2022
min
Σ WED,
log
1 exp(15 (xxp(("<(1))) +D((V Zodo
e
In analogous fashion the invariant risk minimization approach, (e,, pp, e) is distance function representing the violation of the constraints w.r.. the optimal solution, which in our case corresponds to the parameters maximizing the likelihood of the Gibbs distribution. The gradient of the log- likelihood loss LMLE = EtED, log p(T (&, w.r.t. to &b is then the difference of the feature expectations [pg()(T)] - Eptr (irrtolee ((Tr) ZZebarr et al. 2008). The squared norm of this gradient constitutes the equivalent of the IRM penalty the maximum entropy IRL case.
(D e) = OoVELLO +)1l = || (p(()] - (
(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 neural network (p((s) according to the DeepMaxEnt model (Wulfmeier et al. 2015). The gradient norm penalty is then backpropagated to enforce invariant features sp
4.3 ADVERSARIAL IRL REGULARIZATION
We now present the adversarial version of the invariant regularization in algorithm I 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. distinct fashion as opposed to pooling the
The first is the fact that we use multiple experts demonstrations into one big dataset. The second is the regularization of the discriminator objective (Eq. 3) using the gradient norm penalty (E, é, wi e) = 1.0 corresponds to fixed scalar classifier. The regularized discriminator
9, w;
e) e] in a similar
fashion to Eq. 4 where effectively iterates over tuples (Di:, Dx)) Ve E Esr during adversarial training, where D is dataset of transitions generated by policy
Algorithm 1: IRM regularized adversarial IRI.
Input: Expert trajectories DE assumed to be obtained by intervening on P(T) in settings Result: Reward rg. trained student policy mp
Initialize policy mo and discriminator De.0i
for setting in {l,...,&r, } do for step t in [L.....N N) do
Collect trajectory buffer D (T.JEI by executing the policy me: Update De,0,0(8,0) u) via binary logistic regression by maximizing:
c(5,0,w;e) = CBCE(E, o, wi w)e) + +
using expert-specific tuple (Di, Dx) Update policy He w.r.t. Tv.e log De, (s, a, s') log(1 Ds.o) function of IRM regularized discriminator using policy gradient method (e.g. PPO);
end
end
The training process yields two outputs: the trained reward component r'g 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 measured using the ground truth reward. Furthermore, show that rewards obtained using this procedure elicit better performing policies when trained from random initialization on environments with modified dynamics.

---

Under review as a conference paper at ICLR 2022
5 EXPERIMENTS
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 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, 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 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
described in section 4.
5.1
TRACTABLE SETTING: GRIDWORLD EXPERIMENTS USING FEATURE MATCHING
To illustrate the principle of invariant risk minimization in the tractable IRL setting as discussed in we choose simple gridworld environment illustrated in (Fig. 2h). 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 state around the target is Palip 0.2.
4.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 modified versions of the MDP. The initial and final states of the trajectories are fixed. We introduce selection bias into the IRL feature expectation matching problem by manipulating the trajectory dataset 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, trajectories for 3rd group.
Results In Fig 2 we can observe that both the unregularized MaxEnt IRL algorithm (ERM) (Fig. 2D) and L2-regularized MaxEnt IRL algorithm (ERM) (Fig. ( exhibit overfitting to the expert datasets and partially to recover meaningful reward and respective policy where as the IRM- regularized 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. p Fig. 2
26
(4)
00
Figure 2: Gridworld ERM vs IRM reward recovery. (a) expert trajectory datasets: 1st group (blue)
400 trajectories, 2nd group (white): 25 trajectories. 3rd group (green): trajectories. (b) MaxEnt IRL ERM baseline (c) MaxEnt IRL ERM baseline with 1.2 regularization coefficent 1e-3 (d) MaxEnt IRL.
with IRM penalty, 0.01, (c) MaxEnt IRL with IRM penalty, À 0.05
5.2 ADVERSARIAL SETTING: PYBULLET ROBOT ENVIRONMENTS
Setup In this section, the experiments are performed in PyBullet (Ellenberger 2018 (18-2019) gym (Brockman al. 2016) environments, which open-source alternative to the MuJoCo physics
7

---

Under review as a conference paper at ICLR 2022
simulator. We generate the demonstration datasets by training policies using the PPO algorithm (Schulman et al. 2017) and varying the dynamics of the environments. An enviromment is defined as instance of the variable E, which has impact the prior distribution of transitions (and by extension, trajectories) p(s, a, 8') and p(T) 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.
Custonnaaccha
:000
Eskoddo award
1000
300
10000
400
0.7
0.2
01.
01.8 Timestapa
11.
1.0
1,2
1.4
0.00
0.15 6.50
1775 1:00
1.05
1.50
1.75
200
0.0
0.5
1:01
1.5
2.1
Timestape
Timestaps
Figure 3: Policy performance of the ERM vs IRM methods on out-of-distribution settings in the physical parameter intervention setting
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-regularize version of AIRL outperforms the bascline in a zero-shot generalization setting in both of the locomotion environments (Cust Mooooo and Cust omHHll a MMuJoCC) P. 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 Cust omRea che eecherPyyii rPyBul 1 Env environment, where the demonstrations were recorded by policies acting on the environment with 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 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 -iitr-i-iitiitti testing environment has been chosen to be the hard left (Ytarget = 900). We can see an even stronger improvement in performance 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.
CustonkknPPBdllleW
2000
Egtoods rexge
1750
500
6500
410
4350
3300
750
300
08
6.5
1.5 Timestopo
Timestopo
Figure 4: Policy performance the ERM vs IRM methods on out-of-distribution environments in a goal intervention setting
'stable-baselines3 implementation
2 Mulooo refers the environment (MDP) specification. The simulation engine PyBullet.
8

---

Under review as a conference paper at ICLR 2022
Table Transfer performance on out-of-distribution (OOD) environments. Standard deviation values are computed over different random seeds of the same OOD environment. The statistical significance p-values are 0.0458, 0.0456, 0.0041 and 0.0052 respectively for the four scenarios below (p S 0.05).
Model
Environment
AIRL-ERM
AIRL-IRM
Expert
LunarLander-v2 (bounce)
93.1
156.8453.1
230.3
LunarLander-v2 (engine power)
54.3
134.7+40.2
204.8422.3
CustomAmtMullooov0/ (friction)
542.7+114.5
783.9472.4
1223.2
CustomAnttulJooo-vO0 (leg length)
640.7+87.3
859.3
1394.64154.3
5.3 REWARD TRANSFER
In this experimental scenario, we use trained reward functions extracted from the discriminator model in order to retrain new policy from seratch. The training dynamics of the policy with respect 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 the previous setting. this scenario, we additionally use the LunarLa Lander--v 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 I we can observe that the policies trained on the rewards obtained through regularized adver- sarial training show improved performance compared to the ERM baseline in terms of cumulative reward but fail to reach expert performance which corresponds to 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.
CONCLUSION
We have presented regularization objective for inverse reinforcement learning recover robust reward functions which avoid to learn spurious correlations present in demonstration data sets. The robustness manifests itself improved generalization performance ouittf-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 the correct degree of invariance the model should aim for while avoiding degenerate invariant solutions.
In future work, one could envision 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 æa. The invariant regularization objective imposed on the discriminator and the exploitation this effect by the student policy coupled with the adversarial training dynamics requires further analysis. Explicit use of the regularization term as reward component for the student policy could be another avenue to explore.
REFERENCES
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. arXiv preprint arXiv:2002.04692, 2020.
Daniel Angelov, Yordan Hristov, and Subramanian Ramamoorthy, From demonstrations to task- space specifications. using causal analysis extract rule parameterization from demonstrations.

---

Under review as a conference paper at ICLR 2022
Autonomous Agents and Multt--gent Systems, 34(2):1-19, 2020
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization, 2019.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:/600..11540, 2016.
Shiyu Chang, Yang Zhang, Mo Yu, and Tommi Jaakkola. Invariant rationalization. In International Conference on Machine Learning, pp. 1448-148 1458. PMLR, 2020.
Pim de Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. arXiv preprint arXiv: 1905: 1979, 2019.
Benjamin Ellenberger. Pybullet gymperium.
Pybu gym 2019-2019.
Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. In International conference on machine learning, pp. 49-58. PMLR, 2016.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforce- ment learning, arXiv preprint arXiv: 1710. 1248, 2017.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. arXiv preprint 1406. 2014.
Christina Heinze-Deml, Jonas Peters, and Nicolai Meinshausen. Invariant causal prediction for nonlinear models, 2017.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. arXiv preprint 1r0iv:1606.033 03476, 2016.
Ferenc Huszar. Invariant risk minimization: An information theoretic view. https://www /www
zation/] 2019.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980) 2014.
Sergey Levine. Reinforcement learning and control probabilistic inference: Tutorial and review. arXiv preprini arXiv: 1808.000099 2018
Yunzhu Li, Jiaming Song, and Stefano Ermon. Infogail: Interpretable imitation learning from visual demonstrations. arXiv preprint 1rXiv:1703.008
2017.
Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations:
Theory and application reward shaping. In leml, volume 99, pp. 278-287, 1999.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference using invariant prediction: identification and confidence intervals. arXiv preprint arXiv:150..00333, 2015.
Antonin Raffin, Ashley Hill, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, and Noah Dormann. Stable baselines3.
2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. 2017.
Anoopkumar Sonar. Vincent Pacelli, and Anirudha Majumdar. Invariant policy optimization: Towards stronger generalization in reinforcement learning. In Learning for Dynamics and Control, pp. 21-33. PMLR, 2021.
Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfitting in reinforcement learning. arXiv preprint arXiv: /eti:////2.29999 2019.
10

---

Under review as a conference paper at ICLR 2022
Markus Wulfmeier, Peter Ondruska, and Ingmar Posner. Maximum entropy deep inverse reinforce- ment learning. arXiv preprint 355.::5500.44000 2015.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In International Conference on Machine Learning, pp. 1 11214-11224, PMLR, 2020.
Brian D Ziebart, Andrew L Maas, Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. Aaai, volume 8, pp. 1433-1438. Chicago, IL, USA, 2008.
Brian D Ziebart, Andrew Bagnell, and Anind K Dey. Modeling interaction via the principle of maximum causal entropy. In ICML,,010
Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarej, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation learning. arXiv preprint 11000:111. 2019.

---

Under review as a conference paper at ICLR 2022
A DERIVATION OF THE PENALTY TERM FOR MAXIMUM ENTROPY IRL
The gradient of the log-likelihood loss LMLE = ErED, logp(ri0) w.r.t. to v is computed as follows:
MLE Σ
log ())
Σ
log(exp(1)) (p())) log Z0. = Σ
o (((T) - log Vuup
TED.
TED
TED
Differentiating w.r.t. v yields:
θLMLE ED, (p(r)
/ Zo EDxle((T]] Extr((y(p(TT]
1 (xp(W"P(())))(( ((T)))(())
OGp
The gradient penalty term from Eq. 9 with respect to the features & is derived as follows:
a|l (r((voo)))))))
v=1.0|
v V Vii==10 (r (15, (1)
0e
We employ the chain rule:
aL (r (1,,4))
OT. (r (v),)) Ör
a (v),)))
OL (r (v,
Pub
Ör
where the last equality holds because assume linear reward with respect to the features yp: r (v,, s0) = v0t sp. Also, in section 4.2 we showed that
OL (r (vb,)) EDa [p(()] Ep(riw)((((T)
Ōr
where ED (x(r) the feature statisties of the expert and E [px] are the feature statistics of the student policy. For the sake of simplicity we define:
:= (BDa(p(T]) Eptr w)()(()) which is independent of yp: Then:
aL (r (1b, = Cy
We obtain the IRM penalty for the feature matching maximum entropy IRL case the following term:
0||
8L (r(s.e))
M 1
a [|ppp
V V Vol==I.LL (v, w]}l|²
dΦ
Oe
2 O[Cp)" [Cip]
Oip
CT c Opp
2|(C112 22|pE BBpp
Oep
12

---

Under review as a conference paper at ICLR 2022
B STRUCTURAL CAUSAL MODELS
structural causal model (SCM) is defined as a tuple 6 = (S, P(e)), where P(c) = IIes P(e;) is
product distribution over exogenous latent variables E, and S = ffi...., f..,...'''' is set of structural mechanisms where pa(i) denotes the parent nodes of variable z:;
3i e((a((z)),
for i [SS]
(10)
& induces directed acyclic graph (DAG) over the variables nodes and entails joint observational distribution Ps = I hta.://aaxxxx over variables x; conditioned on the parents of z; for some probability distribution p( (p)))))) describing the mechanism f;. Interventions on @ constitute modifications one or more structural mechanisms f. yielding interventional distributions Ps. In this paper, we have considered interventional distributions expert trajectories (T)
C MODEL ARCHITECTURE AND TRAINING DETAILS
For the gridworld experiments, we use Eax Wuulfmeier et al. 2013) formulation of the IRL problem. The state features parametrized by 2-layer MLP with I-dimensional hidden layer. We use RMSProp as an optimizer with a learning rate of le 3. 3.
For the adversarial learning experiments, neurons M1 pPol icy from 3able-baselines3 (Raffin et al., 2019)) We use Adam (Kingma & Ba 2014) as the optimizer with initial learning rate of 3e 4 and an L2 weight decay coefficient of = le 4. We perform one update of the discriminator network for every 3 updates of the actor-critic networks. We use the tuned hyperparameters from b-aaaaalleeeee leseiines--eeee zoo- 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.
actor and critic network with layers of 64
D INTERVENTIONAL SETTINGS
Table Interventional settings for the experimental environments (ID: in distribution, OOD: out-of- distribution)
Environment
Intervention
Values (ID)
Values (OOD)
LumarLander-v2
Bounce
0.1-0.5
0.9
Friction
0.2-0.4
0.001
CustomAmtMuloco-v00
Mass
0,5x, 2x
4x
Friction
(1.5,0.1,0.1), (2.00.,..)))
(2.5,0.7.0.7)
Hind Leg Length
1x,2x
3x
Direction angle
4(-10,0)
10
CustomHHallChccaahMuluoC-ooo
Mass
0.5x, 2x
4x
Friction
(1.5,0.1.0.1), (2.0.0.6.0.6)
(2.5,0.7.0.7)
Leg Length
0.7-1.3
1.5
Direction angle
U(-10,0)
10
CustomRReccheryyBllleeenn---
Gravity
0.5x,2x
4x
13


Reviews:
[{"review_id": "osE1e77UH0e", "reviewer": "Reviewer_nEw9", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* The work tackles the important problem of overfitting in inverse reinforcement learning. Drawing insights from invariant risk minimization the work proposes a well-motivated and theoretically supported objective. \n\n* Figure 2 is convincing in showing that larger regularization strength results in better rewards with the IRM compared to the ERM objective. \n\n* The paper was well written and had a good background work section. \n\nWeaknesses:\n* The method is greatly limited by the necessity of the interventions and knowing which demonstrations belong to which intervention sets. Inverse reinforcement learning (IRL) typically only requires a set of expert demonstration trajectories. This work requires significant additional information in the demonstration trajectories being partitioned as being generated from different \"state interventions\". This is not a realistic setting as the demonstrations need to be generated with special care for different state interventions and the learner has ground truth knowledge of these different settings. \n\n* Insufficient comparison with regularized baselines. Recent work [4], has demonstrated that spectral normalization and mixup normalization can greatly improve the generalization capability of imitation learning agents. I imagine the same can be true for inverse-RL. This comparison is especially important because the proposed method is a form of regularization. The authors should compare to these other forms of regularization. The L2 normalization baseline also is missing from all experiments except the grid world experiment.\n\n\n* Insufficient comparison with prior work in IRL. More recent methods than AIRL, such as f-IRL [3], have been proposed which perform better than AIRL. The authors should compare to the state of the art approaches in IRL. \n\n* Insufficient comparison with prior work in generalization in imitation learning [1,2]. While the authors mention that [2] do not study IRL, it seems like the techniques from [1,2] can easily be extended to IRL? \n\n* Insufficient experimental evidence. The main experiments only show results on a small handful of tasks. Furthermore, the Ant and HalfCheetah environments seem similar. I think experiments on more environments are necessary (such as a larger set of the Gym MuJoCo tasks). \n\n* I cannot find any written details about the method architecture or training details either in the main paper or supplementary material. \n\n* Presentation clarity lacks in some aspects. $ \\omega $ is never defined and its role in $ \\mathcal{L}_{\\text{BCE}}(\\xi, \\phi, \\omega ; e) $ is never clearly defined. What do the authors mean by a \"dummy classifier\"? \n\n\nQuestions / Clarifications: \n\n* In what realistic settings might we have access to this intervention information? \n\n* Why are methods trained for fewer steps in the goal intervention setting? \n\n* Why is the variance of IRM much higher? \n\nReferences:\n\n[1] Xu, Danfei, and Misha Denil. \"Positive-unlabeled reward learning.\" arXiv preprint arXiv:1911.00459 (2019).\n\n[2] Zolna, Konrad, et al. \"Task-relevant adversarial imitation learning.\" arXiv preprint arXiv:1910.01077 (2019).\n\n[3] Ni, Tianwei, et al. \"f-irl: Inverse reinforcement learning via state marginal matching.\" arXiv preprint arXiv:2011.04709 (2020).\n\n[4] Orsini, Manu, et al. \"What Matters for Adversarial Imitation Learning?.\" arXiv preprint arXiv:2106.00672 (2021).\n", "summary_of_the_paper": "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. ", "summary_of_the_review": "The paper has several major issues from the large assumption of interventions in the expert dataset, to insufficient comparisons to baselines, and insufficient experimental evidence.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "jaRvdt4UE7M", "reviewer": "Reviewer_33fy", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "As the paper is mainly empirical (no theoretical justification for many choices), the experimental parts should provide more insights or relevant experiments to be convincing Overall the claim that \"the regularization objective for inverse reinforcement learning recover robust reward functions which avoid to learn spurious correlations present in demonstration data sets\" is not clearly established from the paper. Beyond the scores over a few environments that can support the generalization claim, there isn't any clear elements in the paper to support the claim about spurious correlations. \n\nSome other elements from the paper are unclear. For instance, the following sentence is given without reference and does not seem obvious from the existing literature : \"As with many neural network approximators, the discriminator model absorbs spurious correlations and this learning effect poses a serious problem. These correlations coincide with the binary label information encoding the optimality of the expert.\" Can you clarify, e.g. by providing relevant literature?\n\nIs the source code provided?\n\nAdditional remarks:\n- Section 3: the state space and action space are not clearly defined. Are these discrete or continuous spaces?\n- In page 4, $\\mathcal L^e$ is not defined, except as \"individual losses\"?\n- $\\mathcal L_{BCE}$ in 4.3 takes different arguments than the one defined from Equation 3.\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper does not clearly backup all the claims and some elements are unclear.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "H0Ika8Van2", "reviewer": "Reviewer_YvDw", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- The paper addresses a very important issue of IRL, i.e., the risk of overfitting the demonstrated trajectories. I think that the approach presented in the paper is very interesting and represents a starting point for a comprehensive treatment of the topic.\n- The experimental evaluation, especially the gridworld experiment, effectively shows that the proposed approach is able to generalize very well, recovering a very smooth reward function. Moreover, the recovered rewards, in the LunarLander experiment, display better transferability properties compared to plain maximum entropy IRL.\n\nCons:\n- The construction of the interventions, as the authors acknowledge, might be complex and the paper does not provide any suggestion on how to discover them\n- Eq. (9) and Eq. (5) differ because in Eq. (9) there is a summation over the trajectories also in the constraint, while in Eq. (5) the constraint is stated for every intervention (universal quantification). Can the authors clarify?\n\nMinor issues\n- The ticks on the plot axis are very small\n- Pag. 4: The definition of \\mathcal{E}_{tr} is not clear. Is it a set of datasets, like \\mathcal{E}_{tr} = \\{D_{e_1},...,D_{e_m} \\}?", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, although the paper is not free of weaknesses, I think it addresses a very important problem of IRL making a first step towards the understanding of the generalization of reward functions recovered by IRL.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1Yk348pSTUU", "reviewer": "Reviewer_bhYo", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The intersection of IRL and IRM is quite novel as far as I know, and at a high level it seems like a promising approach. But this paper is let down by a glaring lack of details and some paper structuring issue, with more fundamental issues lurking beneath them.\n\nThe biggest oversight is the complete lack of an appendix. Many vital details (e.g. network architecture, environment reward structure, baseline methodology) are missing from the text, to the extent that it lowers my confidence in the rest of my critique -- I'm having to guess in order to evaluate potential experimental issues. This soon be remedied as soon as possible, and I will not vote for acceptance in its absence. Code release is appreciated, but is by no means a substitute for this; the text should be able to stand alone.\n\n1) Paper structure issues\nAs too much time is spent on irrelevant details of prior work, which leaves little room for critical aspects of the approach. Everything about structural causal models seems to serve no purpose. It certainly motivated IRM, but that can be left to the IRM paper itself; I see no place where this information clarifies or motivates this current method. Similarly, half a page is devoted to Figure 1, which is standard RL as inference fare, and is only used to the extent it motivates maximum entropy IRL (which is only used in the first experiment). In contrast, the jump from Equation 5 to 6 is quite vast; this is the part of the IRM paper that deserves repeating, since it adapting this regularizer that constitutes this paper's method. Similarly, equation 11 deserves unpacking. Converting the gradient of the log-likelihood loss into the difference of feature expectations is exactly the sort of prior work that needs to be understood, rather than merely cited. It is also unclear where these features come from (\"output of a neural network\" is far too vague), which ties into my larger complaint about the lack of relevant detail.\n\n2) Motivational issues\n2.1) Adversarial IRL/imitation learning\nThe combination of relying on an adversarial approach as your backbone IRL algorithm is fundamentally incompatible with your expressed desire to learning reward functions. There is a reason GAIL and related approaches are typically framed as 'imitation learning' rather than IRL: asymptotically their reward function should converge to being constant (i.e. when the policy is so good than one can't distinguish it's state occupancy from the experts'). This renders Table 1 meaningless, as better imitation should degrade the quality of the reward function. I'd suggest either switching to an non-adversarial IRL approach or dropping the idea of generating useful reward functions from the motivation.\n2.2) I agree the IRL approaches tend to overfit when expert data is limited, but this is separable from the issue of generalizing to novel MDPs (i.e. different state transition dynamics). This paper conflates these two issues throughout; seemingly motivated by the first, the experiments only address the second. And this might be hard to fix. Your theory states that all experts are equally optimal, but under maximum entropy RL/RL-as-inference this would make the experts identical if the state transition dynamics of the MDP are held fixed (this is kind of the point of max Ent IRL).\n\n3) Experimental Issues\n3.1) For the non-gridworld experiments it is unclear if the baseline method is just your method without the regularization or some completely different max Ent IRL approach. I'm hoping it is the former (if not this should be included), but the text the latter and I don't understand what would motivate this.\n3.2) Figure 3 results don't appear to be significant. It is also disconcerting that novel environments are being introduced with the literature is fully of similar environments with published results you could compare against. \n3.3) Figure 4 is very confusing. If the experts are optimal for completely different reward functions, what should be expect imitating their aggregate policy to yield? And what is the y-axis (reward function) here?\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "An interesting direction let down by poor paper structuring, a lack of detail, and a mismatch between motivation and algorithmic implementation. I highly suggest adding an appendix before addressing the other issues, as it is highly likely I've made some incorrect inferences about these missing details.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work addresses the issue of learning reward functions that overfit less/are invariant to irrelevant features of expert demonstrations. 
The proposed algorithm builds on top of adversarial imitation learning (AIRL) and proposes to include a regularization principle that is based on invariant risk minimization. The proposed algorithm is evaluated both in grid worlds as well as continuous control tasks. Both zero-shot policy transfer, as well as transfer of the reward function to learn out-of-distribution tasks from scratch.

**Strenghts**
This work is well motivated and addresses an important problem
The proposed method is well motivated, and provides theoretical foundations 

**Weaknesses**
The manuscript had many missing details/no appendix 
only one baseline is provided, while many relevant IRL algorithms exist
The evaluation is very limited in actually evaluating the invariance properties of the learned reward function 
poor alignment between how the proposed algorithm is motivated (learning invariant reward functions), and on what most of the experimental evaluation is focussed (zero-shot transfer of policy)**(more details on this below). 


**Rebuttal**
the authors have updated the manuscript to include an appendix and were able to address most structural issues and provided many of the missing details. 
No additional baselines were provided, and the experimental evaluation remains limited/poorly aligned with the initial motivation

**Summary**
This manuscript addresses an important problem and proposes a promising algorithm. My major remaining concern is the  experimental evaluation that seems not well aligned with the main contribution of this paper. As the authors state in their rebuttal the main supporting evidence for their claim is provided in Section 5.3, with only one set of experiments on using the reward function to learn policies on OOO tasks and very little analysis (< quarter of a page). While the majority of the evaluation (Section 5.2) is focussed on zero-shot transfer of the learned policy (which is trained during the IRL training phase). These zero-shot transfer experiments are not motivated in the context of the "learning invariant reward functions", so it's unclear what these results show. If these results are still relevant in showing that the proposed algorithm learns "invariant rewards", then this needs to be explained. Furthermore, more baselines would have been required (e.g algorithms that are focussed on learning a good policy by learning a "pseudo"-reward - such as GAIL). 
Because of this, my recommendation is that this manuscript is not quite ready yet for publication.
Decision: Reject