I NVARIANT C AUSAL R EPRESENTATION L EARNING FOR G ENERALIZATION IN I MITATION AND R EINFORCEMENT L EARNING
A BSTRACT
A fundamental challenge in imitation and reinforcement learning is to learn policies, representations, or dynamics that do not build on spurious correlations and generalize beyond the specific environments that they were trained on.
We investigate these generalization problems from a unified view.
For this, we propose a general framework to tackle them with theoretical guarantees on both identifiability and generalizability under mild assumptions on environmental changes.
By leveraging a diverse set of training environments, we construct a data representation that ignores any spurious features and consistently predicts target variables well across environments.
Following this approach, we build invariant predictors in terms of policy, representations, and dynamics.
We theoretically show that the resulting policies, representations, and dynamics are able to generalize to unseen environments.
Extensive experiments on both synthetic and real-world datasets show that our methods attain improved generalization over a variety of baselines.
I NTRODUCTION
Imitation learning (IL) (Argall et al., 2009) concerns learning a policy via expert demonstrations without access to a reward function, while reinforcement learning (RL) (Sutton & Barto, 2018) focuses on learning a policy via interaction and reward feedback.
One of the fundamental challenges with state-of-the-art IL and RL approaches is their limited ability to generalize outside of the specific environments they were trained on, because they learn easy-to-fit spurious correlations which are prone to change between training and testing environments.
A natural approach to generalization across environmental changes is constructing a representation that consistently predicts the target well across environments.
That is, such a representation comprises only those features that describe true correlations of interest with the target that are stable across environments, excluding any features presenting environment-varying spurious correlations with the target.
By exploiting the varying degrees of spurious correlations naturally present in training data collected from multiple environments, one can try to identify stable features and build invariant predictors (e.g., invariant policies and dynamics) that are able to generalize to unseen environments (Peters et al., 2016; Arjovsky et al., 2019; Lu et al., 2022).
In both IL and RL, much effort has been made in this direction.
Specifically, in the IL setting, prior work has shown that only a policy relying solely on the true causes of expert actions can robustly model the mapping from states to optimal/expert actions, which is stable under environmental shift (de Haan et al., 2019;
Zhang et al., 2020c; Bica et al., 2021;
Samsami et al., 2021; Kumor et al., 2021).
In the RL setting, one usually considers three types of generalization problems.
First, it has been shown that a representation consisting only of all the causal ancestors of the reward is the minimal sufficient representation for policy learning (Zhang et al., 2020a; Huang et al., 2022; 2021).
This is because in RL, we seek to model return (i.e., cumulative reward) rather than solely rewards, which requires a representation that can capture multi-timestep interactions.
Such a representation can generalize to unseen testing environments with unseen reward functions, as long as the new reward functions are causally dependent on a subset of the same causal ancestors that determine the original reward function in training environments (Zhang et al., 2020b; Huang et al., 2022).
Second, in some scenarios where testing and training environments further share the same reward function, the policy relying solely on such representation, which is learned from training environments, can even directly generalize to unseen testing environments (Higgins et al., 2017; Harrison et al., 2020; Zhang et al., 2020b).
Third, when the rewards between training and testing environments neither share the same function nor depend on a subset of the same causal ancestors, it is generally impossible to generalize either representations or policies.
However, it is still possible to learn generalizable (local) dynamics models by building invariant predictors on a per-state-variable level 1 (Tomar et al., 2021).
Apparently, this also relies on identifying the true causes of each state variable from those state variables in the previous timestep.
In this paper, to distinguish the three types of RL generalization problems above, we call them representation generalization , policy generalization , and dynamics generalization , respectively.
A graphical illustration is presented in Fig.
3 of the appendix.
In this work, we attempt to look at the different generalization problems in IL and RL from a unified view, and propose a general framework, under more relaxed assumptions over the distributional shift and the underlying causal structures, to tackle them, with theoretical guarantees on both identifiability and generalizability.
The key idea is that by taking advantage of structural relationships between environmental variables (i.e., observations, states, actions, and rewards), we formulate the generalization problems in both IL and RL in the framework of invariant causal representation learning (iCaRL)
(Lu et al., 2022).
This framework offers a tool, with theoretical guarantees under rather general assumptions over the underlying causal diagram, to first identify direct causes of a given target from data and then use those causes to build invariant predictors that are able to generalize to unseen testing environments, which will be briefly reviewed in Appendix C.1.
While our methodology builds significantly on the supervised learning iCaRL framework, the application of this methodology to problems in RL and IL is new and the results obtained are promising.
Our contributions are summarised as follows: (1) We investigate different generalization problems in
IL and RL in a unified perspective; (2) We propose a general framework to tackle the generalization problems in IL and RL, with theoretical guarantees on both identifiability and generalizability; (3)
We propose general assumptions over the distributional shift and the underlying causal structures, covering many real-world scenarios in IL and RL; (4) We show that our framework has theoretical guarantees for OOD generalization in terms of policy, representation, and dynamics in IL and RL.
G ENERALIZATION IN I MITATION L EARNING 2. P ROBLEM F ORMULATION
We consider the BC approach (Widrow, 1964; Pomerleau, 1989; 1991; Bain & Sammut, 1995; Schaal, 1999; Muller et al., 2006; M ¨ ulling et al., 2013; Bojarski et al., 2016; Mahler & Goldberg, 2017;
Bansal et al., 2018) to learning an imitation policy on the basis of expert demonstrations collected from multiple environments, with the aim to generalize it to unseen environments.
Technically, we consider a family of environments M E all =
{ ( X e , A , P e , R e , γ ) | e ∈ E all } , with observations x e ∈ X e ⊆ R d , actions a ∈ A , a transition dynamics P e ≡ p e (( x ′ )
e | x e , a ) , a reward function
R e ( x e , a ) ∈ R , and a discount factor γ ∈ [0 , 1) .
Note that, the action space and discount factor do not change across all environments E all .
We assume only access to a dataset of recorded demonstrations
D E tr = {{ τ ei } N e i =1 | e ∈ E tr } from a set of training environments E tr ⊂ E all , where each demonstration τ e consists of a sequence of environment specific observation-action pairs
τ e = ( x et , a t ) t =0 ,... that are drawn from an expert policy π ∗ .
Our goal is to learn a policy π from D E tr so that it is able to mimic the expert behaviour across E all that share a certain structure.
Specifically, we seek a policy that generalizes well across E all by solving the optimization problem:
[Equation]
(1) where each ℓ e is a choice of environment specific loss function and D e
= { τ ei } N e
i
=1 is a set of expert demonstrations collected from the training environment e ∈ E tr .
Unless stated otherwise, for simplicity of notation, we drop the superscript e when referring to the union over all the environments
E all in the rest of this paper.
As discussed in Section 1, prior work has shown that only a policy replying solely on the true causes of expert actions can generalize well to E all .
Hence, solving (1) is reduced to discovering the causal features for expert actions from observations across E tr , and then, using them, build an invariant policy.
In the IL context, this can be also interpreted in the way that a generalizable policy depends only on some shared components of the true latent states denoted by s t , rather than their corresponding observations x t , of the environments (cf.
Fig. 3).
Note that, in this work we do not address the partial-observability problem explicitly (Thrun et al., 2005).
Instead, we approximate stacked consecutive observations as the observation x t , that is, x t loses no information about s t (Hausknecht & Stone, 2015; de Haan et al., 2019; Gelada et al., 2019; Zhang et al., 2020b):
Assumption 1 (Full-observability) .
Observations x contain all the information about states s .
2.1. A SSUMPTIONS ON THE C AUSAL D IAGRAM
[Figure]
We denote the data representation (true state) of observation x t by s t =
( s p 1 t , . . .
, s p r t , s c 1 t , . . .
, s c k t ) ∈ R n , where n = r + k , and { s it } i ∈ I p .
= { p 1 ,...,p r } and { s jt } j ∈ I c .
= { c 1 ,...,c k } are multiple scalar causal factors and non-causal factors of a t , respectively.
We denote s pt .
= ( s p 1 t , . . .
, s p r t ) and s ct .
=
( s c 1 t , . . .
, s c k t ) for the ease of clarification.
It is assumed that s is of lower dimension than x , that is, n ≤ d .
We assume that all the environments
E all share some latent structure between x t , s t and a t , and consider different degrees to which this structure may be shared, which is encapsulated in Fig. 1 and summarized as below.
Note that, at this moment we do not explicitly consider the node c t − 1 .
= ( s t − 1 , a t − 1 ) and its both incom- ing and outgoing arrows (marked in blue), and defer them to Appendix F.1.
Assumption 2.
(a) The causal graph containing s t and a t is a DAG; (b) x t ⊥⊥ a t , e | s t , implying that p ( x t | s t ) is invariant across E all ; (c) a t ⊥⊥ e | s pt , implying that p ( a t | s pt ) is invariant across E all .
Assumption 2, together with its corresponding causal diagram in Fig.
1, is flexible enough to cover most scenarios in the IL setting, which is detailedly explained in Appendix F.2.
2.1. A SSUMPTIONS ON THE P RIOR
When the underlying causal diagram across E all satisfies Assumption 2b, it is straightforward to obtain the following primary assumption over the prior p ( s t | a t , e ) leading to identifiability of the latent variables s t by directly substituting ( z , y ) with ( s t , a t ) in Assumption 2 of Lu et al.
(2022) (cf.
Assumption 10 in Appendix C.1).
Assumption 3.
p T , λ ( s t | a t , e ) belongs to a general exponential family with parameter vector given by an arbitrary function λ ( a t , e ) and sufficient statistics T ( s t ) =
[ T f ( s t ) T , T NN ( s t ) T ] T given by the concatenation of a)
the sufficient statistics T f
( s t ) =
[ T 1 ( s 1 t ) T , · · · , T n ( s nt ) T ] T of a factorized exponential family, where all the T i ( s it ) have dimension larger or equal to 2, and b) the output
T NN ( s t ) of a neural network with ReLU activations.
The resulting density function is thus given by
[Equation]
where Q is the base measure and Z the normalizing constant.
A neural network with ReLU activation has universal approximation power.
Therefore, the term
T NN ( s t ) in Assumption 3 will allow the prior to capture arbitrary dependencies between the latent variables s t .
This is important because for any two s it and s jt , we usually have that s it ̸⊥⊥ s jt | a t , e .
The conditional dependencies between them are mainly due to the confounder c t − 1 (cf.
Fig.
1) (de Haan et al., 2019).
Another possibility is that there might exist instantaneous effects between s it and s jt (i.e., an arrow between s it and s jt ) in some scenarios (Peters et al., 2017; Sutton & Barto, 2018), which is not considered in previous work (Bica et al., 2021).
Now we are ready to address the policy generalization problem in IL, by following the three phases in iCaRL described in Appendix C.1.
For ease of reference, this approach is called iCaRL-IL , which is described in Appendix F.3 and whose corresponding theoretical results are given in Appendix H.
G ENERALIZATION IN R EINFORCEMENT L EARNING 3. P ROBLEM F ORMULATION
Unless stated otherwise, we follow the same notations used in the IL setting as described in Section 2.1.
In the RL setting 2 , we consider three generalization problems: policy generalization, representation generalization, and dynamics generalization.
As analysed in Section 1, they all can be resolved in the iCaRL framework.
Specifically, we first identify the true states s t from observations x t (Phase 1), then discover the causal ancestors of rewards or the true causes of each state variable among the identified states s t (Phase 2), and finally base on them to build invariant predictors in terms of policy, representation, and dynamics (Phase 3).
Apparently, all the three generalization problems are same in Phase 1 but not in Phases 2&3.
Hence, we first describe what they have in common in Phase(Section 3.1) and then separately describe Phases 2&3 for each (Sections 3.2,3.3&3.4).
3.1. A SSUMPTIONS ON THE C AUSAL D IAGRAM
Unlike in the IL setting where the reward is not given, we assume here without loss of generality that s pt and s ct are multiple scalar causal factors and non-causal factors of r t
+1 .
We also assume that all the environment E all share some latent structure as encapsulated in Fig.
4 of the appendix and summarized below.
Note that, since a t is always a causal factor of r t +1 , it is not placed in the figure for simplicity.
Assumption 4.
(a) The causal graph containing s t and r t +1 is a DAG; (b) x t ⊥⊥ r t +1 , e | s t , implying that p ( x t | s t ) is invariant across E all .
The practicality of Assumption 4 can be detailedly explained in a similar way to that of Assumption 2.
A more in-depth explanation can be found in Appendix G.1.
3.1. A SSUMPTIONS ON THE P RIOR
When the underlying causal diagram across E all satisfies Assumption 4b, it is straightforward to obtain the primary assumption (Assumption 12 in Appendix E) over the prior p ( s t | r t +1 , e ) leading to identifiability of the latent variables s t by substituting a t with r t +1 in Assumption 3.
Under Assumptions 4b&12, we can follow the exact same steps as described in Appendix F.3.
to identify the latent variables s t with similar theoretical guarantees on identifiability, by directly substituting a t with r t +1 in all the corresponding equations and theorems.
This resulting NF-iVAE model in the RL setting is called NF-iVAE-RL .
3. P OLICY G ENERALIZATION
After estimating s t for x t in Phase 1, let us first consider the ideal generalization case in which an optimal policy learned from multiple training environments can zero-shot generalize to unseen testing environments.
As discussed in Section 1, it has proved that a representation only consisting of all the causal ancestors of the reward is the minimal sufficient representation (MSR) for policy learning.
Hence, in Phase 2 we need to discover the causal ancestors of the reward, denoted by s At ≡ AN ( R ) , from the identified state s t .
Due to Markovianity of the dynamics in RL, this can be implemented by recursively applying the two-step method described in Appendix F.3.2.
Precisely, we first use the method to discover the direct causes s pt of the reward r t +1 from s t and then reuse it to find the direct causes of s pt +1 from s t , whose union produces s At .
After obtaining s At , to learn an invariant policy that generalizes well to unseen environments in Phase 3, we need to further assume that all the environments E all share the following latent structure:
Assumption 5.
(a) p ( r t +1 | s pt , a t ) is invariant across E all ; (b) p ( s At +1 | s At , a t ) is invariant across E all .
Assumption 5a states that all the environments share the same reward function.
Assumption 5b means that the latent structure induced by MSR s A t is shared across E all .
Thus, it is straightforward that under Assumption 5, the learned policy across E tr is guaranteed to generalize well to E all .
We therefore have the following result with proof in Appendix H.
Proposition 1.
Under Assumptions 4,12&5 and the assumptions of Theorems 7&8, the policy learned across E tr in the limit of infinite data has optimal OOD generalization across E all .
Note that, unlike the IL setting where we learn the invariant policy in the supervised way, in the RL setting we follow Zhang et al.
(2020b)
and combine our MSR s
At with the soft actor-critic (SAC) algorithm (Haarnoja et al., 2018) to devise a practical RL method, termed iCaRL-RL-P , as shown in
Algorithm 1 of the appendix (see more in Appendix G.3).
3. R EPRESENTATION G ENERALIZATION
When Assumption 5 is violated, it is theoretically impossible that a policy learned from E tr can generalize well to E all .
As discussed in Section 1, however, it is still possible that the representation function Φ learned from E tr can generalize to unseen testing environments E te , as long as the new reward functions across E te are causally dependent on a subset of the same causal ancestors s A that determine the original reward function across E tr .
This condition can be formalized as follows:
[Figure]
[Table]
Assumption 6.
AN ( R e ∗ ) =
s A for any e ∗ ∈ E te .
Then, we have the following result whose proof is in Appendix H.
Proposition 2.
Under Assumptions 4,12&6 and the assumptions of Theorems 7&8, the representation function Φ learned across E tr in the limit of infinite data is able to generalize to E te .
3. D YNAMICS G ENERALIZATION
When Assumption 6 is also not satisfied, then generalizing Φ is hopeless as well.
However, it is still possible to learn generalizable (local) dynamics models by building invariant predictors on a per-state-variable level.
As discussed in Section 1, the rationale behind is that in many world scenarios, each state variable only depends on a small subset of those state variables in the previous timestep, which is summarized as below.
Assumption 7.
Given states s t , s t +1 and action a t , we have p ( s t +1 | s t , a t ) =
(cid:81)
i p ( s it +1 | s P i t , a t ) , where s it +1 denotes the i -th dimension of state s t
+1 and s P
i t is a set of the parents of s it +1 in s t .
To learn invariant (local) dynamics models that generalize well to unseen environments, we need to further assume that all the environments E all share the following latent structure:
Assumption 8.
Given states s t , s t +1 and action a t , there exists some (local) dynamics model p ( s it +
| s P i t , a t ) that is invariant across E all .
Graphically, Assumption 8 indicates that s it +1 ⊥⊥ e | s P i t , as shown in Fig.
5 of the appendix.
Without loss of generality, we similarly assume that s P i t .
=
( s p
i 1 t , . .
, s p ir t ) and s C i t .
=
( s c i 1 t , . . .
, s c ik t ) are multiple scalar causal factors and non-causal factors of s it +1 .
After estimating s t for x t in Phase 1, it is straightforward to use the two-step approach described in Appendix F.3.2 to discover the direct causes s P
i t of s it
+1 from s t .
Then, under Assumption 8, we can follow the same step depicted in
Appendix F.3.3 to learn the invariant (local) dynamics model p ( s it +1 | s P i t , a t ) in Phase 3 by directly substituting s pt with ( s P i t , a t ) and a t with s it +1 in Eq.
This approach is called iCaRL-RL-D .
We further have the following result (proof in Appendix H).
Proposition 3.
Under Assumptions 4,12,7&8 and the assumptions of Theorems 7&8, the (local) dynamics learned across E tr in the limit of infinite data has optimal OOD generalization across E all .
E XPERIMENTS
We compare our approach with a variety of baselines on both synthetic and real-world datasets.
In all comparisons, unless stated otherwise, we average performance over ten runs and show the mean results with standard deviations.
Due to space limit, we only highlight some key results as shown in
Fig.
2 on the widely used control tasks and refer readers to Appendix I for more details.
Our methods consistently attain improved generalization over a variety of baselines in both IL and RL settings.
C ONCLUSION AND R ELATED W ORK
We investigated the different generalization problems in terms of policy, representation, and dynamics in IL and RL from a unified view, proposing a framework to tackle them with theoretical guarantees on both identifiability and generalizability under mild assumptions over environmental changes.
Experimental results show that our methods attain improved generalization over a variety of baselines.
In Appendix B, we discuss the related work on different generalization problems in both IL and RL.
A B ASICS OF I MITATION AND R EINFORCEMENT L EARNING
We describe an environment as a standard Markov decision process (MDP) (Puterman, 1994; Sutton & Barto, 2018) given by a tuple M =
( X , A , P, R, γ ) , with observations x ∈ X ⊆ R d , actions a ∈ A , a transition dynamics P ≡ p ( x ′ | x , a ) , a reward function R ( x , a ) ∈ R , and a discount factor γ ∈ [0 , 1) .
A policy π ( ·| x ) defines a distribution over actions conditioned on the observation x .
At time t , an agent is provided with an observation x t ∈ X and chooses an action a t ∈ A according to a policy a t ∼ π ( x t ) .
The agent receives a reward r t +1 =
R ( x t , a t ) and then the environment yields next observation x t +1 ∼ p ( ·| x t , a t ) .
Under a policy π , the value function , denoted by V π ( x ) , represents the expected cumulative discounted rewards (a.k.a., return) following π from observation x , defined by V π ( x ) =
E π [ (cid:80) ∞ t =
γ t r t +1 | x 0 = x ] .
Similarly, we also can define the action-value function for policy π , denoted by Q π ( x , a ) , meaning that the value of taking action a in observation x under a policy π : Q π ( x , a ) =
E π [ (cid:80) ∞ t =
γ t r t +1 | x 0 = x , a 0 = a ] .
In the RL setting, the agent’s goal is to find an optimal policy π ∗ that maximizes the value function: π ∗ = arg max π V π ( x ) for all x ∈ X .
s p i 1 t · · · s p ir t s i t + s c i 1 t · · · s c ik t x t e
In the IL setting, the primary difference from the previous RL setting is that the reward is neither known nor observed.
Instead, we only have access to a dataset of expert demonstrations D = { τ i } Ni , where each demonstration τ consists of a sequence of observation-action pairs τ =
( x t , a t ) t =0 ,... that are drawn from an expert policy π ∗ .
We consider behavioural cloning (BC) approaches (Widrow, 1964; Pomerleau, 1989; 1991;
Bain & Sammut, 1995; Schaal, 1999; Muller et al., 2006; M¨ulling et al., 2013; Bojarski et al., 2016; Mahler & Goldberg, 2017; Bansal et al., 2018) that reduce policy learning to supervised learning by training a discriminative model to predict expert actions given observations.
That is, the goal of BC is to seek a policy π that imitates the expert policy π ∗ by solving the optimization problem:
[Equation]
(3) where ℓ is a choice of loss function.
B R ELATED W ORK
Generalization in IL.
IL through behavioural cloning has been extensively studied (Widrow, 1964;
Pomerleau, 1989; 1991; Bain & Sammut, 1995; Schaal, 1999; Muller et al., 2006; Kober et al., 2010;
Bojarski et al., 2016; Mahler & Goldberg, 2017; Bansal et al., 2018).
However, none of existing methods consider the problem of learning policies robust to spurious correlations so that they can generalize to unseen environments.
There are also some work studying the problem of domain adaptation and transfer learning for IL.
However, they assume access to either demonstrations from testing environments (Sermanet et al., 2017;
Liu et al., 2018;
Kim et al., 2020) or online interaction (de Haan et al., 2019; Lu & Tompson, 2020; Swamy et al., 2021), or they study the different problem of hidden confounding (de Haan et al., 2019; Etesami & Geiger, 2020; Zhang et al., 2020c; Kumor et al., 2021; Swamy et al., 2021).
Another line of work is in the field of meta-learning whose goal is to generalize learned policies to new tasks, but they also require access to demonstrations from the new tasks (Finn et al., 2017a;b; Duan et al., 2017;
Yu et al., 2018; James et al., 2018; Sharma et al., 2019).
Perhaps the most related to ours is the one in Bica et al.
(2021) that also tackles the generalization problem in IL by learning a representation shared across environments.
However, they have no theoretical guarantees on both identifiability and generalizability.
Also, their assumptions over the underlying causal graph are restricted, e.g., they assume that the causal factors s pt and non-causal factors s ct must be independent from each other, which is not necessary in ours; etc.
Policy Generalization in RL.
The goal of policy generalization in RL is to learn policies from multiple training environments so that they can zero-shot generalize to unseen testing environments.
The most common approach to policy generalization is by adding different forms of regularization, such as dropout and batch normalization (Cobbe et al., 2019) and information-theoretic regularizer
(Goyal et al., 2018; Pacelli & Majumdar, 2020).
Although easy to implement, these methods do not explicitly exploit any (causal) structure of the RL problem.
Data augmentation and domain randomization also show their potential in the sim-to-real generalization problem (Akkaya et al., 2019; Peng et al., 2018;
Urakami et al., 2019), but they are complementary to our methods and could be potentially used to generate a diverse set of training environments for ours.
While the approaches based on PAC-Bayes (Majumdar et al., 2021) and adversarial perturbations (Sinha et al., 2018) provide provable generalization guarantees, they require an a priori bound on how much the test environments differs from the training environments.
Inspired by the work of Arjovsky et al.
(2019) and Ahuja et al.
(2020), Krueger et al.
(2021) propose a risk-extrapolation method and Sonar et al.
(2021) propose an approach referred as to invariant policy optimization for OOD generalization in RL, but both have no theoretical guarantees on identifiability and generalizability.
Saengkyongam et al.
(2021) tackle the problem of environmental shifts in offline contextual bandits from a causal perspective.
There are also much work on quantifying generalization in RL (Mnih et al., 2013; Nichol et al., 2018;
Song et al., 2019; Cobbe et al., 2019; 2020; Ahmed et al., 2020; Kirk et al., 2021).
Representation Generalization in RL.
This line of work focuses on representation generalization in RL that aims to learn representations from multiple training environments so that they can zeroshot generalize to unseen testing environments.
Much work has been done on reconstruction-based representations (Lange & Riedmiller, 2010; Lange et al., 2012;
Watter et al., 2015; Hafner et al., 2019;
Gelada et al., 2019; Huang et al., 2021), contrastive-based representations (Oord et al., 2018;
Chen et al., 2020; Laskin et al., 2020), and bisimulation-based representations (Larsen & Skou, 1989;
Taylor et al., 2008; Ferns et al., 2011; Ferns & Precup, 2014; Castro, 2020).
However, none of these methods explicitly consider the OOD generalization problem in terms of representation across environments.
Zhang et al.
(2020a) propose a method of invariant prediction to learn model-invariance state abstractions that generalize to novel observations in the multi-environment setting, but their method is limited to a family of environments represented by block MDPs (Du et al., 2019) and also has no theoretical guarantees on identifiability and generalizability.
Lee et al.
(2021) propose an approach for structure and transfer learning of robot manipulation policies, but it requires performing/simulating the effect of interventions in the environment.
Agarwal et al.
(2021) incorporate the inherent sequential structure in RL into the representation learning process to improve generalization.
Zhang et al.
(2020b) leverage bisimulation metrics to learn generalizable representations which encodes only the task-relevant information from observations.
Similarly, both of the methods above provide no theoretical guarantees on identifiability and generalizability.
Dynamics Generalization in RL.
For the tasks of dynamics generalization in RL, the goal is to learn (local) dynamics models from multiple training environments so that they can zero-shot generalize to unseen testing environments.
Boutilier et al.
(1999) detailedly explore structural assumptions and computational leverage on the underlying MDP for decision-theoretic planning.
The most related setting is of factored MDPs, but most existing work in this field is either to assume a known causal structure for the transition dynamics (Kearns & Koller, 1999; Jonsson & Barto, 2006;
Strehl et al., 2007; Osband & Van Roy, 2014) or to not learn states abstractions (Kearns & Koller, 1999; Strehl et al., 2007; Misra et al., 2020).
Hallak et al.
(2015) discuss learning the factored structure in the dynamics of the environment under the factored MDP assumption.
Volodin et al.
(2020) consider the problem of inferring a causal model of the environment, but intervention is required.
Note that, none of these methods consider the generalization problem in terms of dynamics.
Also note that, most factored MDP works include the factored reward condition, which is not necessary in ours.
Another thing in which we differ from these methods is that we only focus on generalizing the local dynamics model on a per-state-variable level.
The most related to our work is the one in Tomar et al.
While attempting to learn generalizable local dynamics model, they have no theoretical guarantees on identifiability and generalizability.
C V ARIATIONAL A UTOENCODERS
We briefly describe the framework of variational autoencoders (VAEs), which allows us to efficiently learn deep latent-variable models and their corresponding inference models (Kingma & Welling, 2013; Rezende et al., 2014).
Consider a simple latent variable model where x ∈ R
d stands for an observed variable and z ∈ R n for a latent variable.
A VAE method learns a full generative model p θ ( x , z ) =
p θ ( x | z ) p θ ( z ) and an inference model q ϕ ( z | x ) , typically a factorized Gaussian distribution whose mean and variance parameters are given by the output of a neural network with input x .
This inference model approximates the posterior p θ ( z | x ) , where θ is a vector of parameters of the generative model, ϕ a vector of parameters of the inference model, and p θ ( z ) is a prior distribution over the latent variables.
Instead of maximizing the data log-likelihood, we maximize its lower bound
L VAE ( θ , ϕ ) :
[Equation]
where we have used Jensen’s inequality, and KL ( ·||· ) denotes the Kullback-Leibler divergence between two distributions.
[Figure]
C. I NVARIANT C AUSAL R EPRESENTATION L EARNING ( I C A RL)
Invariant causal representation learning (iCaRL) (Lu et al., 2022) is a framework for out-of-distribution (OOD) generalization in supervised learning (SL), with theoretical guarantees on both identifiability and generalizability.
It aims to construct an invariant predictor that performs well across all environments E all by expploiting data collected from multiple training environments E tr , where E tr ⊂ E all .
The scenarios for OOD generalization in SL that iCaRL mainly considers are the ones whose underlying data generating process can be represented by a general causal diagram as shown in Fig.
In particular, the input data is x ∈ R d and its corresponding target is y ∈ R s . {+
+}
The data representation z
=
( z p 1 , . . .
, z p r , z c 1 , . . .
, z c k ) ∈ R
n are a set of latent variables encoding the input data, where n = r + k , and { z i } i ∈ I p .
= { p 1 ,...,p r } and { z j } j ∈ I c .
= { c 1 ,...,c k } are multiple scalar causal factors and non-causal factors of y , respectively.
We denote z p .
=
( z p 1 , . . .
, z p r ) and z c .
=
( z c 1 , . . .
, z c k ) for the ease of clarification.
It is assumed that z is of lower dimension than x , that is, n ≤ d .
The environment is treated as a random variable e , 4 where e could be any information specific to the environment (Storkey, 2009; Peters et al., 2016; Zhang et al., 2017; Huang et al., 2020).
For simplicity, e is set to be the environment index, i.e., e ∈ { 1 , . . .
, N } , where N is the number of training environments.
It is worth noting that iCaRL allows for arbitrary connections between the latent variables z as long as the resulting causal diagram including y is a directed acyclic graph (DAG).
Dashed lines indicate the causal mechanisms which might vary across environments and even be absent in some scenarios, whilst solid lines indicate that they are invariant across all the environments.
To sum up, we posit that the underlying causal graph (Fig. 6) satisfies the following assumption (Lu et al., 2022):
Assumption 9.
(a
The causal graph
containing z and y
is
a DAG
(b) x ⊥⊥ y e | z , implying that p ( x | z )
is invariant
across E all
(c) y ⊥⊥ e | z p
, implying that p
( y | z
p
) is
invariant across E all . To learn the invariant predictor w that is a function mapping from z p to y , the iCaRL framework consists of three phases.
The idea is to first identify the latent variables ( z ) from the input data ( x ) by using an extended identifiable variational autoencoder (NF-iVAE) model under Assumptions 9b& (Phase 1), then discover direct causes ( z p ) of the target ( y ) among the identified latent variables under Assumption 9a (Phase 2), and finally learn the invariant predictor ( w ) for the target from the discovered causes under Assumption 9c (Phase 3).
We briefly review them below, cf.
also Lu et al.
C. P HASE 1: I DENTIFYING L ATENT V ARIABLES
Technically, under Assumption 9b, the proposed NF-iVAE in Phase 1 is a conditional generative model with the parameters θ = ( f , T , λ ) ∈ Θ , defined as
[Equation]
where ϵ is an independent noise variable with probability density function p ϵ ( ϵ ) .
Importantly, the primary assumption leading to identifiability of the latent variables in NF-iVAE is that the conditional prior p T , λ
( z | y , e ) belongs to a general exponential family with sufficient statistics T and natural parameters λ .
This is formalized as follows:
Assumption 10.
p T , λ ( z | y , e ) belongs to a general exponential family with parameter vector given by an arbitrary function λ ( y , e ) and sufficient statistics T ( z ) =
[ T f ( z ) T , T NN ( z ) T ] T given by the concatenation of a) the sufficient statistics T f ( z ) =
[ T 1 ( z 1 ) T , · · · , T n ( z n ) T ] T of a factorized exponential family, where all the T i ( z i ) have dimension larger or equal to 2, and b) the output
T NN ( z ) of a neural network with ReLU activations.
The resulting density function is thus given by
[Equation]
where Q is the base measure and Z the normalizing constant.
A neural network with ReLU activation has universal approximation power.
Therefore, the term
T NN ( z ) in the above prior distribution will allow us to capture arbitrary dependencies between the latent variables.
Most importantly, Lu et al.
(2022) prove that under Assumptions 9b&10 and the conditions stated in
Theorems 1-3 of Lu et al.
(2022), NF-iVAE can identify the latent variable z .
C. P HASE 2: D ISCOVERING D IRECT C AUSES
After estimating z for each data point, under Assumption 9a, Lu et al.
(2022) propose a heuristic two-step approach to discovering the direct causes z p in Phase 2, which works well in practice.
First, the PC algorithm (Spirtes et al., 2000) is conducted to learn a Markov equivalence class of DAGs, which outputs the direct neighbours of y , denoted by NE ( y ) .
Second, all pairs 5 in NE ( y ) are tested with conditional independence testing (Zhang et al., 2012) to discover z p by comparing p -values from the two tests: IndTest ( z i , z j | e ) and IndTest ( z i , z j | y , e ) , where IndTest denotes (conditional) independence test.
This is based on the observation that under Assumption 9, for any two z i and z j from NE ( y ) , only when both are causes of y does the dependency between them generally increase after additionally conditioning on y .
C. P HASE 3: L EARNING AN I NVARIANT P REDICTOR
After having obtained z p for y across E tr , under Assumption 9c, the invariant predictor w can be learned in Phase 3 by solving the following optimization problem:
[Equation]
where ℓ ( · ) could be any loss.
Since E [ y | z p ] is assumed to be invariant across E all , the learned w is guaranteed to perform well across E all .
When in an unseen testing environment, z p is required to be inferred from x before applying w for prediction.
This is implemented by leveraging the assumed invariant p f ( x | z ) through solving the following optimization problem: [Equation]
[Equation]
where the hyperparameters λ 1 > 0 and λ 2 > 0 control that the learned z p and z c have a reasonable scale, both of which are selected on training/validation data.
D D EFINITIONS
For convenience, we restates some definitions given in Lu et al.
Definition 1.
A structural equation model (SEM) C :=
( S , N ) governing the random vector X =
( X 1 , . . .
, X d ) is a set of structural equations:
[Equation]
where Pa ( X i ) ⊆
{ X 1 , . . .
, X d } \ { X i } are called the parents of X i , and the N i are independent noise random variables.
We say that “ X i causes X j ” if X i ∈
Pa ( X j ) .
We call causal graph of X to the graph obtained by drawing i) one for ii if
Definition 2.
Consider a SEM C :
= ( S , N ) .
An intervention e on C consists of replacing one or several of its structural equations to obtain an intervened SEM C e :=
( S e , N e ) , with structural equations:
[Equation]
The variable X e is intervened if S i ̸
= S ei or N i ̸ = N ei .
Definition 3.
Consider a structural equation model (SEM) S governing the random vector ( X 1 , . . .
, X n , Y ) , and the learning goal of predicting Y from X .
Then, the set of all environments E all ( S ) indexes all the interventional distributions P ( X e , Y e ) obtainable by valid interventions e .
An intervention e ∈ E all ( S ) is valid as long as (i) the causal graph remains acyclic,
(ii) E
[ Y e | Pa ( Y )] =
E [ Y | Pa ( Y )] , and (iii) V [ Y e | Pa ( Y )] remains within a finite range.
Definition 4.
( Exponential family )
A multivariate exponential family is a set of distributions whose probability density function can be written as
[Equation]
where Q : X → R is the base measure, Z ( θ ) is the normalizing constant, T : X → R k is the sufficient statistics, and θ ∈ R k is the natural parameter.
The size k ≥ n is the dimension of the sufficient statistics T and depends on the latent space dimension n .
Note that k is fixed given n .
Definition 5.
( Strongly exponential distributions )
A multivariate exponential family distribution
[Equation]
is strongly exponential, if
[Equation]
where l is the Lebesgue measure.
The density of a strongly exponential distribution has almost surely the exponential component and can only be reduced to the base measure on a set of measure zero.
Note that all common multivariate exponential family distributions (e.g. multivariate Gaussian) are strongly exponential.
Definition 6.
Let Θ be the domain of the parameters θ = { f , T , λ } .
Let ∼ be an equivalence relation on Θ .
A deep generative model is said to be ∼ –identifiable if
[Equation]
The elements in the quotient space Θ \ ∼ are called the identifiability classes.
Definition 7.
Let ∼ A be an equivalence relation on Θ defined by:
[Equation]
where A ∈ R k × k is an invertible matrix and c ∈ R k is a vector.
Definition 8.
Let ∼ P be an equivalence relation on Θ defined by:
[Equation]
where P ∈ R k × k is a block permutation matrix and c ∈ R k is a vector.
E A SSUMPTIONS
Assumption 11.
p T , λ ( s t | a t , c c − 1 , e ) belongs to a general exponential family with parameter vector given by an arbitrary function λ ( a t , c
c − 1 , e ) and sufficient statistics T ( s t ) =
[ T f ( s t ) T , T NN ( s t ) T ] T given by the concatenation of a)
the sufficient statistics T f
( s t ) =
[ T 1 ( s 1 t ) T , · · · , T n ( s nt ) T ] T of a factorized exponential family, where all the T i ( s it ) have dimension larger or equal to 2, and b) the output T NN ( s t ) of a neural network with ReLU activations.
The resulting density function is thus given by
[Equation]
where Q is the base measure and Z the normalizing constant.
Assumption 12.
p T , λ ( s t | r t +1 , e ) belongs to a general exponential family with parameter vector given by an arbitrary function λ ( r t +1 , e ) and
sufficient statistics T ( s t ) =
[ T f ( s t ) T , T NN ( s t ) T ] T given by the concatenation of a)
the sufficient statistics T f
( s t ) =
[ T 1 ( s 1 t ) T , · · · , T n ( s nt ) T ] T of a factorized exponential family, where all the T i ( s it ) have dimension larger or equal to 2, and b) the output T NN ( s t ) of a neural network with ReLU activations.
The resulting density function is thus given by
[Equation]
where Q is the base measure and Z the normalizing constant.
F G ENERALIZATION IN I MITATION L EARNING F. T HE C ONFOUNDER C t −
Note that, as pointed out by de Haan et al.
(2019), the state-action pair c t − 1 =
( s t − 1 , a t − 1 ) in the previous timestep serves as a confounder influencing each state variable in s t .
Hence, we can explicitly consider it in the causal diagram, as shown in Fig.
In practice, we use the observation-action pair ( x t − 1 , a t − 1 ) to approximate c t − 1 (Assumption 1) due to the unobservability of s t − 1 .
Thus, c t − 1 can be thought of as an observed confounder.
In NF-iVAE, such an observed c t − 1 can be treated as an additional auxiliary variable other than a t and e , which provides more information leading to identifiability.
That is, in some scenarios where the additional information of ( a t , e ) is not enough to identify s t , we can include c t − 1 to help further identify s t with theoretical guarantees.
This can be done by directly substituting ( a t , e ) with ( c t − 1 , a t , e ) in the previous assumptions, equations, and theorems (Appendix F.4 & H).
This technique also applies to the different generalization problems in RL that we will describe below.
For simplicity, we will no longer include the node c t − 1 in the subsequent causal diagrams.
F. E XPLAINING THE P RACTICALITY OF A SSUMPTION
Let us explain in more detail how practical Assumption 2 is in the IL setting.
Assumption 2a is a common assumption in causal discovery (Spirtes et al., 2000; Pearl, 2009; Peters et al., 2017).
The reason that we require it is that we need to leverage causal discovery algorithms to discover the causes of expert action a t among s t in Phase 2.
It also makes sense in Assumption 2b
that the generative/causal mechanism p ( x t | s t ) is assumed to be invariant across E all .
First, such a causal mechanism mapping latent variables to observations is widely adopted in the machine learning literature (Thrun et al., 2005; Hyvarinen & Morioka, 2016; Suter et al., 2019;
Hyv¨arinen et al., 2019;
Teshima et al., 2020; Khemakhem et al., 2020;
Sch ¨ olkopf et al., 2021;
Lu et al., 2022; Sun et al., 2021).
Second, we can view the generative process p ( x t | s t ) as a physical mechanism generating the observed variables, which is reasonably stable across E all (Peters et al., 2017; Sch ¨ olkopf, 2019;
Ahmed et al., 2020; Sch ¨ olkopf et al., 2021).
Assumption 2c is a widely-used default assumption in
OOD generalization (Peters et al., 2016; Arjovsky et al., 2019) and is also what we aim for in the
IL setting.
Apparently, Assumption 2, together with its corresponding causal diagram in Fig. 1, is flexible enough to cover most scenarios in the IL setting.
F. P OLICY G ENERALIZATION
Now we are ready to address the policy generalization problem in IL, by following the three phases in iCaRL described in Appendix C.1.
For ease of reference, this approach is called iCaRL-IL .
F.3. P HASE 1: I DENTIFYING S TATES
Under Assumption 2b, it is straightforward to obtain a corresponding generative model by directly substituting ( x , z , y ) with ( x t , s t , a t ) in Eqs.
[Equation]
The corresponding evidence lower bound (ELBO) is
[Equation]
) where q ϕ ( s t | x t , a t , e ) is an approximate conditional distribution for s t given by a recognition network with parameters ϕ .
To guarantee the identifiability result, the prior p T , λ ( s t | a t , e ) is assumed to satisfy Assumption 3 (i.e. Eq.
Considering the unknown normalization constant Z in this prior, we use score matching (Hyv¨arinen, 2005; Vincent, 2011) to learn ( T , λ ) by minimizing
[Equation]
In practice, as described in Lu et al.
, we can jointly learn ( θ , ϕ ) by combining Eq.
(18) and
Eq.
(19) in the following objective:
[Equation]
where ˆ f , ˆ T , ˆ λ , ˆ ϕ are copies of f , T , λ , ϕ that are treated as constants and whose gradient is not calculated during learning.
Importantly, under Assumptions 2b&3, we can follow the above steps to identify the latent variables s t in NF-iVAE with similar theoretical guarantees on identifiability, by directly substituting ( x , z , y ) with ( x t , s t , a t ) in all the corresponding equations and theorems.
All these are given in Appendix H.
F.3. P HASE 2: D ISCOVERING D IRECT C AUSES
After estimating s t for x t , we follow Lu et al.
and use a two-step approach to discover the direct causes s pt of the action a t .
Roughly speaking, we first conduct the PC algorithm (Spirtes et al., 2000) to learn a Markov equivalence class of DAGs, which outputs the direct neighbours of a t , denoted by NE ( a t ) .
Then, all pairs in NE ( a t ) are tested with conditional independence testing (Zhang et al., 2012) to discover z p by comparing their dependences after additionally conditioning on a t , cf.
Lu et al.
(2022) and Appendix C.1.
It is worth noting that whether c t − 1 is included or not, this approach works, because we concern only about the direct causes of a t .
F.3. P HASE 3: L EARNING AN I NVARIANT P OLICY
After having obtained s pt for a t across E tr , under Assumption 2c, the invariant predictor w can be learned by solving the following optimization problem: [Equation]
[Equation]
) where ℓ ( · ) could be any loss.
Since p ( a t | s p ) is assumed to be invariant across E all , the learned w is guaranteed to perform well across E all .
See the theorems in Appendix H.
When in an unseen testing environment, s pt is required to be inferred from x t before applying w for prediction.
This is done by leveraging the invariant p f ( x t | s t )
(Assumption 2b) through solving the following optimization problem:
[Equation]
(21) where the parameters f learned from E tr are fixed, and the hyperparameters λ 1 , λ 2 > 0 control that the learned s pt and s ct have a reasonable scale, both of which are selected on training/validation data.
Note that, here Φ can be viewed as a representation function that gives s pt from x t .
F. M ORE ON THE C ONFOUNDER C t −
Following the above, it is straightforward to obtain a corresponding generative model:
[Equation] [Equation]
The corresponding evidence lower bound (ELBO) is
L ELBOphase1 ( θ ,
[Equation]
(24) where q ϕ ( s t | x t , a t , c t − 1 , e ) is an approximate conditional distribution for s t given by a recognition network with parameters ϕ .
To guarantee the identifiability result, the prior p T , λ ( s t | a t , c t − 1 , e ) is assumed to satisfy Assumption 3 (i.e. Eq.
Considering the unknown normalization constant Z in this prior, we use score matching (Hyv¨arinen, 2005; Vincent, 2011) to learn ( T , λ ) by minimizing
L SMphase1 ( T , λ ) :=
E D E tr (cid:2)
E q ϕ ( s t | x t , a t , c t − 1 , e ) (cid:2) ||∇ s t log q ϕ ( s t | x t , a t , c t − 1 , e ) − ∇ s t log p T , λ ( s t | a t , c t − 1 , e ) || 2 (cid:3) (cid:3) .
In practice, as described in Lu et al.
, we can jointly learn ( θ , ϕ ) by combining Eq.
(24) and
Eq.
(25) in the following objective:
[Equation]
where ˆ f , ˆ T , ˆ λ , ˆ ϕ are copies of f , T , λ , ϕ that are treated as constants and whose gradient is not calculated during learning.
G G ENERALIZATION IN R EINFORCEMENT L EARNING G. M ORE E XPLANATIONS ON A SSUMPTION
The practicality of Assumption 4 can be detailedly explained in a similar way to that of Assumption 2.
Importantly, it is worth noting here that we do not include that r t +1 ⊥⊥ e | s pt (Fig. 4 in green, dashed line).
In other words, p ( r t +1 | s pt , a t ) might change across E all in various ways.
Technically, in an environment e ∈ E all , the causal module p e ( r t +1 | s pt , a t ) . = p ( r t +
[Equation] | s pt , a t , e = e ) can be represented by the following structural causal model (SCM) (Pearl, 2009):
[Equation]
where ϵ is a disturbance term and has a non-zero variance (i.e., the model is not deterministic), and θ denotes the effective parameters in the model/mechanism f .
That they all, except a t , have the superscript e explicitly indicates that all of them could be affected by e (i.e., they might vary across
E all ).
Any change on these three terms on the RHS of Eq.
(26) in the new environment e ∗ ∈ E all will produce a different distribution p e ∗ ( r t +1 | s pt , a t ) , i.e., a different reward function r e ∗ t +1 .
The different ways the reward changes lead to the three types of generalization problems we study in the RL setting, which will be respectively discussed in the subsequent sections.
G. I DENTIFYING L ATENT V ARIABLES IN THE RL S ETTING
Under Assumption 4b, it is straightforward to obtain a corresponding generative model by directly substituting ( x , z , y ) with ( x t , s t , r t +1 ) in Eqs.
[Equation]
The corresponding evidence lower bound (ELBO) is
[Equation]
where q ϕ ( s t | x t , r t +1 , e ) is an approximate conditional distribution for s t given by a recognition network with parameters ϕ .
To guarantee the identifiability result, the prior p T , λ ( s t | r t +1 , e ) is assumed to satisfy Assumption 3 (i.e. Eq.
Considering the unknown normalization constant Z in this prior, we use score matching
(Hyv ¨ arinen, 2005; Vincent, 2011) to learn ( T , λ ) by minimizing
L SMphase
[Equation] [Equation]
In practice, as described in Lu et al.
, we can jointly learn ( θ , ϕ ) by combining Eq.
(29) and
Eq.
(30) in the following objective:
[Equation]
where ˆ f , ˆ T , ˆ λ , ˆ ϕ are copies of f , T , λ , ϕ that are treated as constants and whose gradient is not calculated during learning.
G. P OLITY G ENERALIZATION
Note that, unlike the IL setting where we learn the invariant policy in the supervised way, in the RL setting we follow Zhang et al.
(2020b)
and combine our MSR s
At with the soft actor-critic (SAC) algorithm (Haarnoja et al., 2018) to devise a practical RL method, termed iCaRL-RL-P , as shown in Algorithm 1 of the appendix.
In principle, our MSR could be combined with any RL algorithm, including the model-free DQN (Mnih et al., 2015) or model-based PETS (Chua et al., 2018).
When deploying the learned invariant policy in an unseen testing environment, we first infer s
At for x t by leveraging the assumed invariant p f ( x t | s t ) (Assumption 4b).
This can be done by means of the representation function Φ that is obtained by directly substituting ( s pt , s ct ) with ( s At , s ¬ A t ) in Eq.
(21), where s ¬ A t are the remaining factors in s t that are not the causal ancestors of the reward.
H T HEOREMS AND P ROOFS
Unless stated otherwise, it is straightforward that all the theorems presented in this section can be easily proved using the same technique in Lu et al.
To avoid redundancy, we refer readers to
Lu et al.
for the proofs.
H. G ENERALIZATION IN IL
Theorem 1.
Assume that we observe data sampled from a generative model defined according to
Eqs.
(16, 17, 15), with parameters θ :=
( f , T , λ ) , where p T , λ ( s t | a t , e ) satisfies Assumption 3.
Furthermore, assume the following holds: (i) The set { x t ∈ X| φ ϵ ( x t ) = 0 } has measure zero, where φ ϵ is the characteristic function of the density p ϵ defined in Eq.
(ii) Function f in Eq.
(17) is injective, and has all second-order cross derivatives.
(iii)
The sufficient statistics in T f are all twice differentiable.
(iv) There exist k + 1 distinct points ( a t , e ) 0 , . . .
, ( a t , e ) k such that the matrix
L = (cid:0) λ (( a t , e ) 1 ) − λ (( a t , e ) 0 ) , . . .
, λ (( a t , e ) k ) − λ (( a t , e ) 0 ) (cid:1) of size k × k is invertible, where k is the dimension of T .
Then the parameters θ are identifiable up to a permutation and a “simple transformation” of the latent variables s , defined as a componentwise nonlinearity making each recovered T
i ( s it ) in T f ( s t ) equal to the original up to a linear operation.
Theorem 2.
Assume that the following holds: (i) The family of distributions q ϕ
( s t | x t , a t , e ) contains p θ ( s t | x t , a t , e ) , and q ϕ ( s t | x t , a t , e ) > 0 everywhere.
(ii)
We maximize L phase1 ( θ , ϕ ) with respect to both θ and ϕ .
Then in the limit of infinite data, we learn the true parameters θ ∗ up to a permutation and simple transformation of the latent variables s t .
Theorem 3.
Assume the hypotheses of Theorem 1 and Theorem 2 hold, then in the limit of infinite data, we identify the true latent variables s ∗ t up to a permutation and simple transformation.
Proposition 4.
Under Assumption 2 and the assumptions of Theorems 1 and 2, the imitation policy learned by iCaRL-IL across E tr in the limit of infinite data has optimal OOD generalization across
E all .
H.1. O N THE C ONFOUNDER C t −
Theorem 4.
Assume that we observe data sampled from a generative model defined according to
Eqs.
(16, 17, 15), with parameters θ :=
( f , T , λ ) , where p T , λ ( s t | a t , c t − 1 , e ) satisfies Assumption 3.
Furthermore, assume the following holds: (i) The set { x t ∈ X| φ ϵ ( x t ) = 0 } has measure zero, where φ ϵ is the characteristic function of the density p ϵ defined in Eq.
(ii) Function f in Eq.
(17) is injective, and has all second-order cross derivatives.
(iii)
The sufficient statistics in T f are all twice differentiable.
(iv) There exist k + 1 distinct points ( a t ,
c t − 1 , e ) 0 , . .
, ( a t , c t − 1 , e ) k such that the matrix L = (cid:0) λ (( a t , c
t − 1 , e ) 1 ) − λ
(( a t , c t − 1 , e ) 0 ) , . . .
, λ (( a t , c t − 1 , e ) k ) − λ (( a t , c
t − 1 , e ) 0 ) (cid:1) of size k × k is invertible, where k is the dimension of T .
Then the parameters θ are identifiable up to a permutation and a “simple transformation” of the latent variables s , defined as a componentwise nonlinearity making each recovered T
i ( s it ) in T f ( s t ) equal to the original up to a linear operation.
Theorem 5.
Assume that the following holds: (i) The family of distributions q ϕ
( s t | x t , a t , c t − 1 , e ) contains p θ ( s t | x t , a t ,
c t − 1 , e ) , and q ϕ ( s t | x t , a t , c t − 1 , e ) > 0 everywhere.
(ii)
We maximize L phase1 ( θ , ϕ ) with respect to both θ and ϕ .
Then in the limit of infinite data, we learn the true parameters θ ∗ up to a permutation and simple transformation of the latent variables s t .
Theorem 6.
Assume the hypotheses of Theorem 4 and Theorem 5 hold, then in the limit of infinite data, we identify the true latent variables s ∗ t up to a permutation and simple transformation.
H. G ENERALIZATION IN RL
Theorem 7.
Assume that we observe data sampled from a generative model defined according to
Eqs.
(27, 28, 15), with parameters
θ :=
( f , T , λ ) , where p T , λ ( s t | r t +1 , e ) satisfies Assumption 3.
Furthermore, assume the following holds: (i) The set { x t ∈ X| φ ϵ ( x t ) = 0 } has measure zero, where φ ϵ is the characteristic function of the density p ϵ defined in Eq.
(ii) Function f in Eq.
(28) is injective, and has all second-order cross derivatives.
(iii)
The sufficient statistics in T f are all twice differentiable.
(iv) There exist k + 1 distinct points ( r t +1 , e ) 0 , . . .
, ( r t +1 , e ) k such that the matrix
L = (cid:0) λ (( r t +1 , e ) 1 ) − λ (( r t +1 , e ) 0 ) , . . .
, λ (( r t +1 , e ) k ) − λ (( r t +1 , e ) 0 ) (cid:1) of size k × k is invertible, where k is the dimension of T .
Then the parameters θ are identifiable up to a permutation and a “simple transformation” of the latent variables s , defined as a componentwise nonlinearity making each recovered T
i ( s it ) in T f ( s t ) equal to the original up to a linear operation.
Theorem 8.
Assume that the following holds: (i) The family of distributions q ϕ
( s t | x t , r t +1 , e ) contains p θ ( s t | x t , r t +1 , e ) , and q ϕ ( s t | x t , r t +1 , e ) > 0 everywhere.
(ii)
We maximize L phase1 ( θ , ϕ ) with respect to both θ and ϕ .
Then in the limit of infinite data, we learn the true parameters θ ∗ up to a permutation and simple transformation of the latent variables s t .
Theorem 9.
Assume the hypotheses of Theorem 1 and Theorem 2 hold, then in the limit of infinite data, we identify the true latent variables s ∗ t up to a permutation and simple transformation.
H.2. P OLICY G ENERALIZATION IN RL
Proposition 5.
Under Assumptions 4,12&5 and the assumptions of Theorems 7&8, the policy learned across E tr in the limit of infinite data has optimal OOD generalization across E all .
Proof.
This proof is straightforward.
Under Assumptions 4&12 and the assumptions of Theorems 7&8, Theorems 9 guarantees that the latent variables s can be identified across E tr .
Then, by leveraging the method in Phase 2, we can discover the minimum sufficient states s A .
Since Assumptionindicates that all the environments share the same reward function and transition dynamics on top of s A , the learned policy across E tr is guaranteed to generalize to E all .
H.2.
R EPRESENTATION G ENERALIZATION IN RL
Proposition 6.
Under Assumptions 4,12&6 and the assumptions of Theorems 7&8, the representation function Φ learned across E tr in the limit of infinite data is able to generalize to E te .
Proof.
Under Assumptions 4&12 and the assumptions of Theorems 7&8, it is theoretically guaranteed that the latent variables s can be identified across E tr and the representation function Φ can be also identified.
Then, by leveraging the method in Phase 2, we can discover the minimum sufficient states s A .
Also, under Assumption 4b, Φ is invariant across E all .
Therefore, it is guaranteed that s A can be inferred using Φ across E all .
Under Assumption 6, we also know that all the environments share the same set of s A , which is the only information required for policy learning across E te .
Hence, the learned representation function Φ is guaranteed to generalize to E te .
H.2. D YNAMICS G ENERALIZATION IN RL
Proposition 7.
Under Assumptions 4,12,7&8 and the assumptions of Theorems 7&8, the local dynamics model learned across E tr in the limit of infinite data has optimal OOD generalization across E all .
Proof.
This proof is also straightforward.
Under Assumptions 4&12 and the assumptions of Theorems 7&8, it is theoretically guaranteed that the latent variables s can be identified across E tr .
Then, by leveraging the method in Phase 2, we can discover the direct causes s P i for each state s i under
Assumption 7.
Because Assumption 8 indicates that there exists some (local) dynamics model p ( s it +
| s P i t , a t ) that is invariant across E all .
This shows that the learned local dynamics model has optimal generalization across E all .
I E XPERIMENTAL R ESULTS I. I DENTIFIABILITY A NALYSIS ON S YNTHETIC D ATA
In this section, we empirically verify the identifiability of states in both IL and RL setting.
For this, we conduct a series of experiments on synthetic data generated according to a family of MDPs, as shown in Fig.
3 of the appendix.
Details of the ground truth data generating process are given in Appendix J.1.
We train on two environments with different soft interventions on states.
To measure identifiability, we compute the mean correlation coefficient (MCC) used in Khemakhem et al.
(2020), which can be obtained by calculating the correlation coefficient between all pairs of true and recovered latent factors and then solving a linear sum assignment problem by assigning each
VA E β - V A E
i VA E N F - i
VA E - I L
N F - i V A E - I
L *
M CC IL Setting (a) MCC in IL
VA E β - V A E
i VA E N F - i
VA E - R L
N F - i V A ER L *
M CC
RL Setting (b) MCC in RL
VA E β - V A E
i VA E N F - i VA E - R L
N F - i V A ER L * M CC Instantaneous Effect (c) MCC in RL   20 Number of Trajectories 0.2 0. P e
r
f
o r m a n c e ( S c a l e d )
VDICEEDMBCRCALICILiCaRL -IL (d) BeamRider recovered latent factor to the true latent factor with which it best correlates.
By definition, higher
MCC scores indicate stronger identifiability.
In the IL setting, Fig.
7a shows that the states are better identified by NF-iVAE-I/RL than by those commonly used VAE (Kingma & Welling, 2013) and β -VAE (Higgins et al., 2016).
It is worth noting that when additionally conditioning on the previous observation-action pair c t − 1 , the result, denoted by NF-iVAE-IL ∗ , empirically confirms that it helps further identify the states.
We can obtain the similar conclusion in the RL setting, as shown in Fig.
7b.
Even when there exist instantaneous effects between state variables as mentioned in Section 2.1.2, we can see from Fig.
7c that NF-iVAE-RL works as well.
I. G ENERALIZATION IN I MITATION L EARNING
In the IL setting, we follow Bica et al.
and evaluate generalization of our approach on the following two control tasks from OpenAI gym
(Brockman et al., 2016):
LunarLander (Brockman et al., 2016) and BeamRider (Bellemare et al., 2013) (Fig.
7d in the appendix).
For each task, we also use pre-trained RL agents from RL Baselines Zoo (Raffin, 2018) and
Stable OpenAI Baselines (Hill et al., 2018) to obtain expert policies, and then follow an approach similar to the one in Zhang et al.
(2020a) to obtain datasets with demonstrations from the expert in two different environments.
We follow the same setting of Bica et al.
(2021) to add spurious correlations.
We train on demonstrations from two training environments and test on an unseen testing environment.
Further details can be found in Appendix J.2.
Similar to Bica et al.
(2021), we also compare our iCaRL-IL with 1)
Behaviour Cloning (BC) (Pomerleau, 1991), 2) RCAL (Piot et al., 2014), 3) ValueDICE (VDICE)
(Kostrikov et al., 2019), 4) Energy-based Distribution Matching (EDM) (Jarrett et al., 2020), 5)
IRM (Arjovsky et al., 2019), and 6) ICIL
(Bica et al., 2021).
As shown in Fig.
2a, our iCaRL-IL consistently outperforms the benchmarks and can generalize better to the unseen testing environment.
I. G ENERALIZATION IN R EINFORCEMENT L EARNING I.3. P OLICY G ENERALIZATION
In this section, we report experimental results on cartpole swingup from the DeepMind Control (DMC) suite (Tassa et al., 2018) in two settings (i.e., simple distractors and natural video distractors , see Appendix J.3).
we compare our iCaRL-RL-P against several baselines: 1) Stochastic Latent
Actor-Critic (SLAC)
(Lee et al., 2020), 2) DeepMDP (Gelada et al., 2019), 3) Model-Irrelevance
State Abstractions (MISA) (Zhang et al., 2020a), 4) Deep Bisimulation for Control (DBC) (Zhang et al., 2020b), and 5) Invariant Policy Optimization (IPO) (Sonar et al., 2021).
We first train policies on two cartpole swingup environments with different simple distractors and then evaluate them on the average return obtained by deploying the learned policies on another cartpole swingup environment with natural video distractors .
Fig.
2b empirically verifies that the police learned by our method are able to generalize to the unseen testing environment.
I.3. R EPRESENTATION G ENERALIZATION
We also evaluate the generalization performance of the learned representation function Φ on the DMC suite.
Following Zhang et al.
(2020b), this is done by training SAC with new reward functions on walker run using the fixed representation Φ learned from two walker walk environments with different simple distractors .
Note that, here these new reward functions satisfy Assumption 6.
As shown in Fig.
2c, our method, termed iCaRL-RL-R, learn a representation function which has better generalization power.
I.3. D YNAMICS G ENERALIZATION
We now test the generalization capacities of the learned dynamics on cheetah run from the DMC suite.
We first learn dynamics models on two cheetah run environments with different simple distractors and then follow Zhang et al.
(2020a) to evaluate them in terms of model errors obtained by deploying the learned dynamics models on another cheetah run environment with natural video distractors .
As shown in Fig.
2d, our method is more stable than the others in that the model error is consistently small.
J M ORE D ETAILS ABOUT E XPERIMENTS J. S YNTHETIC D ATA
The synthetic data are generated according to a family of MDPs, as shown in Fig.
Details of the ground truth data generating process are as follows:
[Equation]
where U{·} denotes the discrete uniform distribution, N ( · ) the Gaussian distribution [Equation] and g ( · ) is given by a neural network with 3-dimensional input and 10-dimensional output, whose parameters are randomly set in advance.
We create a dataset consisting of samples from two training environments by drawing 2000 steps from each of the two environments E = { 0 . 2 , 2 } .
Note that, in the experiments with instantaneous effects between states, we include an additional relationship between s 2 t and s 1 t by replacing Eq.
(37) with s
[Equation]
[Equation]
In all the experiments on synthetic data we set the number of the latent variables to n = 3 .
The architecture we used is as follows.
NF-iVAE-I/RL λ f -Nonlinear Prior • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size =NF-iVAE-I/RL Encoder • Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU • Mean Output layer: Fully connected layer, output size =• Log Variance Output layer: Fully connected layer, output size =
NF-iVAE-I/RL Decoder • Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU
• Mean Output layer: Fully connected layer, output size = output dimension • Variance Output layer: 0 .
01 × 1 , where 1 is a vector full of 1 with the length of output dimension
NF-iVAE-I/RL Predictor w
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 6, activation = ReLU • Output layer: Fully connected layer, output size =
J. G ENERALIZATION IN I MITATION L EARNING
Unless stated otherwise, we used the exact experimental setting and implementation details that are described in Bica et al.
(2021) for the benchmarks.
For a fair comparison, we use the same architecture for the encoder in our NF-iVAE-IL.
We reverse the architecture of the encoder as a decoder.
We set the number of the latent variables to n = 50 .
We do the hyperparameter search by exactly following the guides given in Bica et al.
In addition, the architectures of the prior and the predictor are as below.
NF-iVAE-IL T NN -Prior
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 50,
activation = ReLU • Output layer: Fully connected layer, output size =
NF-iVAE-IL λ
NN
-Prior
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 50,
activation = ReLU • Output layer: Fully connected layer, output size =NF-iVAE-IL λ f -Prior • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 50,
activation = ReLU • Output layer: Fully connected layer, output size =
Predictor w
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 100,
activation = ReLU • Output layer: Fully connected layer, output size =
J. G ENERALIZATION IN R EINFORCEMENT L EARNING
For convenience, here we restate the description of the two settings: simple distractors and natural video distractors .
See more in Zhang et al.
(2020b).
Simple Distractors Setting.
We include simple background distractors, shown in Figure 3 of Zhang et al.
(2020b) (middle row), with easy-to-predict motions.
We use a fixed number of colored circles that obey the dynamics of an ideal gas (no attraction or repulsion between objects) with no collisions.
1: while forever do 2: for e ∈ E tr do 3: a t ← π ( x et ) 4: x et +1 , r t +1 ← step ( x et , a t ) 5: store ( x et , a t , r t +1 , x et +1 ) 6: end for 7: for e ∈ E tr do 8:
Sample batch from replay buffer 9:
Estimate the representation function s At = Φ( x t ) according to Section 3.
10: end for 11:
Train policy according to a modified SAC ▷ Algorithm12: end while
1: Get
2: Train
3: Train
4: Train
Natural Video Setting.
We incorporate natural video from the Kinetics dataset (Kay et al., 2017) as background (Zhang et al., 2018), shown in Figure 3 of Zhang et al.
(2020b) (bottom row).
Unless stated otherwise in the main text, we used the exact experimental setting and implementation details that are described in Bica et al.
(2021) and Zhang et al.
(2020a) for the benchmarks.
For a fair comparison, we use the same architecture for the encoder in our NF-iVAE-RL.
We reverse the architecture of the encoder as a decoder.
We set the number of the latent variables to n = 50 .
We do the hyperparameter search by exactly following the guides given in Bica et al.
In addition, the architectures of the prior and the predictor are as below.
NF-iVAE-RL T NN -Prior
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 50,
activation = ReLU • Output layer: Fully connected layer, output size =NF-iVAE-RL λ NN -Prior • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 50, activation = ReLU • Output layer: Fully connected layer, output size =NF-iVAE-RL λ f -Prior • Input layer: Input batch (batch size, input dimension) • Layer 1: Fully connected layer, output size = 50,
activation = ReLU • Output layer: Fully connected layer, output size =
Predictor w
• Input layer: Input batch (batch size, input dimension)
• Layer 1: Fully connected layer, output size = 100,
activation = ReLU • Output layer: Fully connected layer, output size =