Invariant Causal Imitation Learning for Generalizable Policies

1 Introduction . 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 can not 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 . 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—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 can not 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 . 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 , 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 ] .

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.