Invariant Causal Imitation Learning for Generalizable Policies
Year: 2022

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 independent 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.

heading: 1 Introduction
text: Strictly batch imitation learning aims to learn a policy that directly mimics the behaviour of experts, for which we only have access to a set of demonstrations: logged trajectories of observations and actions following the expert’s 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 as 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. In healthcare, learning a generalizable behaviour policy that could achieve expert performance in new environments is an important goal: As a means of providing clinical decision support, it could serve as an “individualized”
⇤Equal contribution.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
clinical guideline for actions that can be taken for different patients—especially in a hospital, region, or patient demographic from which we have no access to data during training. In this endeavor, a principal 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 actions [4–7]. Directly learning an imitation policy from such data may lead to learning those spurious associations, thereby failing to generalize to unseen environments, and perpetuating any biases in the expert’s behaviour.
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. In 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 ⌘e (which may differ across environments) and shared latent state representations s (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 ⌘e and s, but they may also be any non-linear transformation of them. We shall operate in the setting where there are no hidden confounders, i.e. that we observe all variables that are affecting the expert’s actions (and the next states that result from these actions).
In addition to spurious correlations, another difficulty stems from learning to 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 [11, 12].
Contributions: In this paper, we introduce Invariant Causal Imitation Learning (ICIL), a novel method that learns a causal representation of the expert’s actions—which is used to build a generalizable 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, ICIL learns an (shared) invariant causal representation as well as an (environment-specific) noise representation. This accommodates dynamics mismatch across environments, while allowing the imitation policy to be learned by conditioning 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 representations 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.

heading: 2 Related Works
text: 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—more broadly—(3) causality in sequential decision-making.
Strictly Batch Imitation Learning: The simplest approach to imitation learning in the batch setting is behaviour cloning [10] which uses standard supervised learning techniques to learn an imitation policy that minimizes the negative log-likelihood of the observed demonstrator actions. However, behaviour cloning suffers from distributional shift as the learnt imitation policy cannot recover if it reaches a state out-of the distribution of the expert demonstrations [11–14]. To overcome this problem, [1, 14] propose incorporate dynamics-awareness by adding regularization to behaviour cloning by using norm-based penalties on the sparsity of implied rewards. Alternatively, [2] uses a 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-based 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 a 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 a 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 a challenging, bi-level optimization, and [6] propose the IRM-v1 objective which is a practical version to optimize. Through this optimization, the IRM objective should learn a 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 we 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-mismatch between the imitator and expert, existing methods assume access to demonstrations from the target environment [17, 19, 20], assume access to online interaction or simulators in the different environments [16], or focus on the different problem of hidden confounding [18, 21]. 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 or more expert trajectories from the new task [22–27].
Causality in 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
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 performing/simulating the effect of interventions in the environment. Alternatively, [29] use ideas from Invariant Risk Minimization [6, 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 signal, and focus on learning the causal ancestors of 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.

heading: 3 Problem Formalism

heading: 3.1 Imitation Learning
text: We work in the standard Markov decision process (MDP) setting: Let an environment be given by e = (X , A, T, r, ), with observations x 2 X , actions a 2 A, transition function T 2 (X )X⇥A, reward function r 2 RX⇥A, and discount factor . Let ⇡ 2 (A)X be a policy with the induced occupancy measure ⇢⇡(x, a) = (1 ) P 1 t=0 t p(xt = x, at = a|xt⇠T (·|xt 1, at 1), at⇠⇡(·|xt))
of observations and actions, and let ⇢⇡(x) = P a2A ⇢⇡(x, a) be the observation occupancy measure.
Unlike in the reinforcement learning setting, where the aim is to learn a policy ⇡(· | x) 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 = {⌧i}Ni=1 from a demonstrator policy ⇡D, where each trajectory ⌧ ⇠ ⇡D = (xt, at, xt+1)t=0,... consists of a sequence of observation, action, next observation tuples that are sampled as at ⇠ ⇡D(· | xt) and xt+1 ⇠ T (· | xt, at).
The goal of imitation learning is to seek an imitation policy ⇡ that minimizes the following risk: R(⇡) = L(⇡,⇡D) (1)
where L is a choice of loss function. Now, if we were in the online setting, we would have access to the environment (or a simulator), with which we can interactively perform distribution matching by minimizing the divergence between the expert’s state occupancy ⇢D and the imitator’s state occupancy ⇢⇡ [15,32–34]. One example is to use the (forward) KL divergence: L(⇡,⇡D) = DKL(⇢D||⇢⇡) [34]. However, in the offline setting we have no further access to the environment. As noted above, the simplest solution is behaviour cloning (BC) [10, 35, 36], which minimizes the negative log-likelihood of the demonstrator’s actions. However, by disregarding the distribution of the expert’s observations, imitation policies learnt by BC often result in compounding error when deployed in practice [11–14].

heading: 3.2 Imitation Learning from Multiple Environments
text: Consider a family of environments M = {(X e, A, T e, re, ) | e 2 E} with observations xe 2 X e, actions a 2 A, transition function T e 2 (X )X⇥A, reward function re 2 RX⇥A, and discount factor . 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 a dataset of recorded trajectories from the expert policy ⇡D in a set of training environments Etrain⇢E , D = {{(⌧ei } Ne i=1 | e 2 Etrain}. Each trajectory ⌧e ⇠ ⇡D = (xet , at, xet+1)t=0,... consists of a sequence of environment specific observations, expert actions and next observations sampled as at ⇠ ⇡D(· | xet ) and xet+1 ⇠ T e(· | xet , at).
In the presence of multiple environments, our goal is to learn a policy ⇡ 2 (A)X that matches the expert behaviour in all possible environments E 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 2 E—that is, the policy should ideally minimize the imitation risk across them:
max e2E
R e(⇡) = Le(⇡,⇡D) (2)
where each Le explicitly depends on the characteristics of the environment e. Note that since we know nothing specific about E , it is difficult to optimize for this directly. That said, if 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 observations from different environments—on which the expert policy depends. Finding such a structure would let us discard irrelevant factors as inputs to the learnt policy, improving generalization [6,7,37]: Assumption 3.1. (Shared Latent Structure) Consider decomposing the observations xe 2 X e in each environment e 2 E into two components: an invariant representation s 2 S and noise terms ⌘ e 2 Z
e (i.e. spurious correlations), such that xe = q(s, ⌘e) for some invertible transformation q : S ⇥ Ze ! X e. There exists some q such that ⇡D only depends on s, and space S is non-empty.
In other words, we assume that the demonstrator’s policy ⇡D depends only on information that is shared across the environments, i.e. the state variables s are the causal parents of the expert action a ⇠ ⇡D(· | 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 variables. Thus, we allow different environments to have different p(x) marginals (as well as different p(a|x)). This allows environments to have different structure. The only requirement is that the environments are the same as far as the task is concerned. This means that there exists some S such that p(s) marginals should be the same (as well as p(a|s)). Learning such a representation 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 reward functions whatsoever, and that we must learn an imitation policy in a strictly batch setting.
Note that the latent structure induced by the state variables is shared across the different environments. This means that the transition dynamics for the state representation p(st+1 | st, at) 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 pe(⌘et+1 | ⌘et , at) are specific to each environment. Our goal, then, is to learn a generalizable policy ⇡—that is, one that depends only on s.
Structure of Environments: Second, to learn a 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 it 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 a 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 Z 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, in our setting p(x) and p(a | x) can differ between the multiple environments. However, Assumption 3.1 enforces that the environments and tasks are the same modulo noise, i.e. that there exists some non-empty S such that p(s) and p(a | s) are the same between them. In other words, we have a set of environments that are different (i.e. the dynamics of xt are different), but the task being performed by the agent is the same (i.e. the dynamics of st are the same). This setting applies to the case when lightning conditions in a 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 staying the same. We provide additional explanations and definitions of environment interventions in Appendix A.
Figure 2 shows a simple example where each observation xt represents a union of the causal parents of the action (state variables) st = {x1t , x2t } and the spurious correlations (noise variables) ⌘t = {x3t }. To satisfy Assumption 3.2, the different environments need to correspond to interventions on x3t . And to satisfy Assumption 3.1,
x 1 t and x2t must not be intervened on. Our aim is to find a representation s of the causal parents {x 1 t , x 2 t } of the actions, as well as the mapping between them and the actions at, that mimics the expert’s policy.
Finally, similarly to [7], we also assume that the observations xt at timestep t can only affect the actions at and the observations at the next timestep t + 1: Assumption 3.3. (Temporal Causal Mechanism) Let xi and xj be any two components of the observation x at timestep t. Then:
x i t+1 ?? x j t+1 | xt, at (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 xt at time t will contain the causal parents of the action at, and that xt and at are the only factors that determine the next observation xt+1.

heading: 4 Invariant Causal Imitation Learning for Domain Generalization
text: The goal of our Invariant Causal Imitation Learning (ICIL) algorithm is to learn a 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 2 E given the above structural assumptions.

heading: 4.1 Learning Invariant Causal Representations
text: To achieve our goal, we decompose the observations xet in each environment e into a representation st = (xet ; ✓s) for the causal features of the action at, and another representation ⌘et = µe(xet ; ✓e⌘) for the noise variables, where ✓s and ✓e⌘ 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, µe : X ! Ze is environment-specific in order to allow for dynamics mismatch of the noise variables between the different environments.
In order to satisfy the causal diagram in Figure 1 and to learn a minimal causal representation, we need the following conditions to be satisfied: (1) st should be invariant across the environments, (2) st and ⌘et should be dynamics-preserving, and (3) st and ⌘et should be independent from each other.
Firstly, to fulfill condition (1) we train an environment classifier on the shared state representation cs : S ! |Etrain|, parameterized by ✓c using the cross-entropy loss. Similarly to [7], in order to build a state representation that is invariant across domains, we use an adversarial loss [41] that maximizes the entropy of the classifier: H(cs( (xt; ✓s); ✓c)). This gives us the following practical loss function:
Linv(✓s) = X
e2Etrain
Exet⇠⇢eD H(cs( (x e t ; ✓s); ✓c)) (4)
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 dynamics-preserving, we also learn the transition dynamics for the state variables gs : S ⇥ A :! S , such that ŝt+1 = gs(st, at; ✓gs) and the environment specific transition dynamics for the noise variables: ge⌘ : Ze ⇥A ! Ze, such that ⌘̂et+1 = ge⌘(⌘et , at; ✓eg⌘ ). To reconstruct xt+1 we also learn : S ⇥ Z
e ! X such that x̂et+1 = (st+1, ⌘et+1; ✓ ). This yields the following practical loss function:
Ldyn(✓s, ✓gs , {✓ e ⌘, ✓ e g⌘}e2Etrain , ✓ ) =
X
e2Etrain
Exet+1,at,xet+1⇠⇢eDkx e t+1 x̂ e t+1k 2 (5)
Note that while an alternative approach could consider directly building an invertible mapping from x
e to (s, ⌘e), the motivation for decoding st+1 and ⌘et+1 into x̂et+1 is twofold. In addition to learning dynamics-preserving representations, as we will see in Section 4.2, this also allows us to compute the energy of the next state obtained by following the imitation policy and enforcing this to be similar to the distribution of states visited by the expert’s policy.
Finally, to ensure that the state representation and the noise representation are marginally independent per condition (3), we minimize the mutual information between them. We use the Mutual Information Neural Estimation (MINE) framework [42], which provides a way for estimating the mutual information using neural networks. In particular, MINE uses a neural information measure I(U, V ) to approximate the mutual information between random variables U and V . Let T✓m be a statistics network parametrized by ✓m. MINE estimates I(U, V ) by ascending the gradient of the following:
I(U, V ) = sup ✓m EP(n)UV [T✓m ] log(EP(n)U ⌦P(n)V [e T✓m ]) = sup ✓m I(U, V ; ✓m) (6)
where PUV is the joint measure of (U, V ) and PU = R V dPUV , PV = R U
dPUV 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 I(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:
Lmi(✓s, {✓ e ⌘}e2Etrain) =
X
e2Etrain
Exet⇠⇢eDI( (x e t ; ✓s), µ(x e t ; ✓ e ⌘); ✓m) (7)
Note that the parameters ✓m of the statistics network T✓m used for computing the mutual information are updated through gradient ascent on I(U, V ; ✓m).

heading: 4.2 Matching Expert Behaviour in a Strictly Batch Setting
text: On the basis of the causal representation s, we shall learn a generalizable policy ⇡ (parameterized by ✓⇡) in the strictly batch setting, such that it matches the demonstrator’s behaviour. To begin, we first condition ⇡ on the representation st and minimize the negative log-likelihood of the expert’s actions:
L⇡(✓⇡, ✓s) = X
e2Etrain
Exet ,at⇠⇢eD log ⇡(at | (x e t ; ✓s); ✓⇡) (8)
However, having only this objective corresponds to performing behaviour cloning, which has wellknown limitations [11–14]. To mitigate compounding error, we want some form of added regularization to incentivize the imitation policy to stay within the distribution of the expert’s observations.
In the online setting, a 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
of divergence between their induced occupancy measures. However, this requires interactive access to the real environment or simulator to perform rollouts of intermediate policies—which is not possible in our setting. Instead, we propose a method that takes advantage of the learnt transition dynamics. For any current observation xt ⇠ ⇢D, we shall encourage the next observation x̄t+1 obtained by following the imitation policy āt ⇠ ⇡(· | xt) to remain within the occupancy measure of the expert.
Consider approximating the expert’s occupancy measure using an Energy Based Model (EBM) such that ⇢D(x) = exp( E✓̄(x)) Z(✓̄)
where the function E✓̄(x) : X ! R is the energy function and Z(✓̄) = R x E✓̄(x)dx is the partition function. We parameterize E✓̄ by a neural network. It is not possible to train the EBM directly through maximum likelihood because Z(✓̄) involves integrating over the entire input domain of x which is impractical. Instead, we use contrastive divergence to pre-train the energy function E✓̄ [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 occupancy distribution. Refer to Appendix B 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(✓⇡; ✓s, ✓gs , {✓ e ⌘, ✓ e g⌘}e2Etrain , ✓ ) =
X
e2Etrain
E xet⇠⇢eD st= (x e t ),⌘ e t =µ
e(xet ) āt⇠⇡(·|st)
x̄t+1= (gs(st,āt),g e ⌘(⌘ e t ,āt))
E✓̄(x̄ e t+1) (9)
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 xt ⇠ ⇢⇡ .
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.

heading: 5 Experiments
text: We perform experiments on OpenAI 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 ICIL1 against standard methods for strictly batch imitation learning: Behaviour Cloning (BC) [10]; Reward-regularized Classification for Apprenticeship Learning (RCAL), which incorporates dynamics-awareness through a sparsity regularization on the implied rewards [14]; ValueDICE (VDICE) [2], which uses an off-policy objective to estimate distribution ratios needed for distribution matching; as well as Energy-based Distribution Matching (EDM) [3], which jointly 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 we 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 we used the invariance-based 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.

heading: 5.1 Evaluation on OpenAI Gym
text: We perform experiments on the following control tasks from OpenAI 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 OpenAI Baselines [53] to 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 LunarLander [47] we add spurious correlations to the state space of each control task and an environment identifier. The
1The code for ICIL can be found at https://github.com/vanderschaarlab/mlforhealthlabpub and at https://github.com/ioanabica/Invariant-Causal-Imitation-Learning.
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 on an environments with multiplicative factors sampled from U( 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, while the test environment has no 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 on 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 on 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, we generally found that using the IRMv1 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 on overall performance, compare performance on train vs. test environments and also evaluate robustness to increasing the size of the spurious correlations.

heading: 5.2 Evaluation on MIMIC-III
text: 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 a generalizable policy for the action of putting patients on the mechanical ventilator.
We define three environments, two for training and one for testing, each consisting of 2000 independent 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 a 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 a 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.
Since MIMIC-III is an entirely offline dataset, it is not possible to compute average returns for running the policies in the test environment. Instead, we evaluate the benchmarks in terms of 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-recall curve (APR) and their standard deviations over 10 runs. We notice that ICIL learns a 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 actions on the test environment.

heading: 6 Discussion
text: 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, we 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 at least two training environments with different interventions on the noise variables (spurious correlations), and the method cannot be used if such data is 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 societal 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.

heading: Acknowledgments
text: We would like to thank the reviewers for their valuable feedback. The research presented in this paper was 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.

Review summaries:
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.

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.  

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.