Model-Invariant State Abstractions for Model-Based Reinforcement Learning
Authors: Manan Tomar
Amy Zhang
Roberto Calandra
Matthew E. Taylor
Joelle Pineau
Venue: ICLR 2022 Submitted (2022)
Keywords: Reinforcement Learning
Model-based RL
State Abstractions
Generalization in RL

Abstract: Accuracy and generalization of dynamics models is key to the success of model-based reinforcement learning (MBRL). As the complexity of tasks increases, learning accurate dynamics models becomes increasingly sample inefficient. However, many complex tasks also exhibit sparsity in dynamics, i.e., actions have only a local effect on the system dynamics. In this paper, we exploit this property with a causal invariance perspective in the single-task setting, introducing a new type of state abstraction called \textit{model-invariance}. Unlike previous forms of state abstractions, a model-invariance state abstraction  leverages causal sparsity over state variables. This allows for compositional generalization to unseen states, something that non-factored forms of state abstractions cannot do. We prove that an optimal policy can be learned over this model-invariance state abstraction and show improved generalization in a simple toy domain. Next, we propose a practical method to approximately learn a model-invariant representation for complex domains and validate our approach by showing improved modelling performance over standard maximum likelihood approaches on challenging tasks, such as the MuJoCo-based Humanoid. Finally, within the MBRL setting we show strong performance gains with respect to sample efficiency across a host of continuous control tasks. 

MODEL-INVARIANT STATE MODEL-BASED REINFORCEMENT LEARNING
ABSTRACTIONS FOR
Anonymous authors Paper under double-bling review
ABSTRACT
Accuracy and generalization of dynamics models key to the success of model- based reinforcement learning (MBRL). As the complexity of tasks increases, learn- ing accurate dynamics models becomes increasingly sample inefficient. However, many complex tasks also exhibit sparsity in dynamics, i.e., actions only have local effect on the system dynamics. In this paper, we exploit this property with causal invariance perspective in the single-task setting, introducing new type of state abstraction called model-invariance. Unlike previous forms of state abstractions, model-invariant state abstraction leverages causal sparsity over state variables. This allows for compositional generalization to unseen states, something that non- factored forms of state abstractions cannot do. We prove that optimal policy can be learned over exact generalization in a simple toy domain. Next, we propose practical method approximately learn a model-invariant representation for complex domains and validate our approach by showing improved modelling performance over standard maximum likelihood approaches Humanoid. Finally, within the MBRL setting we show strong performance gains with respect to sample efficiency across host of continuous control tasks.
state abstraction and show improved
challenging tasks, such as the MuloCo- based
INTRODUCTION
Model-based reinforcement learning MBRL [4, 15] is a popular framework for data-efficient learning of control policies. At the core of MBRL is learning an environmental dynamics model and using it to: 1) fully plan [14, 11], 2) augment the data used by model-free solver 54], 3) use as an auxiliary task while training [35, 58). However, learning a dynamics model - similar to suffers from the issue of generalization, since the data we train on is not necessarily the data we test on. This is persistent issue that is worsened in MBRL as even a small inaccuracy in the dynamics model or changes in the control policy can result in visiting completely unexplored parts of the state space [1]. This advocates for the need to learn models capable of generalizing well. Various workarounds for this issue have been explored in the past; for example, combining local but simple models with global, more expressive models [36, 22], using an ensemble of models to handle uncertainty in estimates [11,33] and coupling the model and policy learning processes [34] so that the model is always accurate to certain threshold. However, these approaches do not stem from a representation learning viewpoint and thus fail to leverage special structure in dynamics for better generalization.
other supervised learning problems
This paper studies how to improve generalization capabilities through careful state abstraction. In particular, leverage two existing concepts to define new kind of state abstraction. The first concept is that many real world problems exhibit sparsity in the local dynamics - given a set state variables, each variable timestep f. only depends in the previous timestep (see Figure 1). The second concept is the principle of causal invariance, which dictates that given a set of candidate features, we should aim build representations that comprise only those features that are consistently necessary for predicting the target variable of interest across different interventions [46]. In the MBRL context, we can cast model learning with causal invariance as supervised objective where the target variables are the next state variables and input features are the current state and action variables (the probable set of causal predictors of the target).
small subset (i.e., local) of all variables

---

Intuitively, since we learn predictor that only takes into ac- count invariant features that consistently predict the target vari- able well, it is likely to contain the true causal features, and therefore, will generalize well to all possible shifts in the data distribution. The two concepts of sparsity and causality intertwined in that they both are forms of inductive biases that surround the agent dynamies [23]. This paper shows how causal invariance tools can be effectively used to exploit sparsity in dy- namies, leading to improved model generalization. Given basic exploratory assumptions, we analyze this question theoretically and show empirically that we can learn a model that generalizes well on state distributions induced by policies distinct from the ones used while learning it. To do this, we introduce a new state abstraction, model-invariance, which leverages sparsity over state variables. We connect this abstraction viewpoint to a concrete problem of generalization in model-based RL., that of arising spurious correlations in the dynamics model, even when all state variables are task relevant. Having built enough intuition, we then introduce practical method that approximates learning model-ine invvriiatt representation for more complex domains that use function approximation (i.c., neural networks (NNs)). We empirically observe that model invariance leads to better model general- ization for domains such as the MuJoCo-based Humanoid. Our method is simple to implement and shows strong performance on multiple MuJoCo tasks, outperforming state of the art model-based methods in the low-data regime.
+2
x
Figure 1: Graphical model of spar- sity across state variables. The state variable 2241 (shaded in blue) only de- pends on two variables of and of (in the blue box). Predicting 2441 using the other variables (say xi+1) can re- sult in spurious correlations which lead to poor generalization.
2 PRELIMINARIES
RL Setup. We consider the agent's interaction with the environment as a discrete time 7-discounted Markov Decision Process (MDP) [47] M = (X,A, P, RG Mo). where X C IE is a finite but arbitrarily large state space and A is the and action space; P p= P( [[1l|x,, a) is the transition kernel; R = r(x, a) is the reward function with the maximum value of Rmaxi E [0, 1) is the discount factor; and Nn is the initial state distribution. Let क X - AA be a stationary memoryless policy, where A is the set of probability distributions on A. The value function of a policy π at state E X is defined as / (T(x)) = E o (I,, at) Fo = z, x). Similarly, the action-value function of T is defined as Q" (Ix a) = E| 7 (r( at) (mg = x, ag = a, T). The Bellman optimality operator T T X + R xxxxA is defined as TQ a) r(æ, a) + 7 (P(-)o; a), max, Q(: (())) In this work, we assume that all d state variables x' .2. ...., xd are useful for the task hand, i.e., we are given the full state. Furthermore, we assume that the transition dynamics over the full state are factorized. More formally:
Assumption 1. (Transition Factorization) For given full state vectors 22, T++1 E X, action a A. and x² denoting the 1tht dimension of state we have P P(z+41 (2,,0) III P ((
Note that this is a weaker assumption than factored MDPs [31, 24] as we do not assume correspond- ing factorization of the reward function.
Invariant Causal Prediction. Invariant causal prediction (ICP) [46] considers the problem of learning an invariant representation w.r.t. spurious correlations that arise due to noise in the underlying causal model (unknown) describing a given system. The key observation is that if one considers the direct causal parents of response/target variable of interest (Y"), then the conditional distribution of Y given these direct causes PYY does not change across interventions on any variable except Y. Therefore, ICP suggests collecting data into different environments (corresponding to different interventions), and to output the set of variables X; for which a learned predictor of Y remains the same given X, across the multiple environments with high probability.
State Abstractions. State abstractions allow us to map vehaviorally-equivaie states into single abstract state, thus simplifying the learning problem, which then makes use of the (potentially much smaller set of) abstract states instead of the original states [5]. In principle, any function approximation architecture can act as an abstraction, since it attempts to group simílar states together. Therefore, exploring the properties of representation learning scheme as state abstraction help

---

develop stronger intuition for building practical algorithms. In the next section, we build our theory based on this connection. A well known state abstraction is bisimulation [19, 48, 37]. Formally, abstraction : X' + S is bisimulation if for any two states 21, 22 and next state I E X, abstract state 8 E S, a E A where (((1) = ((22), we have R(æ1,a) = R(22, a), and Exee-"()) P(z[(z,,)) a) = E Ex) P(z|xz2 a) Since an exact equivalence is not practical, prior work deals with approximate variants through the notion c-closeness [29].
3 MODEL INVARIANCE AND ABSTRACTIONS
Current state abstractions such as bisimulation do not provide support for sparse structures. Instead, closeness in abstract states is defined based on probabilities of all state variables together, i.e., To remedy this, we introduce a new state abstraction, called model-invariance, which is specific to each state variable x². Formally, an invariant abstraction Φ; is one which has the same transition probability of the hhh next state variable x' for any two given states I and T2, i.e., P(=*I [ a) P( [r2, a). Note that if we assume factored rewards, can define corresponding reward-based invariant abstraction that parallels the bisimulation abstraction more closely, but we focus here on the reward- free setting. Since is impractical to ensure this equivalence exactly, we can use an approximate definition which ensures an e-closeness.
Definition 1. (Approximate Model Invariance) & is ee.p-modell-nvariane for each index i,
sup (0 |P('1001, a) P(x* [2, a) S C.P.
is en-model-invarinnt e X
R(m,, R(22, a)
Based on the definition, it intuitively seems like applying the abstraction over the original transition distribution should be close to the transition distribution over abstract states. This can be precisely written in the following lemma, which is necessary result in ensuring that an optimal policy can still be learnt over a model-invariant abstraction:
Lemma 1. (Model Error Bound) Let & be an G, a.p-appoximmate model-invariant abstraction on MDP M. Given any distributions Ps, : 8; E (x(2)) where ps, is supported on & (s;) and ps = -I[ Pso we define Mo = (o(X), Po, Rs. 7) where Po(s.a) = Ez-pa [P(-(x,a) and R = Expp.
[R(x,a)) Then for any x E X, R. A.
(P((s.a)@@P(C @p(xx a) S > 6,P,
where 4P denotes the lifted version of P. where we take the next-step transition distribution from observation space X' and lift it to latent space S (Proof in Appendix 2). Po refers to the transition probability of a MDP that acts the states ((X), rather than the original MDP. Finally, note that we are particularly concerned with the case where each x² is atomic in nature, i.e., it is not further divisible. Such property ensures that model-invariance does collapse to bisimulation".
CAUSAL INVARIANCE IN MODEL LEARNING
Having defined model-invarianl abstractions, we are now ready provide connections between causal invariance and model learning in RL. For simplicity, we can assume that there exists linear structural equation model [45] that consists of the d state variables and action a as the features X and the next state variable x/41 as the target Y. for each index i. Similar to the ICP setting(Section 2), we can define the different environments as follows:
Assumption 2. (ICP Environments) For each e E E: the experimental setting e arises due to one or several interventions on variables from (æ), ,,,,z, at) but not on 2tt here, we allow for do-interventions [44] or ioft-interventions [18].
In the simplest case index describes each state dimension, i.c.. being atomic. However in general the index could be grouping of different dimensions as well. Consider the case where there is only one index over which we build model invariance abstraction. Now x would then correspond to the entire state Since Definition I would be satisfied trivially, we do not gain anything the sparsity level, thus collapsing to bisimulation.

---

For purposes, each intervention corresponds change in the current state and action distribution. This change in distribution can be realized as a change in the policy. Therefore, each policy я defines an ICP environment It can be shown that under Assumption the direct causes, i.c., parents of -/-1- define a valid support over invariant predictors, namely 5 = ( 2. The key idea therefore is to make sure that in predicting each next state variable x 1 we use only its set of invariant predictors and not all state variables and actions (see Figure 1). With this intuition, becomes clearer why our original model learning problem is inherently tied with learning better representations. in that having access to a representation that discards excess information for each state variable (more formally, a causally invariant representation), would be more robust to spurious correlations and thus, at least in principle, lead to improved generalization performance across different parts of the state space. In fact, such a causally invariant representation obeys the properties of a abstraction. Formally,
Proposition! 1. For the abstraction (x)x) = [a]s,, where S, = PA di is model-invariant
Proof is provided in Appendix 2. It now becomes easy to see that sparsity in dynamics is central to what we have discussed so far, since if do not have sparsity, the causally invariant representation trivially reduces to the original state z, thus resulting in no state aggregation. However, when sparsity is present, following the above definition leads to a representation @ such that any two components Φ; and 6j are decorrelated, since each encodes information about a different subset of state variables. Finally, we show that learning a transition model over a model-invarian abstraction & and then planning over this model is optimal.
Theorem 1. (Value bound) If ó is an ER. C.P approximate model invariant abstraction on MDP M. and Ma is the abstract MDP formed using 4, then we cun bound the loss in the optimal state action value function in both the MDPs as:
vo
Qirylar - Qis 2,4
I 1-
T[Q
2,1
d
-
- TIQMa
12,0 SER + Σ EiP 2(1 ((1-))"
Rmax
where refers to the concentrability coefficient as defined in [9]. Here, an admissible distribution refers to any distribution that can be realized in the given MDP by following a policy for some timesteps, while z refers to the distribution from which the data is generated. Proof and all details surrounding the theoretical results are provided in Appendix 2. This result is important because we can follow this with standard sample complexity arguments that will have a logarithmic lo] dependence, thus guarantecing that learning in this abstract MDP is faster and only incurs the above described sub-optimality.
So far, we have embedded the RL model learning problem in the ICP framework. This highlights the connection between generalization issues in MBRL and arising spurious correlations in identifying the true causal predictors of each state variable. Furthermore, properties of the causally invariant representations which avoid spurious correlations. In the next section, we now show how ICP be used sub-routine in learning a more generalized model in the simple setting of a linear MDP.
have explored the state abstraction
5 LINEAR CASE: CERTAINTY EQUIVALENCE
In the simpler setting of tabular RI,, estimating the model using transition samples and then planning over the learned model is referred to certainty equivalence [6]. Particularly for estimating the transition model, it considers the case where we are provided with n transition samples per state-action pair, (Ie, ax) in the dataset Da,as and estimate the model as
1 Σ [
P(æ or)
SE RED
"The proof follows directly by applying Proposition 1 of Peters al. [46] (which itself follows from construction) each state variable indexed by dimension

---

Now, when next state components do not depend each other given the previous state and action (i.e., Assumption 1), we can re-write P(Z1+1 TT, ) as II. P( [Te, at). Furthermore, we know the causal parents of at+1, we can instead empirically estimate the true transition probabilities as ((2441022,00) = P( (A(4))1))
Σ I(== x(+1) nkk RED
A where D D= UDs.. I ZE E (3). From Proposition we know that such an invariant solution
exists, and is defined by the causal parents of each state variable. Therefore, in the linear dynamics case, given data from multiple environments (different policies), we can use ICP to learn the causal parents of each state variable and then estimate the probability of certain transition using Eq. 1.. We refer to this as the invariant model learner and detail the procedure in Algorithm 1. On the other hand, if we do extract the causal parent variables and instead all state variables in Eq. 1, we would get the standard maximum likelihood (MLE) learner.
Algorithm Linear
MBRI. 1: Input Replay buffer D containing data from multi- plc policies/envs, confidence parameter o;
The invariance based solution avoids spurious correlations more than the MLE learner, result- ing in better generalization. To see this, con- sider simple linear MDP with three state vari- ables (x. x² 2²),, each depending only on its ri), and taking integer val-
2: for state variable = 1, .,...; d do
3: S: [i, D, a::; (Appendix ??)
own (PA (A(441) ues between |-10, 10). The exact details of the MDP are described in Appendix 3.1. We con- sider three different distributions corresponding to three different policies, each describing an ICP environment. In Table 1, we compare our invariant learner with standard MLE learner and show how their error with the true probability of transition varies as the number samples grows.
4: Estimate P. from D using Eq. 1;
5: end for
6: Estimate transition probability kernel --II.. ...,, N do
7: for state j = 1,....
TR(8)) + Pln P,85)
9: end for
Note that Table represents the results for one environment, specified by fixed policy that is used for data collection. If the policy changes, would result in a different environment de- scribed in Section 3. Our ideal scenario is to find a predictive model that is closest to the true model for all environments. We find that the invariant learner quickly converges to approxi- mately the same solution across all training envi- ronments, with just a few data samples. On the other hand, the standard MLE learner results in different solutions for each training environment in the low data regime. The solution provided at test time such a case is an average all such solutions found during training, which is clearly off the true probability (higher error in Table 1).
Samples
MLE
Model-Invariance
100
1
9.3 1 1.3
200
24.6 + 9.23
5.7 1.5
500
14.5 + 2.75
2.7 + 0.7
2000
9.6 1:6:10 1.72
1.8 + 0.3
5000
6.4 = 1.46
1.6 0.3
Table 1: Linear MDP Transition Prediction Error. Consider the simple linear MDP from Appendix 3.1. We compare the estimated transition probability of our invariant learner with the MLE learner (lower is better). The invariant learner converges faster and more stably the true solution. Mean and std. err. over random seeds. Order of magnitude of the errors is 1c-3.
is worth noting that this example assumes linearity in dynamics, which allows us to use the ICP procedure from Peters et al. [46]. In the general non-linear case, this not possible. To that end, in the next section we will describe a practical method that leverages ideas from self-supervised learning exploit sparsity in an end-to-end manner.
NON-LINEAR CASE: LEARNING PRACTICAL MODEL-INVARIANT REPRESENTATIONS
We now introduce a practical algorithm for learning model-invariant representations where we relax the following assumptions: linearity in dynamics, having access to data from multiple environments,
Using Theorem 1, know that performing planning over this estimated dynamics model would be copptmall.

---

Softmax
St
Algorithm 2 Non-linear Model Invara-nn MBRL
p
1: Input Replay buffer D 8: Value and policy network parameters 0g. 0. model parameters Or, O for any MBRL algorithm:
2: for environment steps &
T do
3
3: Take action d, ++][2;), observe r. and 21+1. and add to the replay buffer D:
Zt
0(2,, a)
LMLE
KI xM
4: for Mmodel free updates do
5: Sample batch {Tj, a, "5,"1,55+11 1.W from D;
at
6:
Run gradient update for the model free com- ponents of the algorithm (e.g.. 0x, 02 etc.); end for
h
7:
8:
for Moodal updates do
9
9:
Sample batch {.z. a,, 22, from D;
Be+1
Variable
10:
Update reward model (8,):
11:
Update invariant dynamics model: 0, evrariant_update(( Vo, Ly) (Appendix 3.2 or Eq. 2);
Figure 2: Architecture for learning model-invariant representations. Model-invariang uses two (or more) transition dynamics models, denoted by f and h, over a common representation Φ. The critic 9 provides score for chosen dimension/state variable of the output of f and h models.
12:
end for
13:
end for
and being in the strictly batch setting. Having connected the abstraction viewpoint to the problem of spurious correlations arising in dynamics models, we wish come up with NN representations that abstract away irrelevant information on a vaate-vaaaaee level (i.e., x/., not Iy)) and learn them in an end-to-end manner. To that end, we view the task of learning such representations as a self-supervised objective where we want to be invariant to models that exhibit spurious correlations. Consider two (or more) randomly initialized and independently trained (on different samples) dynamics models attached to common representation c. We view these models as augmented versions of the true dynamics models since each model captures different spurious correlations. We must be invariant to augmentations in the model space per-state-variable level (since spurious correlations arise for individual state variable dynamics).
To instantiate this idea, take inspiration from recent method called ReLIC [40], which uses contrastive loss [41] over augmented views of the same data sample and an invariance loss to enforce consistency. ReLIC variables and enforcing invariance over certain functions (like data augmentations). Here we are exploiting the same connection, with additional modifications. Instead of using augmentations over inputs, we induce augmentations over models and then define the invariance loss for a particular dimension - the goal is to eliminate the contribution of spurious features when predicting a particular state variable. Note that doing exact causal interventions is not possible in the large state-action spaces in RL but we can still make strong connections to the causal literature and also use it in simpler settings, as we do in Section 5 (linear case) where we deploy exact invariance tools.
shown be closely connected to doing causal interventions over input
After randomizing two (or more) identical models at the start of training, model is sampled randomly and is used for minimizing the standard MLE model predictive loss at each optimization step. Simultaneously, an invariance loss defined over the predictions of both models augments the MLE loss. The invariance loss enforces consistency in the prediction of both (all) models by minimizing the difference in similarities between the prediction of one model w.r.f. the prediction of the other model and vice versa (Eq. 2). This similarity is computed for single (randomly selected) state variable at each optimization step, with the specifics being borrowed from the ReLIC objective Mitrovic al. [40]. Finally, further consistency is encouraged because all models have common representation with a a bottleneck structure. Our implementation
detailed in Figure and pseudocode in Appendix 3.2. Although this method is designed to be robust to spurious correlations in dynamics models, in our experiments we will show that the representations learnt for each index i, Φ actually end up being sparse, a property that motivated the model invariance abstraction back in Definition 1. The overall model invariance loss is:
C ExND (f( - + KI. (I((p
f softmax (o (W N)), (2)
MLE Loss
Invariance Loss

---

Hm3
Standard Invariant
Humar old 92
1.0
Test Model Error
Test Model Error
mean corr. 0z
0.8
w/o model-inv.
1.63 1 0.48
w/ model-inv,
1.25 + 0.53
0.6
385
0.0
2.5
5.0
7.5
10.0
Steps (x10*)
Per state Inv
Per dim Inv
Figure 4: Invariant Model Learning on Humanoid-v2. Left plot: test model learning error for different horizon values. Mean and std. err. over 10 random seeds. Middle plot: modelling error when we enforce per-state per-dim invariance. Table: mean correlation metrie; order of magnitude is 10".
where 4 (9 ar)), (('((x,, at)) is the cosine similarity between the predictions for the models f and h for the state variable indexed by The output the inner product is normalized for each sample and thus the resulting vector (z at)) is a probability distribution. In our implementation, the function g is fully connected neural network and is used to project the outputs of the models f and h. It is often called the critic network in self-supervise learning losses [10]. Note that the matrix v is not symmetric since the values any two symmetric indices are different (they are computed for different samples) the KL. loss remains well-defined.
We will use the invariant model learner described above within model-based RL algorithm and compare its policy performance to standard MLE based model learner. A general framework that uses an invariant model learner is outlined in Algorithm 2. For the purposes of this paper, we employ a simple actor-critic setup where the model is used to compute multi-step estimates of the Q value used by the actor learner. A specific instantiation of this idea of model value expansion is the SAC-SV SVG algorithm proposed in Amos al. [2] (see Appendix 3.3 for details). It is important note that the proposed version of model-invariance can be used in combination with any MBRL. method and with any type of dynamics model architecture, such ensembles recurrent architectures.
EXPERIMENTS
Presence Spurious Correlations. We first test the pres- ence of spurious correlations in the Humanoid-v2 [53] task by choosing to predict single dimension (the knee joint) in two contrasting settings; 1) No Mask: when all state and action variables are provided as input and, Mask: when the state variables that are likely uncorrelated to the knee joint are masked (see Appendix 3.5). Having trained different models for the two settings, we observe that 1) No_ Mo_Mask: performs worse than 2) Mask:, for both hori- zon values in {3, 5) (see Figure 3). This verifies that there indeed is an invariant, causal set of parents among the state dimensions and that there could be some interference due to spurious correlations in 1), and thus, it performs worse than case 2). Furthermore, when the dimensions that are likely to be useful in predicting the knee joint are masked,
DM
Mask
No. No.mask
Mesk
No.mask
(a) Horizon=3
(b) 5orizon=5
Figure Efffect of spurious correlation on model learning test loss of Humanoid-v2 for a single dimension (knee joint) with two set- tings: Mask and No mask. 10 seeds, I std. dev. shaded. Y-axis magnitude order is le-3. then the model error the highest (worst; not shown in figure). Note that since all variables are necessary for the task in hand, simply discarding some of them is not solution to the spurious correlation problem.
Invariant Model Learning Humanoid-v2. We compare the invariant model learner to standard MLE leamner for the Humanoid-v? task. To observe the effect of the invariance loss clearly, we decouple the model learning component from the policy optimization component testing the model on data coming from the replay buffer a pre-trained model-fres fodel--free SAC agent. Such a setup ensures that the changes in state distribution according changes in policy still present but any changes

---

in the model do not actively affect this distribution, thus clearly reflecting the effect of model learning only.
We observe that our invariant model Iearner performs much better than the standard model learner, especially when the number of samples available is low, i.c., around the 200K to 500K mark (sce Figure 4, left plot). As the number of samples increases, the performance between both models converges, just as observed in the tabular case. This is expected since in the infinite data regime, both solutions (MLE and invariance based) approach the true model. Furthermore, we observe that the number of samples it takes for convergence between the standard and the invariant model learners increases as the rollout horizon (H in Figure 4) of the model learner is increased. Next, we plot the test error when the invariance loss is computed for all state variables and compare it to when it computed for a particular state variable (Figure 4, middle plot). We see that the per-state version performs worse than the per-state- -ar-sttt-vvrri or per-dim version, showing the importance of enforcing invariance at the state variable level. Finally, we test the mean correlation in the dimensions of the learnt representation & with and without model-invarianne mean correlation lower with model-invariance than without, suggesting that the practical algorithm does induce the sparsity property discussed in the earlier part of the paper.
(Figure 4, right table). We see that the
POPLIN
Cheetah
Walker
Hopper
Ant
*PETS
2288 1:510
282 250
114 a: 311
1165 4: 113
"POPLIN-A
1562 1:568
-105 125
202 481
1148 4: 219
.POPLIN-P
4235 +36
597 239
2055 206
2330 + 160
"METRPO
2283 + 450
-1609 #:38
1272 + 250
282 =
*SAC
4035 134
-382
2020 346
836
SAC-SVG H-3
6530 1380
80 = 472
1 1108 1263
2293 1 397
OURS H-3
7067 4:269
-150
1724 271
3124
199
Table 2: POPLIN Invariant MBRL performance Invariant MBRL performance reported at 200k steps on four MuloCo based domains from POPLIN [55]. represents performance reported hy POPI.NN. Standard етror with 10 seeds reported. We bold the scores with larger mean values.
MBPO
Cheetah
Walker
Hopper
Ant
Humamoid
SAC-SVG H-3
8957 4:32
3795
4:300
3201 # 101
3997 +: 153
1712 #: 415
OURS H-3
9109 +334
3961
4:w
3382
4546 + 286
2443
SAC-SVG H-4
8327 +87
3494
1392
3291 =232
4470 + 307
2404 + 495
OURS H-4
9663 1: 487
4347 1:136
3489
4565 1:201
3506 538
SAC-SVG H-5
5710 1329
2773
1442
3059
276
3808 1:531
3190 601
OURS H-5
5796 8:855
3326
1:430
3207
210
3817
1:488
3446 481
Table 3: MBPO Invariant MBRL performance Invariant MBRL performance reported at 200k steps on five MuJoCo based domains from MBPO [28]. Standard error with 10 sceds reported. We bold the scores with larger mean values.
Invariant Model-based Reinforcement Learning. Finally, we evaluate the invariant model learner within the policy optimization setting of SAC-SVG [2]. We compare the difference in performance to SAC-SVG when the horizon length is varied (see MBPO environments in Table 3 and Appendix 3.4) and then compare the performance of our method against multiple model based methods including PETS [1 POPLIN [55], METRPO [32], and the model free SAC [27] algorithm (see POPLIN environments in Table 2 and Appendix 3.4). The results in Table 3 show a consistent improvement in performance when the invariant model learner (OURS) is used instead of the standard model learner (SAC-SVG) across most tasks, including the Humanoid-v2. Furthermore, we see that as the horizon length is increased, the difference in performance between the invariant and standard learners increases for most tasks. A comparison of model training errors for both cases is provided in Table in Appendix. We see consistently higher training loss for the invariant learner, indicating less over fitting. Finally, in Table 2 we see that the invariant learner outperforms all six baselines on 3 out tasks, with the exception of Walker, where most methods fail to reach a positive score. Combining invariant model learner with other policy optimization algorithms is therefore a promising direction for future investigation.

---

8 RELATED WORK
Planning based on structural assumptions on the underlying MDP have been explored in significant detail in the past [7]. The most closely related setting is of factored MDPs, but approaches that build on the factored MDP assumption have predominantly also assumed known graph structure for the transition factorization [31, 50, 42]. The factored MDPs literature has carried two distinet ideas all along, one of sparsity and the other of factorized transitions/rewards This paper focuses predominantly on sparsity. We rely on factorization for some theoretical results but the practical algorithm does not require factorization to work. Sparsity leads to the abstraction viewpoint introduced in this work, which allows us to scale our formalism to practical methods. It allows us to connect the initial formalism to the concrete problem of neural network generalization (without an abstraction- based formalism, it does not make sense to use neural networks).
On separate axis, papers that deal with learning the graph structure and then doing RL [50, 16, 26], do so by constructing close to optimal transition functions, and then planning. In our practical method, focus only avoiding spurious correlations, which is computationally more efficient since by the time an exact model of the MDP can be learnt (usually done by checking all possible parent set candidates), most model-free methods can just learn the reward function and perform much better. Degris et al. [13] learn decision tree based models online while we focus on learning NN based models. Extending such exact methods NN models remains an open problem. Guestrin et al. [25] state a similar idea to this work but simply assume that one is given basis functions that have limited scope
abstractions).
As mentioned, other forms of state abstraction such as bisimulation [37, 58] do not consider sparsity whereas model-invariance does. Therefore, both forms of abstraction are fact orthogonal each other, and therefore can be combined on top of one another. In similar essence, a lot of works have proposed value-aware model learning objectives, which only model the minimal information required to predict the value function of the task [20, 17, 12]. This is again complimentary idea to model-invariance, since we consider states spaces that are already minimal w.r.t. the value function, ie., no task irrelevant information is present.
CONCLUSION AND FUTURE DIRECTIONS
This paper introduced new type of state abstraction for MBRL that exploits the inherent sparsity present in many complex tasks. We first showed that a representation that only depends on the causal parents of each state variable follows this definition and is provably optimal. Following, introduced a novel approach for learning model-invariant abstractions in practice, which can be plugged in any given MBRL method. Experimental results show that this approach measurably improves the generalization ability of the learnt models. This stands as an important first step to building more advanced algorithms with improved generalization for systems that possess sparse dynamics.
In terms of future work, there remain multiple exciting directions and open questions. First, to enable indel-invariance, we could also look at other kind of approaches proposed recently such as the AND mask [43]. The AND mask specifically requires the data be separated into multiple environments, and thus naturally suits the offline RL setting. Second, moving to pixel based input, the representation learning task becomes two-fold, including learning to abstract away the irrelevant information present in the pixels and then learning model-invariant representation. Third, note that our theoretical results do not involve an explicit dependence on a sparsity measure, for example, the maximum number of parents any state variable could have. Including such dependence would ensure tighter bounds. Fourth, is worth asking how such an explicit constraint on model-invariance can perform as a standalone representation learning objective, considering the strong progress made by self-supervised RI..

---

REFERENCES
[1] Pieter Abbeel, Morgan Quigley, and Andrew Y Ng. Using inaccurate models in reinforcement learning. Proceedings of the 23rd international conference on Machine learning, pp. 1-8, 2006.
[2] Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning, arXiv preprint arXiv:2008. 12775, 2020.
[3] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint arXiv: /907.02899 2019.
[4] Christopher Atkeson and Juan Carlos Santamaria. A comparison of direct and model-based reinforcement learning. In Proceedings of international conference robotics and automation, volume 4, pp. 3557-3564. IEEE, 1997.
[5] D. P. Bertsekas and D. A. Castanon. Adaptive aggregation methods for infinite horizon dynamic programming. IEEE Transactions on Automatic Control, 34(6):589-598, 1989. doi: 10. 100999 24227.
[6] Dimitri P. Bertsekas. Dynamic Programming and Optimal Control. Athena Scientific, Ist edition, 1995. ISBN 1886529124.
[7] Craig Boutilier, Thomas Dean, and Steve Hanks. Decision-theoretic planning: Structural assumptions and computational leverage. Journal of Artificial Intelligence Research, 11:1-94, 1999.
[8] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint 46X6:: 2016.
[9] Jinglin Chen and Nan Jiang. information-theoreti considerations in batch reinforcement learning. In International Conference Machine Learning, pp. 1042-1051. PMLR, 2019.
[10] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International Conference on Machine Learning (ICMI). pp- 1597-1607, 2020.
[11] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamies models. In Advances in Neural Information Processing Systems (NIPS), pp. 4754-4765, 2018.
[12] Brandon Cui, Yinlam Chow, and Mohammad Ghavamzadeh. Control-aware representations for model-based reinforcement learning. arXiv preprint arXiv:2006. 13408, 2020.
[13] Thomas Degris, Olivier Sigaud, and Pieme--Hent Wuillemin. Learning the structure of factored markov decision processes in reinforcement learning problems. In Proceedings of the 23rd international conference Machine learning, pp. 257-264, 2006.
[14] Marc Deisenroth and Carl E Rasmussen. PILCO: A model-based and data-efficient approach to policy search. In International Conference on Machine Learning (ICML), pp. 465-472, 2011 1.
[15] Marc Peter Deisenroth, Dieter Fox, and Carl Edward Rasmussen. Gaussian processes for data-efficient learning in robotics and control. IEEE transactions on pattern analysis and machine intelligence, 37(2):088-42 2013.
[16] Carlos Diuk, Lihong Li, and Bethany R Leffler. The adaptive k-meteorologists problem and its application structure learning and feature selection in reinforcement learning. In Proceedings of the 26th Annual International Conference on Machine Learning, pp. 249-256, 2009.
[17] Pierluca D' Oro, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli. Gradient-awart model- based policy search. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 3801-3808, 2020.

---

[18] Frederick Eberhardt and Richard Scheines. Interventions and causal inference. Philosophy of science, 74(5):981-93 1998 2007.
[19] Eyal Even-Dar and Yishay Mansour. Approximate equivalence of markov decision processes. In Learning Theory and Kernel Machines, pp. 581-594. Springer, 2003.
[20] Amir-massoud Farahmand, Andre Barreto, and Daniel Nikovski. Value- loss function for model-based reinforcement learning. In Artificial Intelligence and Statistics, pp. 1486-1494. PMLR, 2017.
[21] Vladimir Feinberg, Alvin Wan, lon Stoica, Michael I Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based value expansion for efficient model-free reinforcement learning. In International Conference Machine Learning (ICMI.), 2018.
[22] Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pp. 512-519. IEEE, 2016.
[23] Anirudh Goyal and Yoshua Bengio. Inductive biases for deep learning of higher-level cognition. arXiv preprint arXiv:201 15091,, 2020.
[24] Carlos Guestrin, Daphne Koller, and Ronald Parr. Max-norm projections for factored MDPs. In International Joint Conferences Artificial Intelligence (IJCAI), volume pp. 673 182, 2001.
[25] Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificial Intelligence Research, 19:399-466 4689 2003.
[26] Zhaohan Daniel Guo and Emma Brunskill. Sample efficient feature selection for factored mdps. arXiv preprint arXiv: 1703.034544 2017.
[27] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference Machine Learning, pp. 1861-1870. PMLR, 2018.
[28] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, pp. 12519-12530, 2019.
[29] Nan Jiang. Notes state abstractions, 2018. URL http::t://a/:ttaaaa wtt.cc....llnnn illinois. edu/ falu/files/// os 5 98/ note4, pdf.
[30] Anders Jonsson and Andrew Barto. Causal graph based decomposition of factored mdps. J. Mach. Learn. Res., 7:2259-2501 2201,, December 2006. ISSN 1532-4435.
[31] Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored MDPs. In International Joint Conferences on Artificial Intelligence (IJCAI), volume 16, pp. 740-747, 1999.
[32] Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. arXiv preprint arXiv:1802. 1802.10992 2018.
[33] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. arXiv preprint 1111v:1612.0147,, 2016.
[34] Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. Learning for Dynamics and Control (L4DC), pp. 761-770, 2020. URL http:
org/ abs 200000022200 04523.
[35] Alex X Lee, Anusha Nagabandi, Pieter Abbecl, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. arXiv preprint arXiv: / 909900953,... 2019.
[36] Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 13(1):1334-1333, 2016.

---

[37] Lihong Li. A unifying framework for computational reinforcement learning theory. PhD thesis, Rutgers University-Gradunss Graduate School-New Brunswick, 2009.
[38] Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstraction and provably efficient rich-observation reinforcement learning. In International Conference Machine Learning (ICML), pp. 6961-6971, 2020.
[39] Dipendra Misra, Qinghua Liu, Chi Jin, and John Langford. Provable rich observation reinforcement learning with combinatorial latent states. In International Conference on Learning Representations (ICLR), 2021. URL httpt:ttp///cc..ct /opp/roveeieeeee/e / hxl HXF.H
net forum?ici
[40] Jovana Mitrovic, Brian McWilliams, Jacob Walker, Lars Buesing, and Charles Blundell. Repre- sentation learning via invariant causal mechanisms. arXiv preprint arXiv:2010.07922, 2020.
[41] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, arXiv preprint orXiv:1807.03748, 2018.
[42] Ian Osband and Benjamin Van Roy. Near-optimal reinforcement learning in factored mdps. In Advances in Neural Information Processing Systems, pp. 604-612 2014.
[43] Giambattista Parascandolo, Alexander Neitz, Antonio Orvieto, Luigi Gresele, and Bernhard Schölkopf. Learning explanations that are hard to vary. arXiv preprint arXiv:2009.0032) 2020.
[44] Judea Pearl. Causality. Cambridge university press, 2009.
[45] Judea Pearl et al. Models, reasoning and inference. Cambridge, UK: CumbridgeUnniversityyresss 19, 2000.
[46] Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference using invariant prediction: identification and confidence intervals. arXiv preprint
2015.
[47] Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
[48] Balaraman Ravindran and Andrew G Barto. An algebraic approach to abstraction in reinforce- ment learning. PhD thesis, University of Massachusetts at Amherst, 2004.
[49] Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimiza- tion. arXiv preprint arXiv:2010.05761, 2020.
[50] Alexander L Strehl, Carlos Diuk, and Michael L Littman. Efficient structure learning in factored-state mdps. In AAAI, volume 7, pp. 645-6500 2007.
[51] Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4): 160,16,,, 1991.
[52] Yi Tian, Jian Qian, and Suvrit Sra. Towards minimax optimal reinforcement learning in factored markov decision processes. Advances in Neural Information Processing Systems, 33, 2020.
[53] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model- based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 0026- IEEF, 2012.
[54] Hado P Hasselt, Matteo Hessel, and John Aslanides. When use parametric models in reinforcement learning? Advances in Neural Information Processing Systems, 32:14322-11333, 2019.
[55] Tingwu Wang and Jimmy Ba. Exploring model-based planning with policy networks. arXiv preprint arXiv:1906.0 190..09990, 2019.
[56] Ziping Xu and Ambuj Tewari. Near-optimal reinforcement learning in factored mdps: Oracle- efficient algorithms for the non-episodic setting. 2020.

---

[57] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block MDPs. In International Conference Machine Learning (ICML), pp. 11214-11224, 13-18 Jul 2020.
[58] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Invariant representations for reinforcement learning without reconstruction. In International Conference on Learning ieterssennntiionnsssss 2021. URL https: / /wwpeeeeeeeeeeeeeeooeeeee. net/ forum? i = 2FCwDKRREU.


Reviews:
[{"review_id": "wUpoW0Fyo1", "reviewer": "Reviewer_9B8W", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In general, I think that the paper is well-written and I like the core idea of this paper because it points out an important direction to improve the data efficiency and generalization ability of the dynamics model. However, I cannot give the accept score in this procedure because there are some unclear points in this paper, but I am willing to increase my score if the authors fix the below problems in the rebuttal process.\n\n1. The authors only show the effect of spurious correlation on the Humanoid env, which is not enough to validate that removing spurious correlation is important for RL problems. It is better to show more examples.\n\n2. Similar to problem 1, the authors only report the mean correlation metric with and without the proposed model-invariance loss, can the authors provide more evidence to validate the effectiveness of model-invariance loss for removing spurious correlation?\n\n3. Why do not explicitly learn the causal structure of the dynamics? What is the advantage of the proposed implicit model-invariant loss compared with explicitly learning and exploiting the causal structure?\n\n4. Would the performance of the proposed invariance loss be improved if we increase the number of randomly initialized dynamics models?\n\n5. I notice that the difference in performance between the invariant and standard learners reduces when the horizon increases from 3 to 5 for most tasks, which violates the claim \n>the horizon length is increased, the difference in performance between the invariant and standard learners increases for most tasks. \n\n Why is this a phenomenon? Can authors add an explanation for it?\n\n6. The authors only show the training env prediction errors in Table 5, but do not report test env prediction errors, which is insufficient to support the claim\n>We see a consistently higher training loss for the invariant learner, indicating less overfitting.\n\nMinor Points: \nThere are several missing Table/Section references, e.g. Algorithm 1/ Appendix 3.5.\n", "summary_of_the_paper": "This paper points out that the causal structure is important for the generalization ability of a dynamics prediction model, and empirically show that learning a model-invariance representation can implicitly learn the causal sparsity and improve the generalization ability of the learned model.", "summary_of_the_review": "In summary, the paper focus on an important topic of the reinforcement learning area. The paper is also well-written, but the analysis and the experimental support are not strong enough, so I give 5 in this procedure, but I am willing to increase my score if the authors fix them in the rebuttal process.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "edeR6Mz5GyL", "reviewer": "Reviewer_hNYg", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper presents in interesting, useful idea that causality plays in important role in generalization performance in model-based reinforcement learning, and proposes an interesting method to leverage this. The paper rests on a fundamental assumption that in many cases, the full state of a system can be decomposed, with the probabilistic model governing each state dimension being independent given the previous state and action, and furthermore, each state dimension being causally dependent only on a select portions of state vector. While this isn't true generally, I buy the argument that for a lot of systems with high dimensional state representations, such as multi-agent systems or the MuJoCo systems considered in this work, this may be approximately correct. Indeed, even if this isn't exactly true, it may be worth making this assumption for improved generalization performance in the early stages of MBRL where data is sparse. \n\nMy main concerns with this work are the following:\n\n- Disconnect from theory to practice:\n    - The theoretical insights connecting RL to ICP are only valid for the very limited linear MDP setting. While in many cases, it is useful to derive theory in a linear case to inform practical algorithms for the nonlinear case, but in this work, the practical approach put forth in the paper seems disconnected from the linear case. There is fundamentally different strategy for intervention, moving from interventions on the data to interventions in the form of learning models from different data samples. There is no call to a ICP subroutine; instead, the approach adapts the ReLIC architecture for self-supervised intervention-driven representation learning. While the connections to causal prediction are clear for the ReLIC model, the architecture proposed by the authors is different, and indeed designed for a supervised learning problem rather than the self-supervised learning problem tackled by ReLIC. These differences make it hard for me to connect the practical algorithm proposed by the authors to the causal learning ideas they use to motivate it.\n\n- Lack of clarity:\n    - I found several details to be hard to follow in the presentation of this work. For example, in Lemma 1, $P_\\phi(s,a)$ is defined to be a mapping from $s,a$ to a distribution on $\\mathcal{X}$, as it is the expectation over $P(\\cdot | x, a)$, which is a distribution on $\\mathcal{X}$. However, in the text, it is described as a transition model in the lifted state, which to me implies it should be a mapping from $s,a$ to a distribution on the abstract state space. Furthermore, in the nonlinear architecture, the details of the invariance loss are confusing. The text refers to $\\psi$ as the cosine similarity between the predictions of $f$ and $h$, but the definition suggests an inner product between the projections of the predictions. The authors should be more clear if they are using the angled brackets to denote cosine similarity rather than the standard inner production. More importantly, the definition of $\\psi^i(f,h)$ as an inner product suggests it maps to a scalar, so it doesn't make sense to apply a softmax to this result directly. The pytorch-like code in the appendix does not provide any additional clarity, as the code does does not distinguish between elementwise and matrix multiplication and the dimensionality of the variables is not clear. \n    By referring to the ReLIC paper, I was able to determine that the softmax is applied over the batch dimension, implying a auxiliary task of instance identification within the batch, and a desire to enforce invariance to interventions for this task. However, this does not make sense for the approach in this paper, since model learning is a supervised task, so one could enforce invariance directly on the predictions of the next state. Indeed, in the ReLIC paper, the loss includes both a invariance term computed on $p^{fh}$, and a instance identification loss on $p^{fh}$. However, in the proposed approach, there is only a KL loss, so what is stopping the projection network $g$ from learning a trivial projection mapping all inputs to the same vector, such that  $p^{fh} =p^{hf}$ without enforcing any invariance on the actual predictions of the next state?\n\nMinor comments / questions:\n- Do you have references to support why Assumption 2 is reasonable? What is the distinction between $x_{t}$ and x_{t+1}$ over a rollout of $x_1, a_1, \\dots, x_T$.\n- What, specifically, was the linear MDP used in the tabular RL experiment? The specific choices of MDP and policies used to create ICP environments should be detailed in the appendix.\n- Bolding results based purely on mean value is nonstandard and obscures the fact that in many cases the proposed approach is not statistically significantly more performant than the SAC-SVG baseline, e.g. in table 3. Additionally, why is the result for \"Ours H-3\" on Hopper in table 2 bolded? \n- In Figure 4, the test model error does not seem to have converged by the end of the plot. It would be useful to see this asymptotic performance to show whether the invariant model converges to the same quality of model as the standard model. ", "summary_of_the_paper": "This work presents a method for learning representations for model based reinforcement learning by leveraging assumed sparsity in the causal graph to improve the generalization performance of the learned model in the low-data regime. The authors propose a method which, at a high level, aims to utilize data collected under different conditions (e.g. policies) to extract a representation under which the transition model is approximately invariant. Furthermore, they provide a bound on difference between the optimal value function on the learned representation and the true optimal value function as a function of the approximation quality of the invariant model. They show how invariant causal prediction can be used in the the linear tabular setting to yield improved performance relative to learning a model directly in the original state space. Next, they present an approach to achieve a similar objective for the nonlinear case, presenting a neural network architecture based on that proposed in ReLIC, and show that the proposed architecture improves the low-data performance in learning dynamics models, and leads to improved performance relative to baselines.", "summary_of_the_review": "Overall, I found that this paper has a good, theoretically motivated derivation to improve low-data generalization in MBRL, which is a clear bottleneck. However, the proposed practical algorithm was hard to understand, and fundamentally disconnected from the theoretical analysis presented in this paper. Hence, I do not support acceptance as is. In a revised version of the paper, I recommend that the authors draw stronger connections between causal prediction to the practical implementation that they propose.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VPOSHmqysqO", "reviewer": "Reviewer_avfm", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- The idea of combining the two concepts of sparsity and causal invariance is not new and has been frequently mentioned in the causality literature, e.g., sparse shift mechanism etc. See more in [1].\n\n- In Assumption 1, the authors do not consider instantaneous effects, which render such an assumption quite limited in real world applications.\n\n- In the paragraph of **Invariant Causal Prediction** (Section 2), the authors describe ICP as learning an invariant representation. This is misunderstanding about ICP, because the representation is allowed to vary across environments. In fact, in ICP, it is only assumed that the causal mechanism is invariant, instead of the representation. \n\n - In Definition 1, \"$:=$\" under $\\sup$ should be corrected. \n\n- It is NOT true to say that \"This change in distribution (over states and action) can be realized as a change in the policy\". Since $p(x, a)=p(a|x)p(x)$, the change in $p(x, a)$ can result only from a change in $p(x)$, when the policy $p(a|x)$ is invariant. Thus, the subsequent claim \"each policy \\pi defines an ICP environment $e$\" is also wrong. All this might lead to some inconsistencies in the rest of the paper.\n\n- I also have some doubts about the claim that \"when sparsity is present, following the above definition leads to a representation $\\phi$ such that any two components $\\phi_i$ and $\\phi_j$ are decorrelated.\" Because $\\phi_i$ and $\\phi_j$ could have some overlap, in this case the two components are no longer decorrelated. \n\n- I feel that the theoretical results presented in the paper are not quite new, all of which are immediate results of or simple extension to previous work. Specifically, Proposition 1 is just an immediate result of directly applying ICP to the MDP setting. Theorem 1 is also a simple extension to Theorems 2&3 in [2].\n\n- Missing appendix index in Algorithm 1.\n\n- How to select the model $h$? Why are the parameters of $h$ not updated during the training phase? How to determine the number of models? Because two models might not be enough in many scenarios. \n\n- In the linear settings, it is convincing that the learned model could generalize well to unseen environments, because causal variables can be identified by means of ICP. However, this is not the case in the nonlinear settings. As we know, invariance does not imply causation. Hence, the proposed approach (i.e., Eq 2) to approximately learning such abstract states has no guarantees on that causal variables can be identified from the training data. This leads to no guarantees on good generalization performance in unseen environments.\n\n- All the baselines to which the proposed approach is compared are not quite related to invariant predictors or something similar that this paper studies. This makes it hard to judge the validity of the proposed approach. I suggest that the authors should compare it with more relevant methods. For example, a natural and simple baseline to compare with is IRM, where IRM can be directly used to learn the invariant transition function. \n\nReferences:\n\n[1] Schölkopf et al. Towards causal representation learning. 2021.\n\n[2] Zhang et al. Invariant causal prediction for block MDPs. 2020.", "summary_of_the_paper": "The authors introduce a new type of state abstraction called model-invariance by integrating sparsity with causal invariance in the single-task RL scenarios. In the linear settings, they learn the model-invariance states by directly applying ICP and prove that an optimal policy can be learned over them. In the nonlinear settings, they instead propose a practical method to approximately learn such model-invariance states. ", "summary_of_the_review": "On the theoretical side, the main theorems are either immediate results of or simple extension to the previous work. On the technical side, they draw on ICP and ReLIC in the linear and nonlinear settings, respectively. On the experimental side, the baselines are not convincing enough to judge the validity of the proposed approach. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "L_4_qQn8n77", "reviewer": "Reviewer_5dqM", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "===== After rebuttal: the authors didn't address my concerns. I'll decrease my score accordingly. ==== \n\nThis paper is well motivated. The model invariance structure indeed rises naturally in continuous control tasks as argued by the authors in Section 7. Identifying the direct causes using invariant causal prediction methods in reinforcement learning could potentially have a wide range of applications.\n\nThe paper is well written and easy to follow. The toy example in Section 5 illustrates how ICP helps model learning. Technical part of this paper is self-contained and well-explained.\n\nOn the plus side, this paper directly tackles the problem of model learning, which is a central question in model-based RL area. The model learning algorithm proposed by this paper reduces spurious correlation in mode learning and improves state-of-the-art results in low data regime. More importantly, the model learning algorithm is self-contained and can be plugged into almost any existing model-based RL algorithms, and achieves impressive performance improvement. The empirical evaluations are solid.\n\nOn the minus side, my main concerns are the following.\n-\tIn the beginning of Section 3, the invariant abstraction \\phi_i is defined w.r.t. a fixed index i. In this case, it’s expected that the abstraction can be much smaller than the state space. However, in definition 1, the abstraction \\phi must hold for every index i (although w.r.t. different errors \\epsilon_{i}). In this case, the existence of a small abstraction is unclear to me. For example, in the toy example in Section 5 (where there are three state variables and each depends only on its own), what’s the minimal model invariant abstraction according to definition 1? If I understand correctly, the transition probability must be the same for every index i, which implies that different states must have different abstraction.\n-\tIt’s hard for me to follow the intuition behind Algorithm 2. In both the theory and the toy example, PA(x^i) is predicted using the invariance across different policies/environments. Is algorithm 2 using the same invariance? It seems to me that algorithm 2 regularizes the learned feature by requiring the similarity of separately learned transition dynamics. Why this regularization helps reducing spurious correlation (not only the performance. Although the improvement in performance is impressive, the core intuition introduced in previous sections is to reduce spurious correlation), in either theoretical or empirical perspective?\n\nAdditional questions.\n-\tWhat’s |\\phi| in the paragraph below Theorem 1, Page 4? Is it referring to the number of abstract states or the side of hypothesis space of state abstraction?\n-\tThe toy example focus on offline setting with replay buffers from multiple policies. But in the online setting, the policy change is expected, so naturally there will be multiple policies. What’s the reason for considering offline setting? Does the ICP method work for the online setting?\n-\tTable 1 shows that the transition prediction error is lower using model-invariance. But for Mujoco tasks, the error is higher. The authors claim that in both cases, model invariance helps model learning (but for different reasons). Why the results/arguments are not consistent?\n-\tWhat does “linearity in dynamics” mean in the bottom of Page 5? What’s the exact transition probability?\n-\tThe variable k is not defined in Eq. (1).\n\nMinor issues:\n-\tBroken refs: Line 3 of Alg. 1; below Theorem 1 in page 17; first line of Section 3.5 in Appendix, page 20.\n\nSome suggestions that don’t related to my assessment:\n-\tAll of the empirical results focus on low data regime (which, by the way, is perfectly fine because that’s the scope of this paper). It would be helpful to also include a result on the speed of convergence, and also the final performance upon converngence.\n-\tThe inner product notation in the first line of page 7 is a little confusing, because standard inner products are commutative.\n", "summary_of_the_paper": "This paper defines a novel model-invariant state abstraction for factored MDPs. It’s also shown that using invariant causal prediction significantly reduces transition prediction error in both toy example and several continuous control tasks. Inspired by theoretical results, this paper also proposes a novel method for learning model-invariant representations using neural networks. The model learning method can be combined with combined with almost any model-based deep RL algorithm. Empirically, this paper shows that in continuous control tasks, their model learning algorithm improves the performance of MBRL algorithms in the low data regime.", "summary_of_the_review": "This paper is well-motivated, the empirical results are impressive. My main concerns are (1) the existence of small model-invariant abstraction, and (2) the intuition of Alg. 2. As a result, I can only recommend a weak accept at this point. But I’m open to raising my score if my concerns are addressed.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a method to learn representations in MBRL by exploiting sparsity in the model to improve data efficiency. The key idea is to build a representation for which the model is invariant.
The idea is quite interesting, but one weakness of the current draft is that there is a disconnect between the presented theory (linear case) and the relevant experimental setup (non-linear).
The paper is overall well written but would still benefit from a revision to improve clarity as pointed out by the reviewers.
The experimental results are inconclusive due to the choice of weak baselines.
Decision: Reject