Contrastive Unsupervised Learning of World Model with Invariant Causal Features
Authors: Rudra P. K. Poudel
Harit Pandya
Roberto Cipolla
Venue: Submitted to ICLR 2023 (2023)
Keywords: world models
causality
contrastive learning
model-based reinforcement learning
reinforcement learning
out-of-distribution generalisation
sim-to-real transfer
robot navigation

Abstract: In this paper we present a world model, which learns the causal features using invariance principle. We use contrastive unsupervised learning to learn the invariant causal features, which enforces invariance across augmentations of irrelevant parts or styles of the observation. Since the world model based reinforcement learning methods optimize representation learning and policy of the agent independently, contrastive loss collapses due to lack of supervisory signal to the representation learning module. We propose depth reconstruction as an auxiliary task to explicitly enforce the invariance and data augmentation as style intervention on the RGB space to mitigate this issue. Our design help us to leverage state-of-the-art unsupervised representation learning method to learn the world model with invariant causal features, which outperforms current state-of-the-art model-based as well as model-free reinforcement learning methods on out-of-distribution point navigation tasks on Gibson and iGibson dataset at 100k and 500k interaction step benchmarks. Further experiments on DeepMind control suite even without depth reconstruction, our proposed model performs on par with the state-of-the-art counterpart models.

Under review as a conference paper at ICLR 2023
CONTRASTIVE UNSUPERVISED LEARNING OF WORLD MODEL WITH INVARIANT CAUSAL FEATURES
Anonymous authors Paper under double-blind review
ABSTRACT
In this paper we present world model, which learns causal features using the in- variance principle. In particular, we use contrastive unsupervised learning learn the invariant causal features, which enforces invariance across augmentations of irrelevant parts or styles of the observation. The world-model-be moadel-based reinforcement learning methods independently optimize representation learning and the policy, Thus naive contrastive loss implementation collapses due to a lack of supervi- sory signals to the representation learning module. We propose invariant auxiliary task to mitigate this issue. Specifically, we use data augmen- tation style intervention on the RGB observation space and depth prediction as an auxiliary task to explicitly enforce the invariance. Our proposed method significantly outperforms current state-of-the-art model-based and model-free re- inforcement learning methods on -ittr--distribttion point navigation tasks on the iGibson dataset. Morcover, our proposed model excels at the sim-to-real trans- fer of our perception learning module. Finally, we evaluate our approach on the DeepMind control suite and enforce invariance only implicitly since depth is not available. Nevertheless, our proposed model performs par with the state-of- the-art counterpart.
intervention
INTRODUCTION
An important branching point reinforcement learning (RL) methods whether the agent learns with or without a predictive environment model. In model-based methods, an explicit predictive model of the world is learned, enabling the agent to plan by thinking ahead (Deisenroth & Ras- mussen 2011 Silver et al.] 2018 Ha & Schmidhuber] 2018 Hafner et al.] (021). The alternative model-free methods do not learn the predictive model of the environment explicitly as the control policy is learned end-to-end from the pixels. As consequence, model-free methods do not consider the future downstream tasks. Therefore, we hope that model-based methods are more suitable for -itt---diittiittionn (OoD) generalization and sim-to-real transfer.
A model-based approach has to learn the model of the environment purely from experience, which poses several challenges. The main problem is the training bias in the model, which can be ploited by an agent and lead to poor performance during testing (Ha & Schmidhuber 2018). Fur- ther, model-based RL methods learn the representation using observation reconstruction loss, for example variational autoencoders (VAE) (Kingma & Welling 2014). The downside of such a state abstraction method is that is not suited separate the task relevant states from irrelevant ones, resulting in current RL algorithms often overfit to environment-specde 2020 Hence, relevant state abstraction is essential for robust RL model, which the aim this paper.
characteristics Zhang et al.
Causality the study of learning cause and effect relationships. Learning causality in pixel-based control involves two tasks. The first is causal variable abstraction from images, and the second is learning the causal structure. Causal inference uses graphical modelling (Lauritzen & Spiegelhalter 1988), structural equation modelling (Bollen 1989), or counterfactuals (Dawid) 2000). Pearl (2009) provided an excellent overview of those methods. However, in complex visual control tasks the number of state variables involved is high, inference of the underlying causal structure of the model becomes intractable (Peters et al. (0166). Causal discovery using the invariance principle tries to overcome this issue and is therefore gaining attention in the literature (Peters et al. 2016 Arjovsky

---

Under review as a conference paper at ICLR 2023
Environment
Observation
Contrastive learning
Invariant causal features
Auxiliary tasks
Future hallucination
Memory
Dynamics
Controller
Figure I: Flow diagram of proposed World Model with invariant Causal features (WMC). consists of three components: unsupervised causal representation learning, ii) memory, and iii) controller.
al. 2020 Zhang al. 2020 Mitrovic et al. 2021]. Arjovsky al, (2020) learns robust classifiers
based on invariant causal associations between variables from different environments. /Zhang et al.) (0200) uses multiple environments to learn the invariant causal features using a common encoder. Here spurious irrelevant features are learnt using environment specific encoders. However, these methods need multiple sources of environments with specific interventions or variations. In contrast, propose using data augmentation
source of intervention, where samples come from as little as single environment, and we use contrastive learning for invariant feature abstraction. Related work, Mitroovic al. (2021) proposed regularizer for self-supervised contrastive learning. On the other hand, we propose an intervention invariant auxiliary task for robust feature learning.
Model-based RL methods do not learn the feature and control policy together to prevent the greedy feature learning. The aim that the features of model-based RL will be more useful for various downstream tasks. Hence, state abstraction uses reward prediction, reconstruction loss or both (Ha & Schmidhuber 2018 Zhang al., 2020; Hatner et al.) 202T). On the other hand contrastive learning does not use the reconstruction of the inputs and applies the loss at the embedding space. Therefore, we propose causally invariant auxiliary task for invariant causal features learning. Specifically, utilize depth predictions extract the geometrical features needed for navigation, which not dependent on the texture. Finally, we emphasize that depth is not required for deployment, enabling wider applicability of the proposed model. Importantly, setup allows contrastive learning on model-based RL. methods and improves the sample efficiency and the OoD generalization.
popular
In summary, we propose a World Model with invariant Causal features (WMC), which can extract and predict the causal features (Figure []. Our WMC is verified the point goal navigation task from Gibson (Xia et af.) 2018) and iGibson 1.0 (Shen et al.] (021) as well as the DeepMind control suite (DMControl) (Tunyasuvunakool et al. 2020). Our main contributions are:
1. to propose a world model with invariant causal features, which outperforms state-of-ther models out-of-distribution generalization and sim-to-real transfer of learned features.
2. propose intervention invariant auxiliary tasks improve the performance.
3. to show that world model benefits from contrastive unsupervised representation learning.
2

---

Under review as a conference paper at ICLR 2023
2
RELATED WORK
Unsupervised Representation Learning, Learning reusable feature representations from large un- labeled data has been an active research area. the context of computer vision, one can leverage unlabeled images and videos to learn good intermediate representations, which can be useful for a wide variety of downstream tasks. A classic approach to unsupervised representation learning is clustering similar data together. for example, using K-means (Hartigan & Wong 1979). Recently, VAE (Kingma& & Welling 2014) has been a preferred approach for representation learning in model- based RL Ha & Schmidhuber 20188) Since VAE does not make any additional consideration of downstream tasks, invariant representation learning with contrastive loss has shown more promising results (Anand et al., 2019 Laskin et el. 2020a). Further, building on the success of supervised deep learning. getting supervision from the data itself is a preferred approach of representation learning from the unlabelled data, which is known as self-supervised learning. Self-supervised learning for- mulates the learning as a supervised loss function. In image-based learning self-supervision can be formulated using different data augmentations, for example, image distortion and rotation (Dosovit- skiy et al. 2014; Chen et al., 2020). We also use different data augmentation techniques to learn the invariant features using contrastive loss, which we explain below.
Contrastive Learning. Representation learning methods based on contrastive loss (Chopra et al., 2005) have recently achieved ofate-of-the-art performances on face verification task. These methods use contrastive loss to learn representations invariant to data augmentation (Chen et al. 2020 He jet al., 2020). Given a list of input samples, contrastive loss forces samples from the same class to have similar embeddings and different ones for different classes. Since class labels are not available in the unsupervised setting, contrastive loss forces similar embedding for the augmented version of the same sample and different ones for different samples. In contrastive learning setting, an augmented version of the same sample is known as a positive sample and different samples as a negative sample; these are also referred to query and key samples, respectively. There are several ways of formulating the contrastive loss such as Siamese (Chopra et al., 2005). InfoNCE (van den Oord et al. 2018) and SimCLR (Chen et al. 2020). We chose InfoNCE for our contrastive loss in this work.
(van den Oord et al. 2018)
Causal Inference using Invariant Prediction. Learning structured representations that capture the underlying causal mechanisms generating data is central problem for robust machine learn- ing systems (Schölkopf al., 2021). However, recovering the underlying causal structure of the environment from observational data without additional assumptions is a complex problem. A re- cent successful approach for causal discovery, in the context of unknown causal structure, is causal inference using invariant prediction (Peters et al., 2016). The invariance idea is closely linked to causality under the terms autonomy and modularity (Pearl 2009) or stability (Dawid & Didelez, 2010). Furthermore, it is well known that causal variables have an invariance property, and this group of methods try to exploit this fact for causal inference. Our proposed world model also ex- ploits this fact to learn the causal state of the environment using the invariance principle, which is formalized using contrastive loss. Also, this view of self-supervised representation learning us- ing invariant causal mechanisms is recently formalized by Mitrovic al. (2021). In section 3 we explain how we utilized the data augmentation technique to learn the causal state of the environment.
Model-based RL. The human brain discovers the hidden causes underlying an observation. Those internal representations of the world influence how agents infer which actions will lead to a higher reward (Hamrick) 2019). An early instantiation of this idea was put forward by Sutton (Sutton 1990). where future hallucination samples rolled out from the learned world model are used in addition the agent's 's interactions for sample efficient learning. Recently, planning through the world model has been successfully demonstrated in the world model by Ha & Schmidhuber (2018) and DreamerV2 by Hafner et al. (202T). In our work we propose to learn causal features to reduce the training biases and improve the sample efficiency further.
Sample Efficiency. Joint learning of auxiliary tasks with model-free RL makes them competitive with model-based RL in terms of sample efficiency. For example, the recently proposed model- free RL method called CURL (Laskin et al. (220a)) added contrastive loss as an auxiliary task and outperformed the state-of-the-ee
model-bases RL method called Dreamer (Hafner al., 2020 Also, two recent works using data augmentation for RL called RAD (Laskin et al. 2020b) and DrQ Yarats et al. 2021) outperform CURL without using auxiliary contrastive loss. These results
3

---

Under review as a conference paper at ICLR 2023
a
C
S
Figure 2: Observation made content (C), causal variables, and style (S), spurious variables. We want representation learning to extract the content variables only, i.e. cause of the action (a).
warrant that if an agent has access rich stream of data from the environments, an additional auxiliary loss unnecessary, since directly optimizing the policy objective is better than optimizing multiple objectives. However, do not have access to rich stream of data for many interesting problems, hence sample efficiency still matters. Further, these papers do not consider the effect of auxiliary tasks and unsupervised representation learning for model-based RI., which is the main focus this paper.
PROPOSED MODEL
The data flow diagram of our proposed world model is shown in Figure D consider the vi- sual control task as a finite-horizon partially observable Markov decision process (POMDP). We denote observation space, action space and time horizon as O, A and T respectively. An agent performs continuous actions at p(a,loss, act), and receives observations and scalar rewards oc, Ft ~ p(o,, Fz Paloce: acc)) from the unknown environment. We use the deep neural networks for the state abstractions from observation 8; p(s,[st-1. 0,-1, 0z), predictive transition model ~ ( at-1) and reward model т; ~ ((ra(s)). The goal of an agent is to maximize the expected total rewards Fip (ET= T(). In the following sections we describe our proposed model in detail.
3.1 INVARIANT CAUSAL FEATURES LEARNING
Learning the pixel-based controller involves two main tasks, first abstraction of the environment state and second maximizing the expected total reward of the policy. Since visual control is complex task, the number of causal variables that involve are high, making the causal structure discovery difficult task as it requires fitting the graphical model or structural equation. That why we choose invariant prediction for causal feature learning. The key idea is that if consider all causes of an effect, then the conditional distribution the effect given the causes will not change when we change all the other remaining variables of the system (Peters et al.] (016). Such experimental changes are known as interventions in the causality literature. We have used data augmentation as our mechanism of intervention, since usually we do have access to the causal or spurious variables or both of the target environment. We have also explored texture randomization as an intervention in our experiments, which we call action replay.
We have shown the high-level idea of the proposed causal features extraction technique in Fig- ure [ The main idea is observation is made of content (C), causal variables, and style (S), spurious variables. We want our representation learning to extract the content variables only, which are the true cause of the action. other words we want our control policy to learn P(a|c ineeariant_encdder efficient, and robust to OoD generalization and sim-to-real transfer. We have chosen the contrastive learning technique (Chen et al. 2020 Grill et al. 2020) to learn the invariant causal features, which means embedding features of a sample and its data augmented version should be the same. We spatial jitter (crop and translation), Gaussian blur, color jitter, grayscale and cutout data aug- mentation techniques (Chen et al.) 2020 Laskin et al.] 2020b] Yarats et al., (02T) for contrastive learning. Since the intervention of the style is performed at the image level, do not need to know
but not P(a s) encoder(o)) as causal variables are sample
4

---

Under review as a conference paper at ICLR 2023
the location of the interventions. Further, the data also come from observation of the different experimental environments. The theoretical guarantee of causal inference using invariant prediction is discussed in Peters et al. (2016) and Mitrovie et al. (202I). However, proposed method does not consider the hidden confounding variables that influence the target effect variable.
3.2 WORLD MODEL
The proposed World Model with invariant Causal features (WMC) consists of three main compo- nents: i) unsupervised causal representation learning, ii) memory, and iii) controller. The represen- tation learning module uses the contrastive learning for invariant causal state learning. The memory module uses a recurrent neural network. It ouiputs the parameters of a categorical distribution, which used to sample the future states of the model. Finally, the controller learns the action probabil- ity to maximize the expected reward using an actor critic approach [Williams) 2021]. We have adopted DreamerV2 (Hafner et al. (02T) to test our proposed hypothesis, which describe below.
1992 Hafner et al.)
Unsupervised causal features learning. Invariant causal feature extraction from the RGB image observation is key component of our model. As described previously, we learn invariant causal features by maximizing agreement between different style interventions of the same observation via contrastive loss in the latent feature space. Since the world model optimizes feature learning and controller separately to learn better representative features for downstream tasks, our early experi- ment with only rewards prediction was poor, which verified in our experiments. Another reason for separate training of the world model and controller is that most of the complexity resides in the world model (features extraction and memory), so that controller training with RL methods will be easier. Hence, we need a stronger loss function to learn a good state representation of the environ- ment. That why we propose depth reconstruction as an auxiliary task and do not use image resize data augmentation to keep the relation of object size and distance intact. We used InfoNCE (van den Oord et al. 2018) style loss learn the invariant causal feature. Hence, our encoder takes RGB observation and task specific information as inputs and depth reconstruction and reward prediction as targets. The invariant state abstraction is enforced by contrastive loss. The proposed invariant causal features learning technique has the following three major components:
A style intervention module that uses data augmentation techniques. We use spatial jitter, Gaussian blur, color jätter, grayscale and cutout data augmentation techniques for style intervention. Spatial jitter implemented by first padding and then performing random crop. Given any observation o,, our style intervention module randomly transforms it into two correlated views of the same observations. All the hyperparameters are provided in the appendix.
We use an encoder network that extracts representations from augmented observations. We follow the same configurations DreamerV2 for fair comparison, and only contrastive loss and depth reconstruction tasks are added. We obtain š = invariant encoder(o)) then follow the DreamerV2 obtain the final state 8, p(se 8t-1, 1, 8.). We use the contrastive loss immediately after the encoder 6t.
Contrastive loss defined for a contrastive prediction task, which can be explained as a differentiable dictionary lookup task. Given a query observation q and a set K (kg. k2n)) with known positive {k+] and negative {k_} keys, the aim of contrastive loss is to learn representation in which positive sample pairs stay close to each other while negative ones are far apart. In contrastive learning literature g, K, k4 and k are also referred anchors, targets, positive and negative samples. We use bilinear products for projection head and InfoNCE loss for contrastive learning (van den Oord et al. 2018). which enforces the desired similarity in the embedding space:
log
(1)
exp(g" + M exp(g"
2((B-)
This loss function can be seen the log-loss of a 2B-way softmax classifier whose label is k4. Where B is a batch size, which becomes 2B after the style intervention module ran- domly transforms in two correlated views of the same observation. The quality of features using contrastive loss depends on the quality of the negative sample mining, which is a dif- ficult task in unsupervised setting. We slice the sample observations from an episode at
5

---

Under review as a conference paper at ICLR 2023
each 5th time-step reduce the similarity between neighbouring negative observations. In summary, causal feature learning has the following component and is optimized by Equa- tion[]
Invariant causal model:
(0(())
Future predictive memory model. The representation learning model extracts what the agent sees each time frame but also is achieved with the memory model and implemented with a recurrent neural network. Further, the transition model learns predict the future state using the current state and action in the latent space only, which enables future imagination without knowing the future observation since we can obtain the future action from the policy if know the future state. Hence, this module is called a future predictive model and enables efficient latent imagination for planning (Ha & Schmidhuber 2018 Hafner et al. 2020). In summary, dynamic memory and representation learning modules are tightly integrated and have the following components,
agent remember the important events from the past. This
Representation model:
Depth prediction model:
@o (oj
(3)
Reward model:
qo (Ta(s))
Predictive memory model:
qo (8, 18-1 Qt-1).
All the world model and representation losses were optimized jointly, which includes contrastive, depth prediction, reward and future predictive KL regularizer losses respectively,
LwM = Ep Σ (o + +In + (ng((ee(s) BEKL)
(4)
where, ekc = -KiiiKK(ppsssssi 0
Controller. The objective of the controller is to optimize the expected rewards of the action, which optimized using an actor critic approach. The actor critic approach considers the rewards beyond the horizon. Since we follow the V2 an action model and value model are learnt in the imagined latent space of the world model. The action model implements a policy that aims predict future actions that maximizes the total expected rewards in the imagined environment. Given H as the imagination horizon length, 7 the discount factor for the future rewards, action and policy model are defined as follows:
Action model:
as(aelse)
(5)
Value model: )
t 7,T -T-1r
3.3 IMPLEMENTATION DETAILS
We have used the publicly available code of D2eamerV2 (Hafner et al. 2021) and added the con- trastive loss on top of that. We have used default hyperparameters of the continuous control task. Here, we have explained the necessary changes for our experiments. Following MoCo (He et al.) 2020) and BYOL (Grill et al., (2020) we have used the moving average version the query encoder to encode the keys K with momentum value of 0.999. The contrastive loss is jointly optimized with the world model using Adam (Kingma & Ba 2014). encode the task observations we used two dense layers of size 32 with ELU activations (Clevert al. 2015). The features from RGB im- age observation and task observation are concatenated before sending to the representation module of the DreamerV2 V2. Replay buffer capacity is 3 for both 100k and 500k steps experiments. All architectural details and hyperparameters are provided in the appendix. Further, the training time of WMC is almost twice than that of DreamerV2 but inference time is the same.

---

Under review as a conference paper at ICLR 2023
Table 1: Experiment results PointGoal navigation task from iGibson 1.0 dataset. Even though data augmentation (DA) improves the DreamerV2 with depth (D) reconstration rather than RGB image(I), proposed WMC further improves the results by significant margin.
Ihlen_0.nm intt
Ihlen Lint
Rs int
Env Avg
Models
Steps
SR
SPL
SR
SPL
SR
SPL
SR
SPL
RAD
100k
0.5
0.01
0.1
0.00
0.8
0.01
0.5
0.01
CURL
100k
8.0
0.07
0.6
0.00
5.4
0.05
4.6
0.04
DreamerV2
100k
1.7
0.01
0.5
0.00
1.6
0.01
1.3
0.01
DreamerV2 DA
100k
7.2
0.05
1.5
0.01
7.7
0.05
5.5
0.03
Dreamer 1
100k
1.9
0.01
0.9
0.00
2.8
0.01
1.9
0.01
DreamerV2- D DA
100k
8.3
0.05
2.1
0.01
10.8
0.07
7.0
0.04
WMC
100k
28.9
0.22
7.9
0.05
30.2
0.22
22.3
0.16
RAD
500k
48.7
0.44
11.6
0.11
48.5
0.44
36.3
0.32
CURL
500k
40.8
0.36
11.4
0.09
41.9
0.36
31,4
0.27
DreamerV2
500k
1.3
0.01
0.8
0.00
2.2
0.01
1.4
0.01
DreamerV2 DA
500k
7.2
0.04
1.1
0.01
9.7
0.05
13.0
0.08
Dreamer V2 D
500k
3.3
0.02
0.9
0.00
4.1
0.02
2.7
0.01
DreamerV2 I+ D DA
500k
38.0
0.25
9.0
0.06
52.7
0.35
33.2
0.22
WMC
500k
58.6
0.44
15.7
0.11
67.4
0.51
47.2
0.36
Table 2: Sribson-to-Gibsod dataset: sim-to-real perception transfer results on navigation task.
Ihlen
Muleshoe
Uvalda
Noxapater
McDade
Models
Steps
SR
SPL
SR
SPL
SR
SPL
SR
SPL
SR
SPL
RAD
100k
0.0
0.00
0.0
0.00
0.0
0.00
0.0
0.00
0.0
0.00
CURL
100k
5.9
0.05
3.8
0.03
5.1
0.04
5.9
0.05
12.8
0.11
WMC
100k
24.4
0.18
20.4
0.15
24.3
0.18
27.3
0.21
40.9
0.31
RAD
500k
26.4
0.23
27.5
0.24
28.5
0.25
28.6
0.25
40.0
0.34
CURL
500k
36.8
0.33
29.3
0.27
33.7
0.30
35.2
0.32
53.8
0.50
WMC
500k
50.0
0.38
50.3
0.38
49.7
0.37
45.5
0.34
50.7
0.38
EXPERIMENTS
4.1 EVALUATION
We evaluate the out-of-distribution (OoD) generalization, sim-to-real transfer of perception learning and sample-efficiency of our model and basclines at 100k and 500k environment steps. Sample efficiency on 100k and 500k steps is common practice (Laskin al.) 2020a b Yarats al., 2021 Hafner et al. 2021). Following (Hafner et al., (021), we update the model parameters on every fifth interctivv step. We used default hyperparameter values of DreamerV2 for our experiments. Similarly, we used official code for RAD and CURL experiments.
4.2 IGIBSON DATASET
We have tested our proposed WMC on random PointGoal task from iGibson 1.0 environment (Shen et al., 2021) for OoD generalization. It contains 15 floor scenes with 108 rooms. The scenes replicas of real-world homes with artist designed textures and materials. We have used RGB, depth and task related observation only. Depth is only used during the training phase. The task related observation includes goal location, current location, and linear and angular velocities of the robot. Action includes rotation in radians and forward distance meters for the Turtlebot. Since iGibson 1.0 does provide dataset splits for Ool) generalization, have chosen five scenes for

---

Under review as a conference paper at ICLR 2023
Table 3: Experiment results DMControl. Results reported averages across 10 seeds.
100k Steps Total Rewards
WMC
CURL
Dreamer
SAC+AE
PSAC
SSAC
Finger, spin
4861191
767-556
341=70
740=164
179:66
811=46
Cartpole, swingup Reacher, easy Cheetah, run
472-667
542--146
326-27
311+11
419+440
835+22
327-198
538--233
314+1555
274+14
145+30
746-25
321178
299 +48
235 137
267:24
197115
616-18
Walker, walk
654
4031:4
277=12
394122
42::12
891
Ball in cup, catch
830+:118
769+43
246+1174
391
312#:63
746-91
500K Steps Total Rewards
Finger, spin Cartpole, swingup Reacher, easy Cheetah, run Walker, walk
471
9261:45
196:1183
8844128
179=166
923=21
675164
841145
762-27
735163
419140
848=15
891#72
929+44
793+1164
627+:58
145+:30
923=24
633470
518+28
570+2533
550+334
197#15
795-30
96511
902:443
897=49
847148
42::12
948
Ball in cup, catch
950:22
9591227
8791 87
794 58
312 63
974
Table 4: Ablation study of the proposed WMC, the results clearly show the benefit of intervention- invariant auxiliary task, depth D, and action replay (AR).
Ihlen_0 Onint
Ihlen_1 iint
Rs_int
Env Avg
Models
Steps
SR
SPL
SR
SPL
SR
SPL
SR
SPL
WMC
100k
28.9
0.22
7.9
0.05
30.2
0.22
22.3
0.16
WMC AR
100k
15.4
0.11
4.2
0.02
19.7
0.12
0.08
WMC - -AR R
100k
0.0
0.0
0.0
0.0
0.0
0.0
0.03
0.00
WMC D+I
100k
15.4
0.11
4.6
0.03
17.0
0.12
12.3
0.09
WMC
500k
58.6
0.44
15.7
0.11
67.4
0.51
47.2
0.36
WMC AR
500k
45.1
0.28
11.5
0.07
26.8
0.18
27.8
0.17
WMC -R-
500k
0.9
0.0
0.2
0
1.2
0.01
0.7
0.00
WMC - DI
500k
28.6
0.12
6.6
0.04
22.3
0.14
19.1
0.10
training and tested on the held-out three scenes and visual textures both. The details are provided in the appendix.
We have trained all models three times with random seeds and report the average Success Rate (SR) and Success weighted by (normalized inverse) Path Length (SPL) on held-out scenes as well visual textures in the Table [ Our proposed WMC outperforms state-of-the-ar model-based RL method DreamerV2 and model-free method RAD and CURL on 100k and 500k interactive steps. Even though depth reconstruction and data augmentation improve the DreamerV2, proposed invariant causal features learning with contrastive loss further improves the results. All methods perform poorly the difficult Ihlen int scene. Further, our experimental results confirm that, similar to the model-free RL. methods (Laskin et al. 20206), data augmentation improves the performance of the model-based RL.
4.3 1GIBSON-TO-GIBSO DATASET
We use the Gibson dataset (Xaa et al. 2018) for sim-to-real transfer experiments of the perception module, representation learning module of the world model; however please note that the robot controller is still a part of the simulator. Gibson scenes are created by 3D scanning of the real scenes, and it uses neural network to fill the pathological geometric and occlusion errors only. We have trained all models on the artist created textures of iGibson and tested on five scenes from the Gibson. The results shown in Table 2 Our proposed WMC outperforms RAD and CURL
8

---

Under review as a conference paper at ICLR 2023
100k and 500k interactive steps, which shows that WMC learns more stable features and is better suited for sim-to-real transfer.
4.4 DEEPMIND CONTROL SUITE
The results for the DMControl suite (Tnyasuvunakool.c et al. 2020) experiments are shown in Table 3 We replaced the depth reconstruction of WMC with the original RGB input reconstruction to adopt the DMControl suite. WMC achieved competitive results, and the key findings are: i) even though depth reconstruction is an important component to enforce the invariant causal features learn- ing on WMC explicitly, the competitive results, even without depth reconstruction show the wider applicability of the proposed model; ii) WMC is competitive with CURL, the closest state-of-the- RL method with contrastive learning.
4.5 ABLATION STUDY
In many real world environments, action replay with texture randomization is not possible to per- form. Hence, we have also experimented WMC without action replay (AR) feature. The results are shown in Table These experiments confirm that WMC learning with only standard self-supervvsed learning is better than just data augmentation. However, in 500k steps measure data augmentation technique on DreamerV2 with depth reconstruction is doing better in Rs_int environment. This result suggests that model optimization with an additional constraint makes the optimization task harder. However, collecting more interactive data in the real environment is difficult in many see- narios. Since given the simulation to real gap is reducing every year (Shen et al. 2021) and WMC is performing better in 5 out of 6 results, additional contrastive loss still important for the model- based RL. Further, this result warrants that we need better method for intervention on image style or spurious variables. Since the separation of those variables from the pixel-observation is not an easy task, we do need some new contributions here, which we leave as future work.
The standard formulation of contrastive learning does not use reconstruction loss (Chen al. 2020
Grill et al. 2020 Laskin et al. 2020aa] He et al. (2020). Since model-based RL does not optimize the representation learning and controller jointly, contrastive loss collapses. Hence, to validate our proposal of style intervention at observation space and intervention invariant depth reconstruction as an auxiliary task, we have done experiments without depth reconstruction and action replay (WMC
AR - D). The results are shown in Table 4 The degradation in performance without depth recon- struction loss much worse than only without action replay. In summary, in a reasonably complex exxel-bsedd control task, WMC is not able to learn meaningful control. This shows the value of our careful design of the WMC using contrastive loss and depth reconstruction as an auxiliary task.
Further, WMC with RGB image reconstruction rather than depth (WMC- D + D) reduces the success rate by 50% and SPL by almost two third. These results confirm that our proposal of doing an intervention RGB observation space and adding intervention invariant reconstruction of depth as an auxiliary task is one of the key contributions of our paper.
5 CONCLUSION
In this work we proposed method to learn World Model with invariant Causal features (WMC). These invariant causal features are learnt by minimizing contrastive loss between content invariance interventions of the observation. Since the world model learns the representation learning and policy of the agent independently, without providing the better supervisory signal for the representation learning module, the contrastive loss collapses. Hence, proposed depth reconstruction auxiliary task, which is invariant to the proposed data augmentation techniques. Further, given an intervention in the observation space, WMC can extract well as predict the related causal features.
an
Our proposed WMC significantly outperforms the state-of -fe---ffthhhheeee the-art models on out-of -iitt--iiittitttionn gen- eralization, sim-to-real transfer of perception module and sample efficiency measures. Further, our method works on a sample-level intervention and does not need data from different environments to learn the invariant causal features.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, volume 32, 2019.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization, 2020.
Kenneth Bollen. Structural equations with latent variables, volume 210. John Wiley & Sons, 1989.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning. 2020.
S. Chopra, R. Hadsell, and Y. LeCun. Learning a similarity metric discriminatively, with application to face verification. In IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2005.
Djork-Ame Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (clus). arXiv preprint arXiv:151 11122/1....79., 2015.
A. Dawid. Causal inference without counterfactuals. Journal of the American Statistical Associ- ation, 95(450):407-422 2000.
A. Philip Dawid and Vanessa Didelez. Identifying the consequences of dynamic treatment strategies:
A decision-theorett overview. Statistics Surveys, 4:184 231, 2010.
Mare Deisenroth and Carl Rasmussen. Pilco: model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML), pp. 465-472, 2011.
Alexey Dosovitskiy, Jost Tobias Springenberg, Martin Riedmiller, and Thomas Brox. Discrimina- tive unsupervised feature learning with convolutional neural networks. In Advances in Neural Information Processing Systems, 2014.
Lean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, koray kavukcuoglu, Remi Munos, and Michal Valko. Bootstrap your own latent approach to sule-supe learning. In Advances in Neural Information Processing Systems, 2020.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2021.
Jessica B Hamrick. Analogues of mental simulation and imagination in deep learning, Current Opinion in Behavioral Sciences, 29, 2019.
J. A. Hartigan and M. A. Wong. Algorithm as 136: k-means clustering algorithm. Journal of the Royal Statistical Society. Series C (Applied Statistics), 1801):000-108, 1979. ISSN 00359254, 14679876.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xic, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1442,6988, 2014.
10

---

Under review as a conference paper at ICLR 2023
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. 2014.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. Proceedings of the 37th International Conference on Machine Learning, Vienna, Austria, PMLR 1 11, 2020a. arXiv:0004.04336.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. Advances in Neural Information Processing Systems, 2020b.
L. Lauritzen and D. J. Spiegelhalter. Local computations with probabilities on graphical struc- tures and their application to expert systems. Journal of the Royal Statistical Society: Series B (Methodological), (00 1988.
Jovana Mitrovic, Brian McWillias, Jacob Walker, Lars Holger Buesing, and Charles Blundell. Representation learning via invariant causal mechanisms, In International Conference on Learn- ing Representations, 2021.
Judea Pearl. Causality: Models, Reasoning, and Inference. Cambridge University Press, 2 edition, 2009.
Jonas Peters, Peter Büihlmann, and Nicolai Meinshausen. Causal inference by using invariant pre- diction: identification and confidence intervals. Journal of the Royal Statistical Society: Series (Statistical Methodology), 78(5):947-1012, 2016.
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 10965)x612-33, 2021.
Bokui Shen, Fei Xia, Chengshu Li, Roberto Martin-Martin, Linxi Fan, Guanzhi Wang. Clau- dia Pérez-D' Arpino, Shyamal Buch, Sanjana Srivastava, Lyne Tchapmi, Micael Tchapmi, Kent Vainio, Josiah Wong. Li Fei-Fei, and Silvio Savarese. igibson 1.0: A simulation environment for interactive tasks in large realistic scenes. In IEEE International Conference on Intelligent Robots and Systems (IROS), 2021.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Si- monyan, and Demis Hassabis. A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play. Science, 362(6419):11400)
2018. ISSN 0036-8075.
Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approxi- mating dynamic programming. Machine learning proceedings 1990, pp. 216-224. Elsevier, 1990.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for continuous control. Software Impacts, 6, 2020.
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv: 180700787 2018.
J J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3):229-22) 256, 1992.
Fei Xia, Amir R. Zamir, Zhi-Yang He, Alexander Sax, Jitendra Malik, and Silvio Savarese. Gibson env: real-world perception for embodied agents. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2018.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representa- tions, 2021.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block MDPs. In Proceedings of the 37th International Conference Machine Learning, 2020.

---

Under review as a conference paper at ICLR 2023
A
APPENDIX
A.1 QUALITATIVE RESULTS OF WMC
D
DG
Figure 3: The out-of- of-iistribttion generalization tests of proposed WMC on held-out scenes and visual textures from iGibson 1.0 environment. Green circle random Goin blue circle is random starting point and blue line represents the travel path of the Turtlebot robot.
12

---

Under review as a conference paper at ICLR 2023
A.2
HYPER PARAMETERS
Table Hyper parameters of proposed WMC.
Name
Symbol
Value
World Model
Dataset size (FIFO)
3 10a
iGibson input image size
B
O
120 160
Batch size
50
Sequence length
L
50
Discrete latent dimensions
32
Discrete latent classes
32
RSSM number of units
1024
KL loss scale
3
1.0
World model learning rate
3 10
Key encoder exponential moving average
0.999
Behavior
Imagination horizon
H
15 1 10-4
Actor learning rate
Critic learning rate
1
Slow critic update interval
100
Common
Policy steps per gradient step
4 4 400
Policy and reward MPL number of layers Policy and reward MPL number of units Gradient clipping
100
Adam epsilon
10-5
Encoder and Decoder
MLP encoder sizes of task obs Encoder kernels sizes
32, 32
4, 4, 4, 4,
Decoder kernels sizes Encoder and decoder feature maps Encoder and decoder strides
5, 5, 4, 5,
32, 64, 128, 256, 512
2, 2, 2, 2
Decoder padding
none, 0-1, none, none, none
Data Augmentation
Padding range
10
Hue delta
0.1
Brightness delta
0.4
Contrast delta
0.4
Saturation delta
0.2
Gaussina blur sigma min, max
0.1,
Cutout min, max
30,
13

---

Under review as a conference paper at ICLR 2023
A.3 IGIBSON 1.0 TRAINING AND EVALUATION SPLITS
Table Train-test scenes splits for iGibsion 1.0 dataset (Shen et al. 2021).
Phase
Scene names
Training
Beechwood jint, Beechwood
Benevolence.0 1a______nn iint, Benevolence. Lint, Benevolence.2.in
Merom_O_in Liint, Merom_l_int,
Pomaria 0 0_int, Pomaria. _int, Pomaria 2.int,
Wainscott Jainscott_O.int, 0_int, Wainscott.
Testing
Ihlen 0 int, Ihlen 1 int, Rs int
Table 7: iGibson 1.0 environment (Shen et al.
2021 held out texture ids for test.
Material category
Held-out texture ids for test
asphalt
06, 15
bricks
08, 19
concrete
06, 15, 17
fabric
01, 02, 28
fabric_carpet
02, 05, 13
ground
13, 19
leather
03, 12
marble
02, 03
metal
10, 19
metal siriulddamodd.ppp. ppate
04
moss
01, 03
paint
05
paving stones
24, 38
planks
07, 09, 16
plaster
03
plastic
04, 05
porcelain
02, 04
rocks
04
terrazzo
06, 08
tiles
43, 49
wood
02, 05, 16, 22,
wood_floor
06. 10, 17,28
14

---

Under review as a conference paper at ICLR 2023
Table Examples of textures split for training and testing from the proposed split in Table 7
Train
concrete
Test
Train
fabric
Test
Train
planks
Test
Train
wood
Test
15


Reviews:
[{"review_id": "uHXcO5RuRQ1", "reviewer": "Reviewer_tiQV", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper provides a novel way towards achieving better representation learning in world models, which is essential to robust policy learning, by unsupervised causal representation learning. The authors of the paper propose depth prediction and data augmentation techniques to reach this goal.  Their model is called World Model with invariant Causal features (WMC). Some of the main contributions and findings of the paper include: 1) Depth reconstruction to achieve better representations. 2) Data augmentation on RGB image input as intervention for contrastive learning. 3) They allow training in a single environment leading to sample efficiency, 4) They also allow tackling training bias in the model. 5) Good results in out-of-distribution generalization in iGibson 1.0 (navigation task) and sim-to-real transfer in iGibson-to-Gibson. 6) Wider applicability of proposed model without depth reconstruction is showed by the DMControl suite experiments\n", "summary_of_the_review": "The paper is well-written and presents a novel approach to representation learning problems in world models with extensive experimentation results. Addressing the highlighted minor weaknesses would make the paper stronger. I recommend acceptance.", "strength_and_weaknesses": "Strengths\n- Well-written and organized\n- Related work section gives a good comparison between the proposed model and previous work\n- Performed experiments for benchmarking are extensive and suitable\n- Ablation study is well-designed and confirms the proposed methods\n\nWeaknesses\n- Small typos in the text (see questions and suggestions)\n- Content of the tables is not always easy to follow (see questions and suggestions)\n- Hyperparameters are the same in all experiments\n\nQuestions and suggestions for the authors (section by section): \nRelated work\n\nQ1: In paragraph Sample efficiency: Not the correct paper of the Dreamer method is referenced by the authors when they write about the Dreamer method by Hafner et al. that was beaten by the CURL technique. The referenced paper didn’t even exist that time.\n\nTheir reference: \nDanijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with\ndiscrete world models. In International Conference on Learning Representations, 2021.\n\nCorrect reference:\nDanijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning\nbehaviors by latent imagination. In International Conference on Learning Representations, 2020.\n\n+ the style of the title “Sample efficiency” could be bold and not italic to match the other paragraph titles\n\nProposed model\nQ2: In Section 3.2 World Model you describe briefly the controller in the first paragraph: “The controller maximizes the action probability using an actor critic approach”. I found this sentence confusing because these methods aim to maximize the expected reward.\n\nExperiments section\nQ3: The meaning of the letters (I, D, AR) should be clearly presented in the first experiments to be easier to follow, they can be only known from the Ablation study, which is at the very end of the section.\n\nQ4: Table 1 and Table 3 (probably typos):  I believe it is a typo in Table 1.: the last but one method in the 500k steps is “Dreamer - I + D + DA”, all the others are DreamerV2. Is it DreamerV2 too? In Table 3. the method called “Dreamer” is the same DreamerV2 as in the other experiments or the earlier Dreamer version of Hafner et al.? + The headers could be more consistent in Table 3.: “100k Steps” and “500K step scores”.\nQ5: Table 3: Could you describe what are exactly the presented results in these environments? Are they the obtained rewards?\nQ6: Table 1, 2 and 4: It would be good to present the confidence of these models by reporting the standard deviations of SR and SPL.\nQ7: Table 1 and Table 4: What is the meaning of the column “Total”?\n\nQ8: The exact hyperparameters are usually very task specific. What is the reason for using the same hyperparameters in these diverse environments (iGibson, DMControl)? Have you done any hyperparameter tuning?\n\nQ9: Ablation study: The authors state that their WMC model performs better in 5 out of 6 environments. However, it is 3/6 in the case of 100k steps and 2/6 in 500k steps.\n", "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": {}}, {"review_id": "SdO3KZsf9Kp", "reviewer": "Reviewer_acab", "summary": "", "questions": "", "limitations": "", "rating": 1, "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": "", "summary_of_the_paper": "This paper adds depth prediction (instead of RGB image prediction) and a contrastive learning objective (with data augmentations) to the world model of DreamerV2 and shows that this results in better performance on the iGibson dataset. For contrastive learning, they use InfoNCE objective and create positive pairs using data augmentations (like crop, blur etc.). One important type of augmentation here is texture randomization. They also remove the image reconstruction loss in DreamerV2 world model training objective and add depth prediction instead. Note this requires depth information to be available in training environment. They compare their technique to plain DreamerV2 and other constrastive learning based methods for RL such as RAD and CURL. They show their technique performs better on iGibson dataset and can transfer better to real textures when trained on synthetic ones (one variant of sim-to-real problem).\n", "summary_of_the_review": "Overall the novelty is pretty limited and there are concerns regarding the empirical evaluation, so I don't think the paper is ready for acceptance in its current form unfortunately.", "strength_and_weaknesses": "Unfortunately, I can't find many strengths of this paper. It makes minor modifications to DreamerV2 and shows that these help in a very particular environment.\n\nFirst, the paper doesn't provide a good motivation for removing image prediction and using depth prediction instead. Given that depth is not generally available, this limits the applicability of the approach significantly. If the argument is that depth prediction in general is better than image prediction, then the evaluations should have made that case (by evaluating model trained with depth in more environments and different settings).\n\nAdding a contrastive learning objective that makes use of data augmentations is not really novel. There are many papers that does this and show it helps. It is unfortunately not clear what the novelty of this paper is wrt to contrastive learning in RL.\n\nThe quantitative results on iGibson look promising but it is unclear whether the other contrastive learning techniques (like RAD or CURL) use the same data augmentations. If not, this wouldn't be a fair comparison. Relatedly, is texture randomization used here? If so, is it used by all contrastive learning based techniques?\n\nUnfortunately, the writing needs to be improved as well. Some of the text was difficult to understand (some examples below), and some of the arguments didn't seem to follow. For example, the authors point out (in a couple of places) that their technique learns invariant causal features of the environment. However, this really depends crucially on what interventions are available; if the interventions you have do not cover all of the spurious variables then contrastive learning over data augmentations cannot learn causal features (it will learn a combination of causal and spurious). And in fact there is really nothing special about the technique that allows it to learn these causal features; it is the data augmentations (interventions), which are provided by us (the humans) that allows learning these causal features.\n\nOther points:\n\n- Do the augmentations change all the spurious variables?\n- The authors texture randomization \"action replay\". This seemed rather confusing to me. How is action replay and texture randomization the same thing?\n- Do other techniques use the same augmentations? Especially texture randomization?\n- For sim-to-real results, do all techniques use texture randomization? Otherwise the comparison is not fair.\n- It'd be nice to have std devs with quantitative results.\n- Please improve table captions. For example, table 4 caption should mention what abbreviations mean?\n- There are other self-supervised approaches for RL slike SpR (Data-Efficient Reinforcement Learning with Self-Predictive Representations). It'd be nice to compare to them.\n\nTypos etc.\npg 3. Self-supervised learning formulates the learning as a supervised loss function. Confusing sentence. Perhaps \"as the optimization of a supervised loss function\" maybe?\npg 4. variables involves are high -> involved\npg 5. also know as -> known\npg 6. which is optimizes -> optimized\npg 6. The actor critic approach considers the rewards beyond the horizons. Confusing sentence. Many RL techniques consider rewards beyond the horizon.\n", "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": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "KulCvNYst6", "reviewer": "Reviewer_Hdwq", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "", "summary_of_the_paper": "Thıs paper presents an approach for learning world models with causal features. The authors propose to augment well-established model-based reinforcement learning (MBRL) objectives with a contrastive unsupervised learning objective, which induces causal features and invariances. Unlike existing approaches, the model is expected to predict depth as a side task, which is shown to improve the model performance. The model outperforms the SOTA MBRL method (Dreamer on out-of-distribution navigation tasks. It is also shown to achieve best results on a sim-to-real navigation task.", "summary_of_the_review": "Overall, the idea of improving MBRL with representation learning techniques is a nice idea. Authors opt for perhaps the simplest approach (augmenting the MBRL objective with a contrastive learning loss), which restricts the originality of the work. I also believe the writing requires major updates, especially the grammar and (lacking) references.", "strength_and_weaknesses": "Strengths:\n- The results tables are impressive. Given a fixed compute budget (e.g., environment steps), the method outperforms SOTA or performs on par. \n\nWeaknesses:\n- Lack of originality: The method is very largely built upon Dreamer - a SOTA MBRL method. The major improvements over Dreamer are (i) the contrastive objective and (ii) the depth prediction task. The former seems to be a very straightforward application of Mitrovic'21 - so much so that Mitrovic'21 presents a very similar version of Figure-2. I also believe the depth prediction task is a rather practical aspect of the work, which significantly restricts the dataset characteristics. As such, the overall contribution is a combination of well-known optimization targets, which I believe lacks novelty in order to be published in this venue.\n- The writing should be significantly improved. Both grammar and text structure requires modifications. Here are some tips:\n  - Contribution and motivation are mixed up in the abstract. The first two sentences state the contribution, then the motivation, and then we read the contribution again.\n  - \"Specifically, we utilize depth prediction to explicitly enforce the invariance and use data augmentation as style intervention on the RGB observation space.\" Are these two parts of the sentence related? Otherwise I suggest having two different sentences.\n  - \"Our design leverages unsupervised representation learning to learn the world model with invariant causal features\" was already stated in the abstract.\n  - The first paragraph is a discussion on model-free vs model-based RL, which is not directly related to the rest of the paper.\n  - \"... as the control policy is learned end-to-end.\" What does end-to-end refer to here? Perhaps learning from pixels?\n  - \"suited to separate relevant states from irrelevant ones\" Relevant to what? Policy learning?\n  - \"so inference of the underlining causal structure\" Underlying, perhaps?\n  - Short explanations for the components in Figure-1 would be nice to include in the caption.\n  - \"Zhang et al. (2020) uses multiple environments to learn the invariant causal features using a common encoder. Here spurious or irrelevant features are learnt using environment specific encoders.\" This reads like related work instead of intro, making it more difficult to follow this long paragraph.\n  - \"Model-based RL methods do not optimize feature learning\" This is vague: What does it mean to \"optimize feature learning\"?\n  - \"The aim is that such features will be\" Which features?\n  - \"Hence, we propose a causally invariant auxiliary task for invariant causal features learning.\" What is the connection with the previous part (why hence?)\n  - \"Further, most computer vision data augmentations on RGB image space are invariant to depth.\" Why is this important?\n  - \"Our main contributions are to show ... to propose ...\"\n  - \"Representation learning methods based on contrastive loss (Chopra et al., 2005) have recently achieved state-of-the-art performances\" on what tasks?\n  - \"Learning structured representations that\" capture instead of captures\n  - \"Furthermore, it is well known that causal variables have an invariance property\" This requires elaboration as it might not be familiar to some readers.\n  - \"Sample Efficiency\" should be in bold.\n  - \"model-based RL **method** called\"\n  - \"Further, these papers do not consider the effect of data augmentation\" whereas both RAD and CURL utilize augmentations?\n  - What is \"a\" in Figure 2?\n  - \"the number of causal variables **that involve**\"\n  - \"That is why we choose invariant prediction as our method of choice for causal feature learning.\" Perhaps say \"That is why we choose invariant prediction for causal feature learning\"?\n  - \"Such experimental changes are known as interventions\" Isn't it a bit late to describe interventions? Maybe explain it in the intro?\n  - \"since we do not have access to the causal or spurious variables or both of the environment.\" Which environments?\n  - \"..., which we call action replay.\" Why this name?\n  - \"The main idea is observation is made of content (C), causal variables, and style (S), spurious variables\" This is unclear: aren't C and S latent factors? \n  - \"which means embedding features of a sample and its data augmented version, intervention on style or spurious variables, should be the same\" This is unclear.\n  - \"the data also can come from observation of the different environments\" Unclear, which environments?\n  - \"with different environment-level interventions\" Why would that be the case? Would environment-specific interventions be a by-product of having multiple environments or because of our interventions?\n  - \"memory also know as world model\" Unclear, why is memory known as world model?\n  - \"Memory module uses **a** recurrent neural network\"\n  - \"parameters of a categorical distribution, discrete and multinomial distribution,\" Discrete distribution?\n  - \"The controller maximizes the action probability\" Perhaps learns over maximizes?\n  - \"our early result with only rewards prediction was poor\" What early results? \n  - \"and reward prediction as target.\" Perhaps \"rewards as predicted targets\"?\n  - The list of used data augmentations is mentioned twice.\n  - \"transforms in two correlated views\"\n  - \"that extracts representation **s** from\"\n  - \"We use the contrastive loss immediately after the encoder\" How do you use it?\n  - \"Given a query observation q and a set K = {k0, k1, ...k2B} with known positive {k+} and negative {k−} keys.\" The sentence lacks a verb.\n  - \"Where B is a batch size, which becomes...\"\n\n\n\n- Reference is needed to support many claims in the paper:\n  -  An important reference to one of the first model-based RL papers is missing: Deisenroth, Marc, and Carl E. Rasmussen. \"PILCO: A model-based and data-efficient approach to policy search.\" Proceedings of the 28th International Conference on machine learning (ICML-11). 2011.\n  -  \"model-free methods learn a greedy state representation\" One would argue that model-free approaches do not even learn a state representation, e.g., if trained on noisy state observations. Giving a reference to support this claim would be better.\n  - \"Therefore, we consider model-based methods more attractive for continuous learning, out-of-distribution (OoD) generalization and sim-to-real transfer\" Reference needed to support the claim.\n  - \"resulting in performance degradation of the control policy, even for slight task-irrelevant style changes.\" \n  - \"so inference of the underlining causal structure of the model becomes intractable.\"\n  - \"Recently, VAE (Kingma & Welling, 2014) has been a preferred approach for representation learning in model-based RL\"\n  - Almost the whole MBRL literature is left out of \"Model-based RL\" paragraph. I suggest authors explain their scope here.\n\n", "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": "IYvYdw2lE0U", "reviewer": "Reviewer_12bZ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "", "summary_of_the_paper": "In this paper, the authors essentially learn a DreamerV2 model, where the default reconstruction loss has been replaced by a contrastive loss and an additional auxiliary loss (which is depth reconstruction for datasets with such information available, and reconstruction once again for those without).  The contrastive loss is learned in a similar way to the setup in SimCLR, where two stylistic augmentations of the same frame are encoded and treated as positives (and other frames are treated as negatives).  The main claim is that what is learned from the outcome of this contrastive loss are the causal features, where augmentation interventions would not affect the result.", "summary_of_the_review": "The overall paper was written clearly, and quite easy to follow.  The authors also demonstrate good results on certain tasks, such as sim-to-real.  However, I believe that method-wise, the benefits are limited with the continued use of some form of reconstruction.  Furthermore the utilization of unsupervised and contrastive losses for world models, which appears to be the main contribution in this work, has been previously explored.  Overall, I recommend that this work revisit its novelty and strengthen its contributions before being considered for acceptance.", "strength_and_weaknesses": "A strength of this paper is its sim-to-real perception task performance on iGibson-to-Gibson, which seems to greatly outperform prior model-free results.\n\nWeaknesses of this paper include its novelty (which will be elucidated upon in the following section).  Method-wise, many benefits arise from unsupervised contrastive learning (not requiring labels, focusing on the features that change rather than the ones that remain constant throughout scenes).  However, these benefits seem to disappear when the authors reintroduce reconstruction, whether it be depth reconstruction when such information is available, or regular RGB reconstruction when it is not (e.g. DMControl).  Rather than terming these reconstruction objectives as auxiliary objectives, it seems to me that in fact the contrastive objective is the auxiliary objective added to a default DreamerV2 (at least in the DMControl task).", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject