Invariant Causal Imitation Learning for Generalizable Policies
Authors: Ioana Bica
Daniel Jarrett
Mihaela van der Schaar
Venue: NeurIPS 2021 Poster (2021)
Keywords: imitation learning
representation learning
decision-making

Abstract: Consider learning an imitation policy on the basis of demonstrated behavior from multiple environments, with an eye towards deployment in an unseen environment. Since the observable features from each setting may be different, directly learning individual policies as mappings from features to actions is prone to spurious correlations---and may not generalize well. However, the expert’s policy is often a function of a shared latent structure underlying those observable features that is invariant across settings. By leveraging data from multiple environments, we propose Invariant Causal Imitation Learning (ICIL), a novel technique in which we learn a feature representation that is invariant across domains, on the basis of which we learn an imitation policy that matches expert behavior. To cope with transition dynamics mismatch, ICIL learns a shared representation of causal features (for all training environments), that is disentangled from the specific representations of noise variables (for each of those environments). Moreover, to ensure that the learned policy matches the observation distribution of the expert's policy, ICIL estimates the energy of the expert's observations and uses a regularization term that minimizes the imitator policy's next state energy. Experimentally, we compare our methods against several benchmarks in control and healthcare tasks and show its effectiveness in learning imitation policies capable of generalizing to unseen environments.

Invariant Causal Imitation Learning for Generalizable Policies
Ioana Bica*
Daniel Jarrett"
University of Oxford, Oxford, UK The Alan Turing Institute, London, UK ioana. bi bicaeng. ox ac. uk
University of Cambridge, Cambridge, UK daniel. jarret farretteie serreettmaths..mm cam ac uk
Mihacla van der Schaar University of Cambridge, Cambridge, UK University of California, Los Angeles, USA The Alan Turing Institute, London, UK mv4720cam ac uk
Abstract
Consider learning imitation policy the basis of demonstrated behavior from multiple environments, with an eye towards deployment in an unseen environment. Since the observable features from each setting may be different, directly learning individual policies as mappings from features to actions is prone to spurious correla- tions-ant and may not generalize well. However, the expert's policy is often function of a shared latent structure underlying those observable features that is invariant across settings. By leveraging data from multiple environments, we propose Invari- ant Causal Imitation Learning (ICIL), a novel technique in which we learn a feature representation that is invariant across domains, on the basis of which we learn an imitation policy that matches expert behavior. To cope with transition dynamics mismatch, ICIL learns a shared representation of causal features (for all training environments). that is independent from the specific representations of noise vari- ables (for each of those environments). Moreover, ensure that the learned policy matches the observation distribution of the expert's policy, ICII. estimates the energy of the expert's 's observations and uses regularization term that minimizes the imitator policy's 's next state energy. Experimentally, we compare our methods against several benchmarks control and healthcare tasks and show its effective- ness in learning imitation policies capable of generalizing to unseen environments.
Introduction
Strictly batch imitation learning aims to learn a policy that directly mimics the behaviour of experts, for which we only have access to set of demonstrations: logged trajectories of observations and actions following the expert's 'ss policy [1-3]. We cannot interact online with the environment, let alone query the expert any further, nor do we have reward signals for supervision. This setting is relevant in real-world scenarios where live experimentation is risky or costly-such healthcare and education.
Our aim is to learn an imitation policy in the strictly batch setting that faithfully matches the expert behaviour, while at the same time is able to generalize to unseen environments. healthcare, learning generalizable behaviour policy that could achieve expert performance new environments is an important goal: As means of providing clinical decision support, it could serve as an "individualized"
Equal contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

PL+1
CTV inommen invarianl
π
π
NC-
specific
at
A+41
Figure I: Causal diagram for the structure of environments. Expert demonstrations contain informa- tion about observations 24 and actions az We assume that observations decomposable into (1) state representations & that consist of the causal parents of the actions, and (2) noise representations v) that encapsulate any spurious correlations with the actions. To allow for dynamics mismatch, the transi- tions between the noise representations are specific to each environment. We want to recover the invari- state representation 8, such that the learned policy (-((s) generalizes well to new environments.
clinical guideline for actions that can be taken for different patient demographic from which we have no access to data during training. In this endeavor, a princi- pal challenge is that the sets of expert demonstrations that we have access to may contain variables that induce selection bias, or are otherwise spuriously correlated with the expert's is actions [4-7]. Directly learning an imitation policy from such data may lead to learning those spurious associations, thereby failing to generalize unseen environments, and perpetuating any biases in the expert's behaviour.
especially in hospital, region, or
However, in general it is likely that the expert's actions are only causally affected by a subset of the observed variables or by a shared latent structure [8, 9]. For instance, when imitating ideal driving behaviour, the background scenery might change, but the actions should only depend on car and road features. Another example includes the case when the lightning conditions in a room are changing, but physical dynamics of the environment are staying the same [7]. By leveraging expert trajectories from multiple different environments, our aim is to uncover this shared latent structure that causally determines expert actions, which allows us to eliminate the spurious associations and biases. this way. the learnt policy will better be able to generalize to any unseen environments that share the same latent structure as those used for training.
As illustrated in Figure 1, we assume access to observations and actions from the expert's policy in the different environments e. The observations are functions of noise factors 1f (which may differ across environments) and shared latent state representations (which is invariant across environments) that encapsulate the causal parents of the expert's actions. Note that the observed features for an environment may simply be the union of 17" and but they may also be any non- non-linear transformation of them. We shall operate the setting where there are no hidden confounders, i.e. that we observe all variables that are affecting the expert's 's actions (and the next states that result from these actions).
In addition to spurious correlations, another difficulty stems from learning imitate sequential behavior in the strictly batch setting itself: While behaviour cloning [10] provides an intrinsically batch solution, it ignores important information contained in the expert's roll-out distribution, and the learned policy may drift from the support of the distribution of states visited by the expert [1 12].
Contributions: In this paper, introduce Invariant Causal Imitation Learning (ICIL), novel method that learns a causal representation of the expert's 's actions actions-which is used to build a gen- eralizable imitation policy that matches the expert's behaviour. ICIL operates in the strictly batch setting and does not assume access to data from the target environments. By leveraging expert demonstrations from multiple different training environments, ICII. learns (shared) invariant causal representation as well as an (environment--pecifie noise representation. This accommodates dynamics mismatch across environments, while allowing the imitation policy to be learned by condi- tioning on the invariant causal representations. First, to satisfy the causal relationships in Figure 1, ICIL learns dynamics preserving representations and ensures that the learnt causal and noise repre- sentations are marginally independent by minimizing their mutual information. Second, to encourage the learnt imitation policy to stay within the support of the distribution of states visited by the expert's policy, ICIL estimates the energy of the expert's observations and uses a regularization term that minimizes the imitator policy's next state energy. Third, we evaluate ICIL against benchmarks for batch imitation learning in control and healthcare environments. We also empirically investigate directly using ideas from invariant risk minimization [6] to augment the loss function of existing batch imitation learning methods, and benchmark against their ability to generalize across environments.

---

Dynamics Sensory shill
State distribution matching
No access to farget trajectories
Temporal aspect
Method
Environment Offline mismatch (hidden confounders)
Pomerleau [10]
Model-free
Yos
No
No
No Model rollouts
N/A
No
Inbaiom
Leaming
Ho & Ermon [15]
Model-free No
No
No
NIA
Yes
Kostrikov et al. [2]
Model-free Yos
No
No
Adversarial off-policy matching
NIA
Yes
de Haan ef al. [8]
Model-free No
No
No
No
N/A
Yes
Gepeniiz In
et at. [16]
Model-free
No
Yes
No
Model rollouts
No
Yos
Kim et [17]
Model-based No
Yes
No
Model rollouts
No
Yes
Fisami et al. [18]
Model free
Yes
No
Yes
No
No
Yes
IRM Arjoesky al. [6]
Model
Yes
N/A
Yes No
ICIL (Ours)
Model-based
Yes
Yes
No
Energy-based
Yes
Yes
Table 1: Comparison of our proposed method with related works. ICIL operates in the strictly batch setting, allows for dynamics mismatch, does not require access to target trajectories, and incentivizes the imitation policy to stay in the support of the expert's 's distribution via energy- regularization.
2 Related Works
We tackle the problem of learning generalizable policies in an offline setting using ideas from causal inference. As such, our work straddles the intersection of research in (1) strictly batch imitation, (2) invariant representation learning, and-aore broadly- (3) causality in sequential decision-making.
Strietly Batch Imitation Learning: The simplest approach to imitation learning in the batch setting behaviour cloning [10] which uses standard supervised learning techniques to learn an imitation policy that minimizes the negative bog-likelihood of the observed demonstrator actions. However, behaviour cloning suffers from distributional shift as the learnt imitation policy cannot recover it reaches a state out-of the distribution of the expert demonstrations [1 1-14]. To overcome this problem, [1, 14] propose incorporate dynamics-awareed awareness by adding regularization to behaviour cloning by using norm-based penalties on the sparsity of implied rewards. Alternatively, [2] uses distribution matching approach and propose an offline objective for estimating the distribution ratio of the imitator policy and the expert policy, while [3] jointly learn a policy function together with an energy-basee model of the state distribution. However, none of the existing approaches consider the problem of generalization across environments and learning policies robust to spurious correlations.
Invariant Risk Minimization: In the supervised learning setting, Invariant Risk Minimization (IRM) [6] leverages data from multiple domains to learn data representation that elicits an invariant predictor across the different environments. The training data from each environment corresponds to different interventions on the data generating process. Given data from several training environments, the IRM objective aims to find representation such that there exists a classifier that is optimal across all training domains, i.e. that minimizes the empirical risk in each domain. This represents challenging, bi-level optimization, and [6] propose the IRM-v1 objective which a practical version to optimize. Through this optimization, the IRM objective should learn predictor that only uses the causal parents of the target variable and that is thus invariant across environments. However, directly using IRM for our sequential problem setting is not desirable, since it does not take into account the effect of each action on the subsequent states. Nonetheless, we empirically investigate augmenting existing methods for batch imitation to use the IRM-v1 objective in conjunction with their defined imitation risk, and verify whether they are able to generalize across environments. In our experiments observe that, in general, directly applying the IRM objective in this manner is not good enough.
Generalization in Imitation Learning: The problem of domain adaptation and transfer learning for the imitation learning setting has been tackled by several works so far. However, while they consider problems of dynamics., embodiment-, and/or viewpoint- expert, cxisting methods assume access to demonstrations from the target environment [17, 19, 19,20, assume access to online interaction or simulators in the different environments [16], or focus on the different problem of hidden confounding [18, 21J. Another line of work that is related is learning from demonstrations and meta-learning. While works in meta-learning also aims to generalize learnt policies to new-tasks, they require access to one more expert trajectories from the new task [22-27].
between the imitator and
Causality Imitation Learning and Reinforcement Learning: Several ideas from causality have been used to improve imitation learning and generalization in reinforcement learning. The idea of conditioning the imitation policy on the causal parents has been employed by [8] to avoid the problem of 'causal confusion' when learning a policy for the single environment setting. However, [8] requires
3

---

queering the expert or being able to perform interventions in the environment this is not possible in the batch setting, Similarly. [9.28] also learn causal relationships between the observations, actions and rewards by perforning/iimularing the effect of interventions in the environment. Alternatively, [29] use ideas from Invariant Risk Minimization 30] to learn optimal reinforcement learning policies that generalize across domains. Perhaps the most similar setting to ours is the one in [7] which studies the problem of generalization in reinforcement learning and also learn a representation that is shared across the domains. However, unlike our imitation setting, they assume access to a known reward sig- nal, and focus on learning the causal ancestors that reward to improve reinforcement learning [31].
To the best of our knowledge, we are the first to tackle the problem of learning generalizable imitation policies in the strictly batch setting. Table 1 summarizes main differences with relevant related works.
3 Problem Formalism
3.1 Imitation Learning
We work in the standard Markov decision process (MDP) setting: Let an environment be given by e = -X, T, T,r.?), with observations I E X, actions a E A, transition function T E A x reward function r E R and discount factor 7. Let ग E A(A)* be policy with the induced occupancy measure Pr (r, a) = (1 (1-7) Er=o 7'p(2 x, dg = a b -- (t-1), ag ( (mx)) of observations and actions, and let Px = - LEEA Ps (xr, a) be the observation occupancy measure.
Unlike in the reinforcement learning setting, where the aim is to learn policy T(- xr) that maximizes the cumulative sum of some known reward signal, in imitation learning the reward is neither known nor observed. Instead, we only have access to a dataset of trajectories D { from a demonstrator policy n. where each trajectory T ~ ιn = (Iz, úts Je+ 1 )o..... consists of sequence observation, action, next observation tuples that are sampled as at x xx)) and xz T(- 2e, az).
The goal of imitation learning is to seek an imitation policy π that minimizes the following risk:
R(T) L(T, mD)
where L is choice of loss function. Now, if we were in the online setting, we would have access to the environment (or simulator), with which we can interactively perform distribution matching by minimizing the divergence between the expert's 's's state occupancy PD and the imitator's 's state occupancy Pa [15,32-34]. One example is to use the (forward) KL divergence: L(T,, ((T,TD) = D (pDllps) [34]. However, in the offline setting we have no further access to the environment. As noted above, the simplest solution is behaviour cloning (BC) 11,,3,.36, which minimizes the negative log-likelihood of the demonstrator's actions. However, by disregarding the distribution of the expert's ''s observations, imitation policies learnt by BC often result in compounding error when deployed in practice [11-14].
3.2 Imitation Learning from Multiple Environments
Consider a family of environments . A,T",, ,n) lee
E} with observations x" X actions a E A. transition function T E ((2))xA reward function r E TxA, and discount factor y. This is the primary setting that we shall operate in. Note that the action space and discount factor do not change between environments. For notational simplicity, when considering the union over environments, we shall drop the index e, We assume offline access to dataset of recorded trajectories of training environments Etrain CE, D = (((79) NNe, I e Etrain), Each trajectory 7 ~ TD = (#f, az of4 )t=.... consists of a sequence of environment specific obser- vations, expert actions and next observations sampled a, TD (x)) and xf41 T (- art ax)-
from the expert policy aD in
In the presence of multiple environments, our goal is learn policy E A(A)* that matches the expert behaviour in all possible environments & that share a certain structure for the observations and the transition dynamics. In particular, this involves finding a policy that generalizes well across these related environments e E E-that is, the policy should ideally minimize the imitation risk across them:
max R (T) = L (T, mD) eEE
where each L explicitly depends on the characteristics of the environment Note that since we know nothing specific about ε, it is difficult to optimize for this directly. That said, we make mild assumptions about the "relatedness" of these environments, we can learn policies that generalize well.

---

Structure of Observations: First, we assume there is a shared latent structure underlying the obser- vations from different environments- on which the expert policy depends. Finding such structure would let us discard irrelevant factors inputs to the learnt policy, improving generalization [6, (,37,
Assumption 3.1. (Shared Latent Structure) Consider decomposing the observations xº E X in each environment e E & into two components: an invariant representation s S and noise terms 7f E z (i.e. spurious correlations), such that x" = q(s, n1") for some invertible transformation 9:8xx x 2 - X There exists some q such that TD only depends on s, and space S non-empty.
In other words, we assume that the demonstrator's policy depends only on information that shared across the environments, i.e. the state variables are the causal parents of the expert action ~ TD(- | s). As illustrated in Figure 1, the state variables and the noise terms are responsible for generating the patient observations, but the policy depends only on the state vari- ables. Thus, we allow different environments to have different p(r) marginals (as well dif- ferent p(a[z)). This allows environments to have different structure. The only requirement is that the environments some such that p(s) marginals should be the same (as well as p(a(s)). Learning such a repre- sentation that is invariant satisfies the standard Environment Invariance Constraint [38]. While this set-up is similar to the one in [7], a crucial difference is that we have no access to any re- ward functions whatsoever, and that must learn imitation policy in strictly batch setting.
the same far as the task is concerned. This means that there exists
Note that the latent structure induced by the state variables shared across the different environments. This means that the transition dynamics for the state representation p(s 8t,Ge) Qe) remain invariant across the environments. On the other hand, as different environments may be characterized by different types of spurious correlations, to allow for flexibility in their structure and evolution, we consider that the transition dynamics of the noise terms p" (7 78. at) are specific to each environment. Our goal, then, is to learn generalizable policy π tht is, one that depends only
Structure of Environments:
Second, learn policy that depends only on s, we must assume that the available training environments are actually different, so that we can learn the invariant state representation using the data from these environments and separate from the noise representation:
Assumption 3.2. (Environment Interventions) Each available training environment e corresponds to a hard [39] or soft [40] intervention on one or more dimensions of that environment' 's observation space (where these dimensions do not constitute any causal parents of the demonstrator's actions).
To ensure that generalizable policy actually exists, Assumption 3.1 requires that S be non-empty across all environments. Here, to ensure that the space S can actually be learned, Assumption 3.2 requires that 2 be non-empty across the training environments. Note that we require that the interventions inducing the different environments not be on the causal parents of the action, such that Assumption 3.1 is not violated.
Overall, our setting p(z) and p(a z) can differ between the multiple environments. How- ever, Assumption 3.1 enforces that the environments and tasks are the same modulo noise, iec that there exists some non-empty them. In other words, we have set of environments that are different (i.c. the dynam- ies of 2e are different), but the task being performed by the agent is the same (i.e. the dy- namics of 8; are the same). This setting applies to the case when lightning conditions in room are changing, but physical dynamics of the environment are staying the same [7] or when weather conditions are changing, but driving behaviour and dynamics are stay- ing the same. We provide additional ex- planations and definitions of environment interventions in Appendix A.
such that p(s) and p(a s) are the same between
Figure 2 shows a simple example where each observation zn represents a union of the causal parents of the action (state vari- ables) = { x xx) and the spurious cor- relations (noise variables) т). = { To satisfy Assumption 3.2, the different envi- ronments need to correspond to interven- tions on z². And to satisfy Assumption 3.1,
at
Figure Causal diagram illustrating temporal depen- dencies between causal parents of action {a], x and noise variables {2î] Different environments are in- duced by different interventions on the noise variables.
5

---

zt and 2 must not be intervened on. Our aim is to find representation of the causal parents =f] of the actions, as well as the mapping between them and the actions o,, that mimics the expert's policy.
Finally, similarly to [7], we also assume that the observations 2t timestep can only affect the actions de and the observations at the next timestep t + 1:
Assumption 3.3. (Temporal Causal Mechanism) Let I² and xd be any two components of the observation timestep t. Then:
2 I1 F241 Ie.
(3)
Note that Assumption 3.3 simply serves to place us within the standard MDP setting: It ensures Markovianity of the temporal transitions, that only the observations z; at time t will contain the causal parents of the action dz, and that and ar are the only factors that determine the next observation 2++1-
4 Invariant Causal Imitation Learning for Domain Generalization
The goal of Invariant Causal Imitation Learning (ICIL) algorithm is learn representation of the state variables s that is invariant across domains, and an imitation policy æ that depends on this causal representation and matches the demonstrator's behaviour. We operate in the strictly batch setting, and our aim is for ı to generalize to unseen environments e & given the above structural assumptions.
4.1 Learning Invariant Causal Representations
To achieve our goal, we decompose the observations =f in each environment into a representation 8, o(*f; 0.) for the causal features of the action a, and another representation of (* 0:) for the noise variables, where 6, and 0r are the learnable parameters of ó and Since the causal parents of the action are invariant across the environments, the state representation model Φ : X + S is the same across all environments. On the other hand, nM X + 2 is environment-speccin order to allow for dynamics mismatch of the noise variables between the different environments.
in
In order to satisfy the causal diagram in Figure and to learn a minimal causal representation, we need the following conditions be satisfied: (1) s, should be invariant across the environments, (2) 8, and of should be d'ynamic--reservingg and (3) 8t and rf should be independent from each other.
Firstly, to fulfill condition (1) we train an environment classifier on the shared state representation C2 : S + |Erainl parameterized by 0e using the cross-entropy loss. Similarly to [7]. in order to build state representation that is invariant across domains, we use adversarial loss [41] that maximizes the entropy of the classifier: H (c, (x(x;;; 0,); This gives us the following practical loss function:
Lino (0.) = Σ
Eng~ob H(c.(0((5;:0)))00)
eEEtrain
Out of all possible representations that are invariant, we specifically seek one that also preserves the transition dynamics, fulfilling condition (2). To ensure that the state and noise representations are we also learn the transition dynamics for the state variables 9s S x A : S, such that 8141 ga (8t 0r; 0g.) and the environment specific transition dynamics for the noise variables: sf : 2e x A 2a, such that iff4 st (18 a,; ). To reconstruct xz+1 we also learn s
S x 2 - X such that 1 (8 ph 0z). This yields the following practical loss function:
Ldgn (0z, yy f0 8% 1eEErrin.
0z)=
Σ Exigneeee re |
$1+1l1
(5)
cEEtrian
Note that while an alternative approach could consider directly building an invertible mapping from xf to (s, r, the motivation for decoding 8t+11 and n/f into 4 1 twofold. addition to learning dynnamics-pesserving representations, we will see in Section 4.2, this also allows to compute the energy of the next state obtained by following the imitation policy and enforcing this be similar to
the distribution of states visited by the expert's 's policy.

---

Representation leamnnn
Transition dynamics learning
97
17
2
Layn
9s
2t
7 pm
Lyus
9
t
π
?/
2
F5
Cenerax
Is
8
Freecn paromcter
Dine
Policy learning
Figure 3: Block diagram of our model. ICIL decomposes the observations x into an invariant causal representation 8; and an environment specific noise representation ry". To obtain invariant representation, we maximize entropy of an environment classifier that receives as input & (Cino). Moreover, the state and noise representations the prediction error of the next observation (Caym) and independent by minimizing their mutual information (Cm.)). We learn generalizable imitation policy that is conditioned the invariant causal representation (L+) and to ensure that the learnt policy matches the distribution of the expert's observations, minimize the imitator's policy next state energy (Cenergy).
learnt to be dynamics preserving by minimizing
Finally, to ensure that the state representation and the noise representation are marginally independent per condition (3), minimize the mutual information between them. We the Mutual Informa- tion Neural Estimation (MINE) framework [42]. which provides a way for estimating the mutual information using neural networks. In particular, MINE uses neural information measure I [U, V) to approximate the mutual information between random variables U and V. Let Te.,, be a statistics network parametrized by 8m. MINE estimates (U, V) by ascending the gradient of the following:
I(U, V) = sup Epon) [Te..] - log(E ar(y) [e sup [UU V:
(6)
O.. UV
Om
where Pov the joint measure of (U, V) and IPr = Jn dPvv, Py = Jus dPov are the marginal distributions. p(n) denotes the empirical distribution associated with n i.i.d samples. As noted in [42], the neural information measure [U, V) can approximate the mutual information with arbitrary accuracy. We therefore add the following practical loss function to our optimization objective, which seeks to minimize the mutual information between the state representation and noise representation:
(mi (0m)eE
Σ eEEcruin
/ 0.), u(x;; (f):;mm))
Note that the parameters om of the statistics network To... used for computing the mutual information are updated through gradient ascent on T TU, V:Om).
4.2 Matching Expert Behaviour in Strictly Batch Setting
On the basis of the causal representation s, we shall learn a generalizable policy æ (parameterized by 0,) in the strictly batch setting, such that it matches the demonstrator's 's behaviour. To begin, first condition π on the representation 8, and minimize the negative log- likellood of the expert's actions:
L C. = eEEtrain
Exjarop log (
((x;; 0.): 0x)
(8)
However, having only this objective corresponds to performing behaviour cloning, which has well- known limitations [1 [11]. To mitigate compounding error, we want some form of added regulariza- tion incentivize the imitation policy stay within the distribution of the expert's observations.
In the online setting, popular approach is to make sure that the rollout distribution of the imitating policy matches the rollout distribution of the expert's policy-for instance, by minimizing some form
7

---

of divergence between their induced occupancy measures. However, this requires interactive access to the real environment or simulator to perform rollouts of intermediate policies-whicl is not possible in our setting. Instead, we propose a method that takes advantage of the learnt transition dynamics. For any current observation It ~ PD. we shall encourage the next observation F++1 obtained by following the imitation policy 0 ~ x( xy) to remain within the occupancy measure of the expert.
Consider approximating the expert's occupancy measure using an Energy Based Model (EBM)
such that pr(x) = expl Fa(2)) where the function Eş(z) X + IR is the energy function and
2(0)
Z - Eg (z) dz is the partition function. We parameterize Eg by a neural network. It is not possible to train the EBM directly through maximum likelihood because ((0) involves integrating over the entire input domain of which is impractical. Instead, we contrastive divergence to pre-train the energy function Eg [43,44]. Contrastive divergence lowers the energy of the observations coming from the expert's occupancy distribution and increases the energy of the observations outside of the expert's 's occupancy distribution. Refer to Appendix for details on how we train the EBM.
To incentive the imitation policy to stay within the distribution of the expert's observations, we train it to minimize the energy of the next observation obtained by following given the current observation:
Lenergy (0x; 0z, Og.: (0 0% JoeE 0p)
E
Eg(
EE 2EErrain
=p (-1a))
(g (.), (mf.)))
This effectively assigns a high "reward" to the imitation policy for staying within high-density areas of the expert's occupancy measure, and a low "reward" for straying from it. This can be seen as an adaptation of online imitation methods [45, 46] where the expectation would be instead over I; Pa.
We illustrate in Figure 3 all of the components of the our ICIL model. Further details and the full algorithm for optimizing ICIL can be found in Appendix C.
Experiments
We perform experiments OpenAl gym tasks [47] and on an ICU dataset from the MIMIC III database [48]. In both cases, we generate data from multiple domains by augmenting the feature space with noise variables (spurious correlations).
Benchmarks We compare ICII.' against standard methods for strictly batch imitation learning: Be- haviour Cloning (BC) [10]; Reward-regularized Classification for Apprenticeship Learning (RCAL), which incorporates dynamics-awarenee through a sparsity regularization on the implied rewards [14]; ValueDICE (VDICE) [2], which uses off-policy objective to estimate distribution ratios needed for distribution matching; as well as Energy-basee learns the imitator policy with an energy model of its state distributions. These methods seek to find a policy that approximately matches the expert's behaviour from a single environment, and were not designed with generalization in mind. Hence augment these benchmark by using the IRMv1 objective [6] in conjunction with their originally defined imitation risk to obtain the additional benchmarks: BC-IRM, RCAL-IRM, VDICE-IRM, and EDM-IRM. More details about how used the invariance- -basad penalty from IRM [6] to augment these existing methods such that they may learn generalizable policies can be found in Appendix D. Implementation details about all benchmarks and the hyperparameter settings used can be found in Appendix F.3.
Distribution Matching (EDM) [3], which jointly
5.1 Evaluation on OpenAT Gym
We perform experiments the following control tasks from OpenAl gym [47]: Acrobot [49], Cartpole [50], LunarLander [47] and BeamRider [51]. For each task, we use pre-trained RL agents from RL Baselines Zoo [52] and Stable OpenAl Baselines [53] obtain expert policies. We then follow an approach similar to the one in [7] to obtain datasets with demonstrations from the expert in two different environments. In particular, for Acrobot [49], Cartpole [50] and LunarL [47] add spurious correlations to the state space of each control task and an environment identifier. The
'The code for ICIL can be found at https::pp //t:: /htbtb.ttuuu...gg cwtp//aaneeesccaaaaaaabbbo wwtaaaabmmloreaaaaaa forheal thlabpub and https::/// //gi //gggtttuu thub.. com/ homm//aaaaicc/////caoaaaaa Caaiiaaaaaaaaaaa Invar ant Causal Laausal-Imitatiinn Learning.

---

(a) Acrabot
(b) Cartpole
(c) LanarLLande
(d) BeamRider
Figure 4: Evaluation OpenAl gym environments. r-axis indicates the number of trajectories (in 11,5, 10, 15, 20)) with expert demonstrations from each training environment given as input to each benchmark and y-axis represents the average return of the learnt imitation policy on the test environments, scaled between (expert performance) and (random policy performance).
spurious correlations in each environment are different multiplicative factors of a subset of variables in the original state space. The invariant causal state is represented by the original variables in the state space of each control task. We train the benchmarks on demonstrations from two environments with 1 and 2 multiplicative factors for the spurious correlations and we test with multiplicative factors sampled from ((-1, 1). For BeamRider [51], similarly to [7], different camera angles are used for the training and testing environments, In particular, we use two training environments where the game frames are rotated by 10 degrees to the left and to the right respectively,
environments
while the test environment has rotation. The rotation is applied to the entire frame for all trajectories in each environment. However, note that, despite the rotation, the dynamics for the state variables and how they influence the action stay the same. Further details about the train and test environments can be found in Appendix F.3.
We vary the number of demonstrated trajectories from each environment that we give as input to each benchmark and we evaluate them on the average return obtained by deploying the learnt imitation policies the test environment. Figure 4 shows the mean results and standard errors obtained across 10 runs where for each run we train the benchmarks on different trajectories from the train environments and we evaluate a test environment with newly sampled multiplicative factors for computing the spurious correlations. We notice that our method consistently outperforms the benchmarks and is capable of generalizing better to the unseen target environments. Moreover, generally found that using the IRMv1 I objective [6] together with existing methods for strictly batch imitation learning did not improve performance and resulted in more unstable training. For additional results, see Appendix F where we perform ablation studies to investigate the impact of the different terms in the loss function used to train ICIL overall performance, compare performance on train vs. test environments and also evaluate robustness to increasing the size of the spurious correlations.
5.2 Evaluation on MIMIC-III
We also perform experiments on a healthcare dataset with Intensive Care Unit (ICU) patients extracted from the Medical Information Mart for Intensive Care (MIMIC-III) database [48]. The dataset consists of trajectories of clinical measurements (e.g. heart rate, respiratory rate) recorded every hour. The aim is to learn generalizable policy for the action of putting patients the mechanical ventilator.
We define three cnvironments, two for training and one for testing, each consisting of 2000 indepen- dent patient trajectories from MIMIC-III. We augment the original feature space by adding spurious correlations (noise variables) that are the same as the expert actions with probabilities p = 0.1 and p 0.2 in the training environments and with probability p = 0.8 in the testing environment.
In real setting, such spurious correlations are commonplace. For instance, consider some hospitals (i.e. training environments) where selection bias is present, such that patients with certain otherwise irrelevant comorbidity happen to receive a treatment more often [54-57]. However, learning an imitation policy that takes into account such a comorbidity when assigning the patient's treatment would fail to generalize to hospitals where fewer patients suffer from this comorbidity but should still receive the treatments. More details about the dataset can be found in Appendix F.3.

---

Mechanical ventilator AUC
Benchmark
ACC
APR
BC
0.783 + 0.001
0.762 + 0.002
0.692 = 0.001
RCAL
0.790 0.002
0.771 0.002
0.697 = 0.002
VDICE
0.794 + 0.001
0.784 + 0.001
0.716 = 0.001
EDM
0.786 0.003
0.741 0.011
0.682 = 0.005
BC-IRM
0.791 + 0.002
0.767 + 0.003
0.696 + 0.002
RCAL-IRM
0.789 0.002
0.766 0.003
0.694 0.003
VDICE-IRM
0.766 + 0.001
0.730 + 0.001
0.694 = 0.001
EDM-IRM
0.781 0.004
0.717 0.015
0.673 0.007
ICIL
0.855 0.003
0.856 0.004
0.789 0.004
Table 2: Evaluation on MIMIC-III in terms of action-matching. We compare the actions selected by the benchmark imitation policies with the ones from the clinical expert policy the test environment and report the accuracy (ACC), the area under the receiving operator characteristic curve (AUC) and area under the precision-recal curve (APR).
Since MIMIC-III is an entirely offline dataset, it is not possible to compute average returns for rumning the policies in the test environment. Instead, we evaluate the benchmarks in terms action matching on the test environment. We report in table 2 the mean accuracy (ACC), the mean area under the receiving operator characteristic curve (AUC), the mean area under the the precision-rees recall curve (APR) and their standard deviations over 10 runs. We notice that ICIL learns policy that best discards the spurious correlations present in the training environment to learn a generalizable policy for putting patients on the mechanical ventilator that best matches the expert's 's actions on the test environment.
Discussion
In this paper, we tackle the problem of learning generalizable imitation policies in the strictly batch setting. Our ICIL model leverages ideas from causality and learns an invariant state representation that minimizes the presence of spurious correlations. By conditioning the imitation policy on this state representation, we obtain a policy that generalizes to environments with the same shared latent structure, but with different noise distribution and dynamics. ICIL also matches expert behaviour by incentivizing the learnt imitation policy to stay within the expert's observations distribution.
In terms of limitations, believe that future work should consider providing theoretical insights and error bounds on the generalization error. In addition, to be able to learn an invariant state representation, our method requires demonstrated trajectories from least training environments with different interventions on the noise variables (spurious correlations), and the method cannot be used if such data not available in practice. Finally, we bear in mind that-as with any other imitation learning method that aims to match the expert's policy- ICIL can have potential negative socictal impacts if the expert's policy is flawed in the first place. Thus, in sensitive applications such as clinical decision support, care must be taken to prevent potentially negative feedback loops.
Acknowledgments
We would like to thank the reviewers for their valuable feedhack. The research presented in this paper supported by The Alan Turing Institute, under the EPSRC grant EP/N510129/1 by Alzheimer's Research UK (ARUK), by the US Office of Naval Research (ONR), and by the National Science Foundation (NSF) under grant number 1722516.
References
[1] Bilal Piot, Matthieu Geist, and Olivier Pietquin. Bridging the gap between imitation learning and inverse reinforcement learning. IEEE transactions on neural networks and learning systems, 1814):1814-1822 2016.
10

---

[2] Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off- aff-ppoiiy distribu- tion matching, arXiv preprint arXiv:1912.05033 2019.
[3] Daniel Jarrett, Ioana Bica, and Mihaela van der Schaar. Strietly batch imitation learning by energy-based distribution matching. Advances in Neural Information Processing Systems, 2020.
[4] Amy Zhang, Yuxin Wu, and Joelle Pineau. Natural environment benchmarks for reinforcement learning. arXiv preprint arXiv:/81 1.06033,, 2018.
[5] Antonio Torralba and Alexei A Efros. Unbiased look at dataset bias. In CVPR 2011, pages 1521-1528. IEEE, 2011.
[6] Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint urXiv:/907.02500 2019.
[7] 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 Muchine Learning, pages 1214-111244 PMLR, 2020.
[8] Pim de Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. arXiv preprint arXiv: /900 / 1979, 2019.
[9] Timothy E Lee, Jialiang Zhao, Amrita S Sawhney, Siddharth Girdhar, and Oliver Kroemer. Causal reasoning in simulation for structure and transfer learning of robot manipulation policies. arXiv preprint w11.:::2103.17777 2021.
[10] Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural computation, 3(1):88-9) 97, 1991.
[11] Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pages 661 JMLR Workshop and Conference Proceedings, 2010.
[12] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth interna- tional conference on artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
[13] Francisco S Melo and Manuel Lopes. Learning from demonstration using mdp induced metrics. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 385-401, Springer, 2010.
[14] Bilal Piot, Matthieu Geist, and Olivier Pietquin. Boosted and reward-regularizedd classification for apprenticeship learning. In Proceedings of the 2014 international conference Autonomous agents and multi-agent systems, pages 1249-1256, 2014.
[15] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. arXiv preprint Ttff:16006034770 2016.
[16] Yiren Lu and Jonathan Tompson. Adail: Adaptive adversarial imitation learning. arXiv preprint 2000::2008.1266 2020.
[17] Kuno Kim, Yihong Gu, Jiaming Song, Shengjia Zhao, and Stefano Ermon. Domain adaptive imitation learning. In International Conference on Machine Learning, pages 5286-5295. PMIR, 2020.
[18] Jalal Etesami and Philipp Geiger. Causal transfer for imitation learning and decision making under sensor-shift In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 10125,, 2020.
[19] Pierre Sermanet, Corey Lynch, Jasmine Hsu, and Sergey Levine. Time-contrastive networks: Self-supervised learning from multi-view observation. In 2017 IEEE Conference on Computer Vision and Patlern Recognition Workshops (CVPRW), pages 486-487. IEEE, 2017.
[20] YuXuan Liu, Abhishek Gupta, Pieter Abbeel. and Sergey Levine. Imitation from observation: Learning to imitate behaviors from raw video via context translation. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pages 1118-111 125. IEEE. 2018
[21] Junzhe Zhang, Daniel Kumor, and Elias Bareinboim. Causal imitation learning with unobserved confounders. Advances in Neural Information Processing Systems, 33, 2020.

---

[22] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnoste meta-learning for fast adapta- tion of deep networks. In International Conference on Machine Learning. pages 1126-1155 1135. PMLR, 2017.
[23] Yan Duan, Marcin Andrychowicz, Bradly C Stadie, Jonathan Ho, Jonas Schneider, llya Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. arXiv preprint arXiv:1703.07326, 2017.
[24] Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot visual imitation learning via meta-leaming. In Conference on Robot Learning, pages 357-368. PMLR, 2017.
[25] Tianhe Yu, Chelsea Finn, Annie Xie, Sudeep Dasari, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot imitation from observing humans via domain- -adatiive meta-learning, arXiv preprint arXiv:1802.01553 2018
[26] Stephen James, Michael Bloesch, and Andrew J Davison. Task-embedded control networks for few-shot imitation learning. In Conference on Robot Learning, pages 783-795. PMLR, 2018.
[27] Pratyusha Sharma, Deepak Pathak, and Abhinav Gupta. Third-person visual imitation learning via decoupled hierarchical controller. arXiv preprint
2019.
[28] Sergei Volodin, Nevan Wichers, and Jeremy Nixon. Resolving spurious correlations in causal models of environments via interventions. arXiv preprint arXiv:2002.05221,, 2020.
[29] Anoopkumar Sonar, Vincent Pacelli, and Anirudha Majumdar. Invariant policy optimization: Towards stronger generalization in reinforcement learning. arXiv preprint arXiv:2006.01096 2020.
[30] Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. In International Conference on Machine Learning, pages 145-155. PMLR, 2020.
[31] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, pages 1861-1870. PMLR, 2018.
[32] Nir Baram, Oron Anschel, and Shie Mannor. Model-based adversarial imitation learning. arXiv preprint 1612.02 179, 2016.
[33] Chelsea Finn, Paul Christiano, Pieter Abbeel, and Sergey Levine. A connection between generative adversarial networks, inverse reinforcement learning, and energy-based models. arXiv preprini arXiv: 1.031.0330 1.03852, 2016.
[34] Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. divergence mini- mization perspective on imitation learning methods. In Conference on Robot Learning, pages 1259-1277 PMLR, 2020.
[35] Michael Bain and Claude Sammut. A framework for behavioural cloning. In Machine Intelli- gence 15. pages 103-129, 1995.
[36] Umar Syed and Robert E Schapire A reduction from apprenticeship learning to classification. Advances in neural information processing systems, 23:2253-2261, 2010.
[37] Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society, Series B (Statistical Methodology), pages 947-1012, 2016.
[38] Elliot Creager, Jörn- Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference on Machine Learning, pages 2189-2200 PMLR, 2021.
[39] Judea Pearl. Causality. Cambridge university press, 2009.
[40] Frederick Eberhardt and Richard Scheines. Interventions and causal inference. Philosophy of science, (4(5):981-9995 2007.
[41] Eric Tzeng, Judy Hoffman, Kate Saenko, and Trevor Darrell. Adversarial discriminative domain adaptation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7167-7176. 2017.
[42] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International Conference Machine Learning, pages 531-540. PMLR, 2018.
12

---

[43] Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800, 2002.
[44] Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. 2019.
[45] Siddharth Reddy, Anca D Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards. arXiv preprint arXiv:1905.11008, 2019.
[46] Minghuan Liu, Tairan He, Minkai Xu, and Weinan Zhang. Energy-based imitation learning. arXiv preprint urXiv:2004.080 0.999, 2020.
[47] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprini arXiv: 1606.0154, 2016.
[48] Alistair EW Johnson, Tom Pollard, Lu Shen, H Lehman Li-Wei, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark. Mimic-iii, a freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
[49] Alborz Geramifard, Christoph Dann, Robert Klein, William Dabney, and Jonathan How. Rlpy: value- sunbe-ffuncton-basee reinforcement learning framework for education and research. J. Mach. Learn, Res., 1573-11777,577 2015.
[50] Andrew G Barto, Richard S Sutton, and Charles W Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE transactions on systems, man, and cybernetics, (5):834-86,, 1983.
[51] Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279 2013.
[52] Antonin Raffn.. RI baselines zoo. https://github htwps/hwwccccomoaaaaaaaaaaaaaaaa 2018.
[53] Ashley Hill, Antonin Raffin, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, Rene Traore, Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, and Yuhuai Wu. Stable baselines. https: / //gituu.. httphhillaasstaleeeeeee wawssttaaebbbbbleeaeeeaeaaabbeee basebaseeen,s.. 2018.
[54] Chloe E Taylor, Helen Jones, Mohammad Zaregarizi, Nigel T Cable, Keith P George, and Greg Atkinson. Blood pressure status and post-exercise hypotension: an example of a spurious correlation in hypertension research? Journal of human hypertension, 24(9):585-592, 2010.
[55] Thomas Desautels, Ritankar Das, Jacob Calvert, Monica Trivedi, Charlotte Summers, David Wales, and Ari Ercole. Prediction of early unplanned intensive unit readmission in uk tertiary care hospital: a cross-sectional machine learning approach. BMJ open, 7()), 2017.
[56] Andrea Soo, Danny Zuege, Gordon H Fick, Daniel Niven, Luc R Berthiaume, Henry Stelfox, and Christopher Doig. Describing organ dysfunction in the intensive care unit: cohort study of 20,000 patients. Critical Care, 23(1):1-15, 2019.
[57] Adarsh Subbaswamy, Roy Adams, and Suchi Saria. Evaluating model robustness and stability to dataset shift. In International Conference on Artificial Intelligence and Statistics, pages 1611-2619 PMLR, 2021.
[58] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-II), pages 681-688. Citeseer, 2011.
[59] Daniel Jarrett and Mihaela van der Schaar. Inverse active sensing: Modeling and understanding timely decision decision-making. International Conference on Machine Learning (ICML), 2020.
13


Reviews:
[{"review_id": "xz9b_o1f8g", "reviewer": "Reviewer_VV1t", "summary": "The authors consider the setting of batch imitation learning with access to expert demonstrations over multiple environments which correspond to interventions on shared, underlying causal structure. They use the framework from [7] to propose a method that disentangles the relevant and irrelevant components of the state for a task with an additional regularization component that minimizes the mutual information between the two components with MINE. For the imitation policy, they propose an energy-based model approach to approximate the occupancy measure of the expert policy as a way to regularize and train the policy. \n", "questions": "", "limitations": "", "rating": 7, "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": "At first glance imitation learning in the batch setting seems like a strange setup, and therefore not so unusual that it is original. However, the authors motivate this setting well by pointing out that in the medical domain, these are reasonable assumptions to make — that you are in an offline setting with access to expert data from different environments with shared latent structure.\n\nAs the authors themselves point out, the setup, assumption, and parts of the method are directly from [7] but there are novel components to modify it for the batch and imitation learning setting. The experiments seem reasonable, although the unusual metric on the Open AI Gym tasks make it difficult to ground these results to those reported in other works.  Perhaps a more relevant benchmark would be D4RL, which has an expert dataset. \n\nThe MIMIC III experiment, given the setting, is well chosen. The results there also are clearly quite significant and with respect to reasonable baselines. It would also perhaps be best to report both train and evaluation performance given the emphasis on the focus being generalization.\n\nOverall this seems like a solid contribution and I would lean towards accept.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "No", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "t_beUBAxAfN", "reviewer": "Reviewer_raJe", "summary": "This paper proposes a method for performing imitation learning from a set of expert demonstrations collected across a set of environment variations, where each variation is an intervention on the noise variables of the environment. The method, ICIL, is demonstrated to successfully control for spurious correlations in a set of OpenAI Gym envronments and a healthcare domain dataset.", "questions": "", "limitations": "", "rating": 7, "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": "The method itself is well motivated with a clear formalization of the problem setting. The work is well positioned with respect to related works, and the specific problem of learning a generalizable policy from offline data is an important one. \n\nMy main concern relates to how the experiments focus on highly contrived toy cases specifically designed to fit the formal assumptions of their method. It would be beneficial to include experiments that test the method in settings where the examples were not hand-designed by the authors, e.g. OpenAI Procgen, where many levels show differences that can be viewed as spurious noise, e.g. background colors; or a version of Mujoco where the background changes across episodes. These more challenging examples would test how well this method can scale in practice, and how robust this method is to domains that do not perfectly match the structural assumptions around per-environment interventions only occuring on noise variables.\n\nIf the authors provide results on a less contrived domain, I will gladly raise my score to at least 7.\n\n### Post-rebuttal update:\nThe authors' rebuttal addresses my concerns, so I am raising my score to a 7.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "- The method makes the assumption that the environments from which demonstrations are collected correspond to interventions only on the noise variables of the environment. It would be valuable for the authors to describe real world scenarios meeting these assumptions, or if such examples are hard to come by, describe how these assumptions are relaxed in example practical domains.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "hf2dtZ_gJzX", "reviewer": "Reviewer_GSwe", "summary": "This paper proposes a new imitation learning algorithm, ICIL, which aims to make imitation learning more robust to irrelevant distractor features (i.e. features that are not causally upstream of the chosen expert action across a range of environments). Starting from BC, ICIL first adds a state encoder that splits states into an action-relevant component $s_t$ and a noise/distractor component $\\eta_t$. Next it adds three additional auxiliary losses to prevent distractors from influencing $s_t$: a \"dynamics-preserving\" loss to ensure that $(s_t, \\eta_t)$ are sufficient to predict the next environment observation; an \"environment invariance\" loss to make $s_t$ marginally independent of the environment; and a mutual information loss that regularises $s_t$ and $\\eta_t$ towards marginal independence. Finally, ICIL incorporates an auxiliary loss to match next-step expert state occupancy under a learnt dynamics model, with the intent of decreasing compounding error. Experiments show that ICIL gets higher return than several baselines in variants of standard benchmarks that have been adapted to include additional distractor features.", "questions": "", "limitations": "", "rating": 7, "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": "In my view, the main strengths of this paper are:\n\n1. Clarity. Writing quality was high on the whole: the problem is motivated in a compelling way, assumptions are clearly stated, and description of the algorithm and experiments makes sense. My only objection is to the assumptions and desiderata that are introduced without enough justification in sections 3-4 (see weaknesses below).\n2. The experiments have a wide selection of baselines. Incorporating IRM into the baselines was particularly appreciated—it's hard to find appropriate baselines for a problem like this, so it's good to see some novel baselines that at least try to solve the same problem in different ways. Moreover, the proposed algorithm manages to improve substantially over the baselines in all the examined environments.\n\nHowever, there are also some weaknesses to the algorithm:\n\n1. The algorithm makes several non-trivial assumptions and has many moving parts, and some of them seemed arbitrary or poorly motivated. For instance, it's assumed that the transitions of decision-relevant variables are the same between \"environments\", but transitions of noise variables are not. What kinds of domains is this true for in practice? The desiderata in Section 4.1 (invariance across environments, dynamics-preserving $s$/$\\eta$, disentangled $s$/$\\eta$) are also vague and get introduced without justification:\n   1. Why is having equal $s_t$ marginals across environments an adequate notion of \"invariance\"? What about environments where the demonstrator really does have different $s_t$ marginals because of some difference in environment structure? It seems like $\\mathcal L_{\\text{inv}}$ objective is going to throw away some relevant information in such cases.\n   2. Having \"dynamics-preserving\" state representations seems too weak—why not just have an invertible mapping from $x$ to $(s,\\eta)$? Such a mapping would still be \"dynamics-preserving\" while also satisfying Assumption 3.1 directly.\n   3. What does it mean for $s$ and $\\eta$ to be \"disentangled\"? Why is this desirable? \"Disentangled\" is a loaded word, and statistical independence (which an MI penalty incentivises) is neither necessary nor sufficient to obtain the kind of disentanglement that you see in carefully hand-designed state representations (for a counterexample, see Theorem 1 of \"Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations\" by Locatello et al.).\n2. I also had several concerns about experiments:\n   1. The various losses (there are five of them!) do not seem to have been examined separately in an ablation, so it's not clear which are responsible for the improvement over baselines.\n   2. The experiment domains all have manually-inserted distractors that are deliberately constructed to meet the assumptions of the algorithm. The experiments thus do not provide evidence that the assumptions are adequate to handle the sort of distractors present in real-world tasks.\n   3. The experiments report return/accuracy, but do not directly evaluate how sensitive the various models are to the inserted distractor features. In some ways this would be more informative than just looking at overall performance, since it would distinguish between gains from reduced reliance on distractors and gains from improved optimisation, etc.\n\nOn the whole I feel that this is a well-written paper, but I am currently ambivalent about acceptance because I do not feel the justification and evaluation of the algorithm is adequate given how complex the algorithm is. Clarifying points raised in the weaknesses above would help.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations and societal impact were discussed adequately.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.5"}, "extra_text": {}}, {"review_id": "cJ9M0EKFVB", "reviewer": "Reviewer_YxyH", "summary": "The paper presents an imitation learning approach that can generalize across demonstrations acquired from multiple environments. The approach projects the observations in an invariant shared latent space across environments. The variability across environments is modeled as noisy perturbations to reflect spurious correlations with the actions. The noisy perturbations and invariant state representations are learned such that the transition dynamics are preserved, while the mutual information between them is minimized. Experiments on OpenAI gym benchmarks and MIMIC III derived dataset suggest performance improvement over other competing methods.", "questions": "", "limitations": "", "rating": 6, "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 paper is well-written for the most part. The problem formulation is interesting and the rationale for the choice of loss functions is well-motivated. Fig. 3 is helpful to explain the ideas of the paper. Experimental results are sound, and interesting for a wider imitation learning community. \n\n- The problem of generalization in imitation is studied in several different contexts. Environments can vary across pixel distributions (such as sim or real [1]), dynamics [2], reward functions and so on. It would be useful to investigate these cases independently, and in combination with each other. For example, situations where the dynamics and the reward function are the same but pixel distributions are different, does \\eta^e converge to a trivial solution ? \n\n- As the paper addresses dynamics invariance, it would be useful to specify the generalization context clearly (including the title).\n\n- Transfer learning and domain generalization are closely related to the proposed approach. Related work can be better grounded in literature, in particular with [7]. I am also not convinced if the paper is the first to tackle generalizable imitation learning in a batch setting.\n\n- It would be useful to report ablation studies. While the loss functions are well-motivated, they likely have competing objectives which can make the learning unstable. Results reporting the role of each loss function alone and in combination would give furhter insights about the efficacy of different terms.\n\n- Generalizing across different environments here means doing well on average, but the resulting imitation behvaior may not be optimal for a particular environment. Additional experiments to evaluate the imitation performance in a single environment vs multiple environments as a function of learned \\eta would be helpful.\n\n[1] DIRL: Domain-Invariant Representation Learning for Sim-to-Real Transfer, 2020\n\n[2] Off-Dynamics Reinforcement Learning: Training for Transfer with Domain Classifiers, 2020-21\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "see above. No negative societal impact.", "review_text": "", "extra_scores": {"ethical_concerns": "no", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposed a method which exploits principles of causal invariance from causality in order to develop an imitation learning algorithm which is more robust to spurious/distractor features in the expert trajectories. It shows the improved performance of this method in OpenAI Gym tasks and a healthcare dataset. The reviewers were unanimous in recommending acceptance, and no significant argument made by the lone weak accept has made me doubt this paper is worth publishing at NeurIPS. Without much difficulty, I can also recommend acceptance, and hope to see the minor improvements suggested by the reviewers in the final version of the paper.
Decision: Accept (Poster)