Robust Robotic Control from Pixels using Contrastive Recurrent State-Space Models
Authors: Nitish Srivastava
Walter Talbott
Martin Bertran Lopez
Shuangfei Zhai
Joshua M. Susskind
Venue: ICLR 2022 Submitted (2022)
Keywords: contrastive learning
model-based RL
distractions
predictive coding

Abstract: Modeling the world can benefit robot learning by providing a rich training signal for shaping an agent's latent state space. However, learning world models in unconstrained environments over high-dimensional observation spaces such as images is challenging. One source of difficulty is the presence of irrelevant but hard-to-model background distractions, and unimportant visual details of task-relevant entities. We address this issue by learning a recurrent latent dynamics model which contrastively predicts the next observation. This simple model leads to surprisingly robust robotic control even with simultaneous camera, background, and color distractions. We outperform alternatives such as bisimulation methods which impose state-similarity measures derived from divergence in future reward or future optimal actions. We obtain state-of-the-art results on the Distracting Control Suite, a challenging benchmark for pixel-based robotic control.

Under review as a conference paper at ICLR 2022
ROBUST ROBOTIC CONTROL FROM PIXELS USING CONTRASTIVE RECURRENT STATE-SPACE MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
Modeling the world can benefit robot learning by providing a rich training signal for shaping an agent's latent state space. However, learning world models in unconstrained environments over high-dimensional observation spaces such as images is challenging. One source of difficulty is the presence of irrelevant but mode- background distractions, and unimportant visual details of task- relevant entities. We address this issue by learning recurrent latent dynamics model which contrastively predicts the next observation. This simple model leads to surprisingly robust robotic control even with simultaneous camera, background, and color distractions. We outperform alternatives such as bisimulation methods which impose state-similarity measures derived from divergence in future reward or future optimal actions. We obtain state-of-the-an results on the Distracting Control Suite, a challenging benchmark for exee-based robotic control.
INTRODUCTION
For robot, predicting the future conditioned on its actions is a rich source of information about itself and the world that it lives in. The gradient-rich training signal from future prediction can be used shape robot's internal latent representation of the world. These models can be used generate imagined interactions, facilitate model preddictve control, and dynamically attend to mispredicted events. Video prediction models (Finn al., 2016) have been shown be effective for planning robot actions. Model-based Reinforcement I Learning (MBRL. methods such as Dreamer (Hafner et al. 2020) and SLAAC (Lee et al. 2020) have been shown to reduce the sample complexity of learning control tasks from basel observations.
These methods learn an cendii-conddtioned model of the observations received by the robot. This works well when observations are clean, backgrounds are stationary, and only task-relevant in- formation contained the observations. However, in a real-world unconstrained environment, observations are likely to include extrancous information that is irrelevant to the task (see Figure ITap. Extrancous information can include independently moving objects, changes in lighting, colors, as well as changes due camera motion. High-frequeney details, even in task-relevant entities, such as the texture of a door or the shininess of a door handle that a robot is trying to open are irrelevant. Modeling the observations at pixel-level requires spending capacity explain all of these attributes. which is inefficient.
The prevalence of extraneous information poses immediate challenge to reconstruction based methods, e.g. Dreamer. One promising approach is to impose a metric on the latent state space which groups together states that are indistinguishable with respect to future reward sequences (DBC, Zhang et al. (202T)) or future action sequences (PSE, Agarwal et al. (202I)) when running the optimal policy. However, using rewards alone for grounding is sample inefficient, especially in tasks with sparse rewards. Similarly, actions may not provide enough training signal, especially if they are low-dimensional. On the other hand, observations are usually high-dimensional and carry more bits of information. One way to make use of this signal while avoiding pixel-level reconstruction is to maximize the mutual information between the observation and its learned representation (Hjelm) et al. 2019). This objective can be approximated using contrastive learning models where the task is to predict a learned representation which matches the encoding of the true future observation (positive pair), but does not match the encoding of other observations (negative pairs). However, the

---

Under review as a conference paper at ICLR 2022
RSSM
Bahavior Modal
Inverse Dynamics
CC
stg (+)
To nest time-step
6,1
Observation Encaner
(a) Environment Examples
(b) Recurrent Contrastive State-Space Model
Figure I: Left: Examples visual control tasks in canonical (top) and distracting (bottom) environ- ments. Background, color, and camera pose distractions significantly increase observation complexity. Right Our proposed Contrastive Recurrent State-Space Model (CoRe).
performance of contrastive variants Dreamer and DBC has been shown to be inferior, for example, Fig. 8 in Hafner et al. (2020) and Fig. 3 in Zhang et al. (2021).
In this work, we show that contrastive learning can in fact lead surprisingly strong robustness to severe distractions, provided that recurrent state-space model is used. We call our model CoRe: Contrastive Recurrent State-Space Model. The key intuition is that recurrent models such as GRUs and LSTMs maintain temporal smoothness in their hidden state because they propagate their state using gating and incremental modifications. When used along with contrastive learning, this smoothness ensures the presence of informative hard negatives in the same mini-bateh of training. Using CoRe, we get saat--ff-theaa results on the challenging Distracting Control Suite benchmark (Stone al., 2222)) which includes background, camera, and color distractions applied simultaneously. Video results from our model can be seen at supplementaryyindex.hti
2 MODEL DESCRIPTION
We formulate the robotics control problem from visual observations as discrete-time partially observable Markov decision process (POMDP). At any time step t, the robot agent has an internal state denoted by zz receives an observation O,, takes action a,, and moves to state 2(41, receiving reward r2+11 We present
model with continuous-valued states, actions, and observations. straightforward extension can be made to discrete states. The objective is to find policy that maximizes the expected sum of future rewards Ei r'r,, where 7 is discount factor.
2.1 MODEL COMPONENTS
As shown in Figure 1b our model consists of three components: a recurrent state space model (RSSM) (Hafner et al. 2019) which encodes the learned dynamics, a behavior model based on Soft-Actor- Actor-Criitc (SAC) [Haamoja et al. 2018) which is used to control the robot, and an observation encoder.
Observation encoder: This component is convolutional network / which takes observation as input and outputs encoding a4 f(o:). LayerNorm (Ba et 2016) applied in the last layer of the encoder.
Recurrent State Space Model (RSSM): This component consists of the following modules,
Future predictor
p(z, zt- a;-1)),
Observation representation decoder
y: g(zt),
Correction Model
q( #t))
Inverse Dynamics model
a: a(zt, Zz-1);
Reward predictor
F. r(z)).
Each module is parameterized as fully-connected MLP, except the future predictor which uses GRU cells (Cho et al., 2014). The future predictor outputs a prior distribution 2, over the next state, given the previous state and action. This is decoded into yi which corresponds the observation that the agent expects to see next. The correction model is given access to the current observation encoding æ,, along with the prior E, and outputs the posterior distribution z,,
encoding of
2

---

Under review as a conference paper at ICLR 2022
Le., it corrects the prior belief given the new observation. The posterior state 2, is used to predict the reward, and also feeds into the actor and critic models that control the agent's 's behavior. The inverse dynamics model predicts the action that took the agent from state 2;-1 to 22. The RSSM is similar to the one used Dreamer and SLLAC except that it includes an inverse dynamics model and that the observation representation is predicted from the prior latent state instead of the posterior latent state. These differences lead to slightly better results and different probabilistic interpretation of the underlying model which we discuss in Section 2.2 Similar Dreamer, the latent state 2n consists of a deterministic and stochastic component (More details in Appendix A).
Behavior Model: This component consists of an actor ( (1) and two critic networks Qo. (2 a). i {1.2} which are standard for training using the SAC algorithm (Haarnoja et al. 2018
2.2 REPRESENTATION LEARNING WITH CORE
Our goal is to represent the agent's latent state by extracting task-relevant information from the observations, while ignoring the distractions. We formulate this representation learning problem as dynamics- model's 's prediction of the observation's 's encoding ye- Mutual information maximization (Hjelm et al. 2019) has been used extensively for representation learning. In our case, it can be approximated using contrastive learning objective defined over mini-batch of sequences, in way similar to Contrastive Predictive Coding (CPC) (yan den Oord et al.) 2018). At time-step f in example i in the mini-batch, let æi, and yi denote the real observation's ' encoding and the predicted observation encoding, respectively. The loss function is
mutual information maximization between the observations Or and the
exp(XarT,ey..)
exp(Xar,,,
= -
--Elos Log Eveer exp(Xar,,yvie)
Σ log
Ev.or exp( exxp(aar,.... 1r
where À is learned inverse temperature parameter and i' and 1' index over all sequences and time- steps in the mini-batch, respectively. Intuitively, the loss function makes the predicted observation encoding match the corresponding real observation's encoding, but not match other real observation encodings, and vice-versa. In addition, the predicted and corrected latent state distributions should match, which can be formulated as minimizing Ki-divergence term
LKL = KL (g(2) &, xz, ) (p( (1))),)))-1))
(2)
An additional source of training signal is the reward prediction loss. (r, 2. Furthermore, modeling the inverse dynamics, i.e. predicting a, from 2, and 1[+1 is a different way of representing the transition dynamics. This does not provide any additional information because we already model forward dynamics. However, we found that it helps in making faster training progress. We incorporate this by adding an action prediction loss. The combined loss is
Ja (9) L. + lL + a, (r: - f2)2 + aalla; d.11?
where B, ar, Ga are tunable hyperparameters and θ is the combined set of model parameters which includes parameters of the observation encoder, the RSSM, and the inverse temperature X.
Relationship with Action-Conditioned Sequential VAEs The above objective function bears resem- blance to the BLOO from a Conditional Sequential VAE which is used the underlying probabilistic model in Dreamer and SLLAC The objective function there is the sum of an observation reconstruction term ||o, - o, 2| and LKL scaled by 8, corresponding to a 3-VAE formulation (Higgins et al. 2017). This formulation relies on decoding the posterior latent state to compute 6, Since the posterior latent state distribution q( 2(2, ((,,,,,,2) xx.) is conditioned on the true observation 2;, observation reconstruction is an auto-encoue encoding task. Only the KL-term corresponds to the future prediction task. On the other hand, in our case the prior latent state decoded to a representation of the observation, which already constitutes future prediction task. The KL-tterm only provides additional regularization. This makes the model less reliant on careful tuning (empirically validated in Appendix C.3). In fact, setting it to zero also leads to reasonable performance in experiments, though not the best.
2.3 BEHAVIOR LEARNING
Behavior is learned simultancously with the state representation using SAC. This involves learning an actor ne that parameterizes stochastic policy and two critics Qo, i (1,2) which use the latent
3

---

Under review as a conference paper at ICLR 2022
state 2, as the representation of the agent's state. The critic is trained by minimizing Bellman error, =
(0,; ④)
IF
[[Qu, (21, a+) - sg(ri + + VVrrges (2++1)))
(3)
Ze+1mp ma
Virger ((+1) ae 1
E
[min/=1.2 Q6, (Z++1, a(+1) a log no (a141
(4)
where 8, represents exponentially moving average of 0.. 'sg' denotes the stop- nop-gradient operation, and a is the weight on the entropy regularization. Note that Jq is a function of é as well. In other words, the RSSM and observation encoder model are also trained using the gradients from the critic loss. Appendix expelorrre this design choice in more detail. The parameters O are shared between the two critics. The actor updated using the following loss function
Jn
E [a log Te (aglzz) - m,n at)].
(5)
Zt 2t Se~
Overall training consists of three steps: collecting data using the current actor, using the data to update the representation model, and using the data to update the actor and critic as shown Algorithm |
EXPERIMENTS
We use the Distracting Control Suite (DCS) (Stone et al. 2021) and Robosuite (Zhu) al. 2020) benchmark our model's ability to withstand visual distractions. DCS consists of six simulated robotic control tasks derived from the DeepMind Control Suite (Tassa et al., 2018). There are three types of distraction: background, camera, and color. They can either be fixed for the duration of the episode (static) changed smoothly during the episode (dynamic). Three benchmarks are defined: easy, medium, and hard, which have increasing levels of difficulty along all 3 distraction types. We focus our experiments on two settings: crnamic-backgroune
Algorithm 1: The CoRe Model
Require: Environment E, initial θ é. 01, 02 for each iteration do
01 Fresa (), Za 0, + 0:
+ 0]
for each environment step do
Φe fa(or):
21 RSSM(Z-11, ac-1, 2t):
at (at
rt. Oc+1 + Espp(ar):
D DU (o, rt):
end for each gradient step do
B ~ N samples of length T from D;
Train the model: E3 JJrr (0)]
and dynamic- medium. Results on other settings follow the same trends (Appendix B). In the dynamic- background-only case, the backgrounds for end
Train the critic: Eg [o (0., θ)]
Train the actor: Eg Ja (o)]:
end
training and validation environments are drawn from the 60 training and 30 validation videos in the DAVIS dataset (Pont- (Pont-Tuset et al. 2017). Our experiments show that (1) CoRe enables visual robotic control in the presence of severe distractions, performing better than strong baselines, (2) all three elements of CoRe recurrence, contrastive learning, and dynamics modeling are important, (3) recurrent architectures
important to make contrastive learning work well, (4) CoRe can solve a challenging simulated robotic manipulation task, and (5) interpretable attention masks emerge that help visualize how CoRe works. We report the mean validation reward over 10 random seeds and 100 episodes per seed along with the standard error. Details of the network architectures and training
are provided in appendix A
3.1 COMPARISON OF ROBUSTNESS TO DISTRACTIONS
This set of experiments quantifies the robustness of CoRe to visual distractions. We compare to a model-free RL baseline SAC-RAD, a state-of-the-ar method based on bisimilarity metrics PSE (Agarwal et al., 2021). Recon, reconstruction variant of our model. where we replace the con- trastive loss with an observation reconstruction loss (making similar to SLAC), and CURL [Laskin et al. 2020a) which also uses contrastive learning but does not recurrence or future prediction. Figure 2 the comparison as a function of environment steps for of the 6 environments. We use the published code for PSE and CURL and only change the distraction settings.
In the no-distraction setting (Figure 2 top row), we can see that all methods perform well. In the dynnmmic-ackground--lll setting (Figure 2 middle row), CoRe and PSE both perform well overall
4

---

Under review as a conference paper at ICLR 2022
100
100
100
WL 9000 Wiir
220
1000
VOL VILOSSOO. JeALE
100
100
200
100
SAC-IRAD RIA
PSC
Raaon
CURL
Figure 2: Results on pixel-based control tasks from the Distracting Control Suite in various distraction settings. Top: No distractions. Middle: Background distractions only. Bottom Background, color, and camera distractions.
Method
Env Steps
Video dataset
BiC-Catch
C-swingup
C-run
F-spin
R-casy
W-walk
DBC
500K
Kinetics (Driving)
650 = 40
260
40
800
+ 50
DBC
500K
DAVIS (2 vidcos)
170
6
443
74
158 :
PSE
500K
DAVIS (2 vidcos)
821 17
749 : 19
308 = 12
779
+ 49
955 + 10
789 +: 28
CURL
500K
DAVIS (2 videos)
580 = 53
624 : 38
313 = 20
851 + 19
416 +: 68
773 +: 35
CoRe (Ours)
500K
DAVIS (2 videos)
802
483 19
743
:: 22
+
CoRe (Ours)
IM
DAVIS (2 videos)
784
635 +: 23
527 26
752
20
620 90
775 +:
CURL
500K
DAVIS (60 videos)
220 42
697 1 20
332 = 17
847
+ 14
302
1 29
750 41
PSE
500K
DAVIS (60 videns)
667
762 +
392 +
829
943
718 +
CoRe (Ours)
videos)
957
CoRe (Ours)
IM
DAVIS (60 videos)
961 + 4
821 + 15
773 + 14
975
869 81
955
Table 1: Background distractions only. Mean + Standard Error. Bold indicates best results at 500K steps.
and better than the other methods. CoRe does significantly better than PSE on the 'cheetah-run' task. SAC-RAD does not perform well, confirming the finding in Stone et al. (02T) that data- aata-auumented model-free RL methods are insufficient to solve these tasks. Recon also does not work well, as expected, because it wastes capacity
modeling the background. In the dynamic-medium setting (Figure 2, bottom row), we see that the performance of CoRe is far better than other baselines. CoRe is able to solve most of the tasks in this challenging distraction setting while other methods struggle to make progress.
In Table T compare the performance of our proposed model with bisimulation based methods (DBC and PSE) and CURL. Results from DBC and PSE have been reported with different background videos making them hard to compare directly. DBC results were reported using videos from the Car Driving class in the Kinetics dataset (Kay et al. 2017), while PSE results were reported when training on 2 videos from DAVIS and testing on 30 validation videos. To do fair comparison, we train CoRe as well as DBC (using their published code) in the same setting as PSE. This setting is denoted as We can see that here, PSE outperforms DBC. CoRe performs similar to
DAVIS (2 videos) in Table I PSE. However, when we use all 60 training videos, CoRe performs better.
3.2 ABLATION EXPERIMENTS
In this set of experiments, we compare CoRe to the variants shown in Table 2 the "ynamic-medium distraction setting. From Table 3 that model-free methods, SAC and QT-Opt, trained
5

---

Under review as a conference paper at ICLR 2022
Attribute
SAC/QT-Opt
RSAC
RSAC+Dyn
Recon
CURL
RCURL
CoRe
Recurrence
Contrastive Loss
Dynamics Modeling
Table 2: Variants produced by ablating the three major elements of CoRe's design. RSAC and RCURL are variants of SAC and CURL which use recurrent architecture (same as CoRe's). RSAC+Dyn models dynamics in the latent state space. Recon does the same, along with pixel reconstruction. All models use data augmentation.
Method
Eny Steps
Mean
BiC-Catch
C-wingup
C-run
F-spin
R-casy
W-walk
SAC + RAD
500K
8
139
7
192
2
24
93 +
31
OT-OOpt+ RADÍ
500K
103
132
241
52 +
25 +
105 + 10
64
RSAC
500K
198 + 20
257 89
277 + 19
244 + 17
90 +
112 + 17
211 +11+15 15
RSAC+ Dyn
500K
244 26
175 =
430 + 39
270
- 32
66 +
141 + 28
379
55
Rocon
500K
147
211 +11+55 55
219 12
147
147
3+1
105 +
199
199+:20
CURL.
500k
223 223
136 +
320
11
170 10
163 + 37
222 + +-37
328
+ 19
RCURL
500k
232
318 79
280 11
216
154 65
106
321 221:10
CoRc (Ours)
500K
480
706
354
26
354 + 10
540
+ 73
445 + 48
479 +1
CoRe (Ours)
IM
684 24
832 + 22
483 29
490 + 13
810 + 60
801 + 32
686 + 42
Table Results on Distracting Control Suite Benchmark (dynamic-mediun setting). Mean + Standard Error. Bold indicates best results at 500K steps. f indicates baselines reported in Stone et al. (2021).
with data augmentation not able to solve these tasks. Adding recurrence (RSAC) and dynamics modeling (RSAC+Dyn) improves results. However, trying to predict pixels hurts performance (Recon). Contrastive prediction in learned space (CoRe) works much better. CURL and RCURL, which use constrastive loss but do not model dynamics, perform worse that CoRe, showing that just having a contrastive loss is insufficient. More ablation experiments that study the relative effects of modeling rewards, forward and inverse dynamics are presented in appendix C.I
3.3 IMPORTANCE OF HAVING A RECURRENT STATE SPACE MODEL
Contrastive models have been reported to underperform DBC (Zhang et al., 202T). However, in our experiments found that contrastive models work quite well. We hypothesize that this difference is because we use recur- rent state space whereas previous work used contrastive learning with feed-forward encoder operating on stacked observation frames. To validate this hypothesis, we com- pare our model to its non-recurrent version. The models both do contrastive future prediction and use similar net- work architectures, but the non-recurrent one consumes three stacked frames. The training mini-batches are sam- pled similarly using 32 sequences of 32 time-steps each. Therefore, even the non-recurrent model has access to ob- servations that semantically close. The results at 500K environment steps in the dynamic-medium setting on DCS shown in Figure 3 (Top). We can see that, in the ab- sence of recurrence, contrastive learning does not perform well. This indicates that just having nearby observations is not sufficient to provide hard negatives and recurrent network has help fix this problem.
300
ropepea Aqpecc menbeb
An intuitive explanation for this is the following. A recur- rent network (in our case, a GRU-RNN) has a smoothness bias built-in because at each time step, it carries forward previous state and only modifies it slightly by gating and adding new information. This is true large extent even during training. and not just at convergence. Therefore, when CoRe is trained, it generates hard negatives through- out training in the form of nearby future predictions. This is true even when the observations have distractions present which change the overall appearance of the observation dramatically. On the other
deeba
Figure 3: Top: Scores drop if recurrence is removed. Bottom: Using negatives only from the same episode is often helpful.

---

Under review as a conference paper at ICLR 2022
Pardo Lightng
Punda Lighting
1000 10000
400
-000
-20
"LLttoootyye Jword
mmm
100
50K
100K
150K
200K
25.0K
2000
500
100K
TCOK
2006
25.K
3002
50K
130K
1O0K
ZEEK.
220K
200K
500
100K
tox
2034 200K
200K
stops
dtepo
st0p0
depp
REAC
SCURL
Colle
SAC With groundl FLth ShER
Figure 4: Results on the Door Opening Task Robosuite Zhu et al. 2020 Each column corresponds to a different distraction setting robot arm type. Top: Camera observation as well as the proprioceptive state is provided to the agent. Bottom: Only camera observation is provided, requiring the vision system do more work.
hand, starting from a random initialization, feed-forward networks are less likely to map observations that are semantically close but visually distinct to nearby points. Therefore, hard negatives may not be found easily during training.
To confirm this further, we train CoRe with a modified contrastive loss in which for each sample in the mini-batch, only the same temporal sequence is used to sample negatives. As shown Figure) (bottom), this is not harmful (but actually beneficial) performance means that for most tasks, CoRe doesn't need other sequences in the mini-batch to find hard negatives. This avoids major concern in contrastive learning, which is having large mini mini-batches to ensure hard negatives can be found. Essentially, recurrence provides an architectural bias for generating good negative samples locally. Performance degrades on the reacher task because observations there contain a goal position for the arm. Contrastive learning tries to remove this information because is constant throughout the episode. Therefore, the actor and critic may not get access to the goal information, causing the agent to fail. This highlights key limitation of contrastive learning it discourages retaining constant information in the latent space. For that type of information, important to have negative examples coming from other episodes.
all tasks, except reacher. This
3.4 RESULTS ON ROBOSUITE
In this section, present results of our method Robosuite 7Zhu al. 2020). MuloCo-based simulator for robotic manipulation tasks. We experiment with two kinds of arms: Panda and Jaco, and two distraction settings: (1) Color and lighting only, and (2) Color, lighting, and camera. In this experiment, we use the static distraction setting to emulate scenario where data collected within an episode has consistent lighting, color, and camera position, but these factors vary across episodes (depending on time of day, accidental movement of the camera relative to the robot setup, etc). The action space in delta end-effector pose. Control is run 20Hz, with each episode lasting for 500 steps (25 seconds). Architecture and training details are described in appendix A
Figure compares CoRe with recurrent versions of SAC and CURL. In the distraction-free case, we also compare to results from SAC operating on ground-truth object state instead of using image observations, as reported by Zhu et al. (2020). This can be seen as an upper bound on performance because the agent gets perfect access to task-relevant state. In the case where proprioception (the robot's 's state) is provided separately (top row). the performance of CoRe is same as or better than that of CURL. In the harder observation, CoRe significantly outperforms other methods. This shows that CoRe is better at handling more complex vision tasks. Comparison other baselines is included in appendix C.4
where the robot's state must also be obtained from the camera
7

---

Under review as a conference paper at ICLR 2022
Figure 5: Gating masks showing that CoRe learss remove distractions. Green arrows indicate the locations the ball and cup, which are tiny, but CoRe can still find them. In the door opening task in Robosuite (bottom right), the lower edge of the door handle is very bright when the door is closed, indicating strong attention to that part of the image. In addition, the agent attends to the edges of the door and the table, as well as its own base. Since the camera position can change randomly, the agent needs to find the robot's position relative the door, for which the robot's 's base and the door's edges are important. Proprioception provided separately, Therefore, the agent doesn't 't need to attend to the arm itself.
3.5 VISUALIZATIONS
In this section, we present qualitative visualizations to help understand what the model does to cope with distractions. We make visualization easier by modifying the observation encoder include pixel-level gating network. This is done using a four-layered CNN which takes the observation as input and outputs a full-resolution sigmoid mask. This mask is multiplied into the observation image, and the result is fed to the encoder. The gating network is trained along with the rest of the model. We expected this gating to improve performance but found that it actually has no effect. This lets us use it as diagnostic tool that factors out the attention mask that the model is otherwise applying implicitly.
the attention masks inferred by the gating network. We can see that the model learns to exclude distractions in the background, even in the dynamic-medium distraction setting for the ball inll_in_cup task where the task-relevant objects are very small. For the door opening task, the model learns attend to the bottom edge of the door handle, the top and right edges of the door, the edge of the table, and some parts of the robot's base. This includes all the information needed to find the relative position of the door with respect the agent. Proprioception is provided separately. Therefore, the agent doesn't It need to attend to the arm itself.
DISCUSSION AND RELATED WORK
Our work uses contrastive learning in the context of world modeling and pixel-based robotic control. In this section, we review related work in these areas.
Pixel-based robotic control Controlling robots directly from raw environment observations is a general framework for robotic problem solving since avoids the use of specific modules such as object detectors and instead relies on the data and reward function to learn a representation of the observations. Model-free methods such as QT-Opt (Kalashnikov et al. 2018) and SAC (Haarnoja et al. 2018) have been applied to pixel-based control. However, learning encoders for nigh-dimensional observations from reward only is sample inefficient. Simply adding observation autoencoding to SAC as auxiliary task has been shown to work well in Yarats et al. (201)). Instead of using generative model, CURL (Laskin et al. 2020a) uses contrastive learning to learn observation encoders yielding similar results. ATC (Stooke al. (2221) improves CURL by predicting representations for future observations, instead of the current one. Data augmentation techniques RAD (Laskin et al., 2020b) and DrQ Yarats et al. 2021) dramatically improve sample efficiency. However, these methods
8

---

Under review as a conference paper at ICLR 2022
still not sufficient to solve visual tasks in the presence of distractions (Stone et al. 2021 Other approaches shape the learned encoders to produce a latent space that is amenable to locally- -ennerr control algorithms (Watter et al., 2015) Cui et al. 2020; Levine et al. 2020). Improvements in the encoder architecture, for example, using attention Tang et al. (2020) and efficient transformer-basee models (Choromanski et al., (02T) have been proposed. Our proposed model can potentially benefit
from the use of these better architectures. We use a simple convolutional encoder for our experiments and focus on the modeling aspect of the problem.
World Modeling World models aim to learn how the world changes as a result of the agent's is action instead of just learning the optimal policy, Finn et al. (2016) used video prediction parameterized as pixel-level flow to predict the future video frames. Dreamer (Hafner et al. 2020) and SLAC (Lee et al. 2020) model observations using a sequential action-conditioned VAH. However, these methods are not designed to deal with distractions since they would waste capacity on reconstructing them along with the task-relevant portion of the observations. PBI. (Guo et al. 2020) improves on Dreamer using an additional term that predicts the state given observations. In order to avoid observation reconstruction, Gelada et al. (2019) proposed DeepMDP, a method that embeds the given high-dimensional RL problem into latent space which preserves the essence of the problem (i.e. the reward and transition function). While this method only preserves the current reward, Deep Bisimulation for Control (Zhang) al., ) uses the future sequence of rewards extract stronger learning signal. In this case, the learned state space is trained to impose a bisimulation metric, under which similarity between states determined by the similarity in future reward sequences produced from those states. PSE (Agarwal iet al. (02T)) further extends this idea to use similarity in future action distributions. While these models use dynamics to get more training signal, they do not make use of structure in the observation space itself. In our proposed model, we seek to find middle ground between modeling observations and not modeling them at all. This is done in a principled way using mutual information maximization between the observation and its prediction from the internal dynamics. While previous work often found contrastive learning to be ineffective, we show that combining it with recurrent state space models makes it work. Recently, contrastive variant of Dreamer (Okada & Taniguchi| (020) has been proposed which shares the same motivation. Concurrent with our work, [Nguyen et al. (2021) explore formulation similar ours based temporal predictive coding, but do not evaluate on the difficult camera and color distractions we do here.
Contrastive Learning: Learning contrast positive and negative examples is a general representa- tion learning tool. It can be applied in many domains, for example, in face verification (Weinberger al. 2006 Schroff et al. (2015), object tracking (Wang & Gupta 2015). time-contrastive models (Ser- manet et al. Hyvarinen & Morioka 2016), language models such as Log- Billneer models Mnih &
Hinton (2007). Word2Vec (Mikolov al.. 2013) and Skip- tkip-thoughhs (Kiros et al.. 2015). Contrastive Predictive Coding (CPC) (van den Oord et al. 2018)) proposes an auto-regressive model to do future prediction at multiple time steps and applies it to various domains including image patches, audio, and C (Guo et al. (019) adds action-condinn coniiinconiiiooinng to the CPC model. Our model uses a similar contrastive loss as CPCJA. However, it also uses latent state correction using the KL term (Equation 2). Subsequent works have improved on CPC using momentum in the learned encoder (MoCo (He et al., (2200). adding projection heads before computing contrastive loss (SimCLR (Chen et al. 2020)), avoiding negative examples altogether (BOOL (Grill et al. 2020) and SWAV (Caron ct al. )0200). Our model can potentially be improved using these techniques but we found that the simple contrastive loss is already sufficient to solve challenging visual control problems. Most similar to our work is CVRL (Ma et al. (020)) which uses contrastive learning with a recurrent architecture but with the standard sequential VAE formulation instead of the one with prior decoding that is used
in CoRe.
5 CONCLUSION
We show that a contrastive objective can be a viable replacement for reconstructing observation pixels in solving visual robotics control problems using madel-based RL. It works even in the presence of severe distractions. In addition, we show that having recurrent architecture helps contrastive learning work well. Even though our results are only in simulated environments, the surprisingly high degree of robustness indicates that this approach should transfer well to a real-world setting. In future work, we plan to test this approach on real-world robots.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization reinforcement learning. In International Conference on Learning Representations, 2021.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In Advances in Neural Information Processing Systems, 2020.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. International Conference on Machine Learning, 2020.
Kyunghyun Cho, Bart van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches, 2014.
Krzysztof Choromanski, Deepali Jain, Jack Parker-Holdds Holder, Xingyou Song, Valerii Likhosherstov, Anirban Santara, Aldo Pacchiano, Yunhao Tang, and Adrian Weller. Unlocking pixels for rein- forcement learning via implicit attention. CoRR, abs/2102.0435 2021.
Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learn- ing by exponential linear units (elus). In International Conference on Learning Representations, 2016.
Brandon Cui, Yinlam Chow, and M. Ghavamzadeh. Control aware representations for model-based reinforcement learning. ArXiv, abs/2006. 13408, 2020.
Chelsea Finn, Ian Goodfellow, and Sergey Levine Unsupervised learning for physical interaction through video prediction. In Advances Neural Information Processing Systems, 2016.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, 2019.
Jean-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 new approach to self-supervised learning. In Advances Neural Information Processing Systems, 2020.
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A. Pires, and Rémi Munos. Neural predictive belief representations, 2019.
Zhaohan Daniel Guo, Bernardo Ávila Pires, Bilal Piot, Jean-Bastien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap catent-predictioo representations for multitask reinforcement learning. In International Conference on Machine Learning, 2020.
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, 2018
Danijar Hafner, T. Lillicrap, Ian S. Fischer, Ruben Villegas, David R Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. ArXiv, abs/181 1t///811.0000 0ad//181104551, 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
Kaiming He, Haoqi Fan, Yuxin Wu. Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Computer Vision and Pattern Recognition (CVPR), 2020.
10

---

Under review as a conference paper at ICLR 2022
1. Higgins, Loie Matthey, A. Pal, Christopher P. Burgess, Xavier Glorot, M. Botvinick, S. Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations, 2017.
Devon Hjelm, Alex Fedorov, Samuel
Karan Grewal, Phil Bachman, Adam
Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In International Conference on Learning Representations, 2019.
Aapo Hyvarinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. In Advances in Neural Information Processing Systems, 2016.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine. Scalable deep reinforcement learning for vision- based robotic manipulation. In Proceedings of The 2nd Conference on Robot Learning, 2018
Will Kay, João Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanarasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, Mustafa Suleyman, and Andrew Zisserman. The kinetics human action video dataset. CoRR, abs/1705.00650, 2017.
Ryan Kiros, Yukun Zhu, Russ R Salakhutdinov, Richard Zemel, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Skip-thought vectors. Advances in Neural Information Processing Systems, 2015.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl:
Contrastive unsupervised representations for reinforcement learning. International Conference on Machine Learning, 2020a.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. In Advances in Neural Information Processing Systems, 2020b.
Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic. Deep reinforcement learning with a latent variable model. In Advances Neural Information Processing Systems, 2020.
Nir Levine, Yinlam Chow, Rui Shu, Ang Li, M. Ghavamzadeh, and H. Bui. Prediction, consistency, curvature: Representation learning for locally-linear control. ArXiv, abs/1909.01 9006, 2020.
Xiao Ma, Siwei Chen, David Hsu, and Wee Sun Lee. Contrastive variational reinforcement learning for complex observations. In Proceedings of the 4th Conference Robot Learning 2020.
Tomas Mikolov, Kai Chen, Greg S. Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space, 2013.
Andriy Mnih and Geoffrey E. Hinton. Three new graphical models for statistical language modelling. In International Conference on Machine Learning, 2007.
Tung D. Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In International Conference Machine Learning, 2021.
Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by latent imagination without reconstruction. CORR, abs/2007. fass/2007.155335 14535, 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, gegh-performncce deep learning library. In Advances in Neural Information Processing Systems. 2019.
Jordi Pont-Tusct. Federico Perazzi, Sergi Caclles, Pablo Arbeláez, Alexander Corkine-Hornung and Luc Van Gool. The 2017 davis challenge on video object segmentation. 170400070 2017.
Florian Schroff. Dmitry Kalenichenko, and James Philbin. Facenet: A unified embedding for face recognition and clustering. In CVPR, pp. 815-823. IEEE Computer Society, 2015.

---

Under review as a conference paper at ICLR 2022
Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, and Sergey Levine. contrasnne networks: Self-supervised learning from video. Proceedings of International Conference in Robotics and Automation (ICRA).
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite- a challenging benchmark for reinforcement learning from pixels. arXiv preprint (xXiv:2011002771 2021.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference Machine Learning, 2021.
Yujin Tang, Duong Nguyen, and David Ha. Neuroevolution of eeffinntrprettble agents. In Proceed- ings of the Genetic and Evolutionary Computation Conference, 2020.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefraneq. Timothy P. Lillicrap, and Martin Riedmiller. Deepmind control suite. CoRR, abs/ 1801 ////1/8010000,, 2018.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In IROS, pp. 5026-5033. IEFE, 2012.
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR, abs/ /1111107..3333 907.03748, 2018.
Xiaolong Wang and Abhinav Gupta. Unsupervised learning of visual representations using videos. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), December 2015.
Manuel Watter, Jost Tobias Springenberg. J. Boedecker, and Martin A. Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances in Neural Information Processing Systems, 2015.
Kilian Q Weinberger, John Blitzer, and Lawrence Saul. Distance metric learning for large margin nearest neighbor classification. In Advances Neural Information Processing Systems, 2006.
Denis Yarats, Amy Zhang, llya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. 2019.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. In International Conference Learning Represeniations, 2021.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, and Roberto Martin-Martin. robosuite: A modular simulation framework and benchmark for robot learning. In arXiv preprint arXiv:2009. 12293. 2020.
12

---

Under review as a conference paper at ICLR 2022
A ARCHITECTURE AND TRAINING DETAILS
In this section, we present the architecture and training hyper-parameters for the proposed CoRe model in more detail.
A.1 GRU-RNN WITH STOCHASTIC AND DETERMINISTIC STATES
CoRe uses a recurrent state-space model (RSSM) that is based on GRUs (Cho et al. 2014) and similar to the one used in Dreamer (Hatner al. 2020), as shown in the figure below. Latent states at any time-step consist of a (deterministic) recurrent hidden state and a stochastic state. The prior latent state is &, [ $2] and the posterior latent state is 2 [h,, 8t). They share the same recurrent hidden state h; but differ in the stochastic component. & depends on hes whereas depends on he and the new observation features 2t. The feed-forward operation is
RSSM
h = GRUCelI GRU MLP(s,-11.04 a.-1))
Future predictor
Corrector
ü, δ PriorMLP(h.).
$: P + â * *N((0, 1),
GRUCClll
r., PosteriorMLP ObsMLP((:)),
$: = μ + o * (0,,
ClaVIL
where MLP stands for fully-connected multi-layered neural network. The figure above illustrates this operation. It can be seen as a more detailed view of the RSSM shown in Figure 1(b). Note that the recurrent hidden state h, propagated through time without sampling. This important because sampling can make it hard to retain information for long time-scales and destroys smoothness of the latent state. However, information inserted into the RNN through s; is stochastic, making the RSSM capable of generating multiple futures. This prevents mode-averaging and allows the model to generate diverse trajectories when doing rollouts.
A.2 ARCHITECTURE
The model consists of number of components: observation encoder CNN, RSSM (which includes PriorMLP, PosteriorMLP, GRU_M _MLP, GRUCell, ObsMLP, along with the reward, inverse dynamics, and observation representation decoder networks), and the actor and critic networks. Table desseribs the architecture of these networks. CNN layers are denotes as inum filters, kernel size, stride]. We borrow the observation encoder architecture from previous work (Laskin et al., 2020a Yarats et al. 2019) and use as-is to ensure an accurate comparison. ELU non-linearity (Clevert et al.) 2016) is used in all places except the observation encoder (which uses ReLU units).
For experiments with Robosuite, the observation encoder architecture is modified to accept 128 x 128 resolution images. Two convolutional layers with 32 filters of kernel size 3 with stride of were applied to reduce the spatial size from 128 x 128 to 32 x 32. This was followed by 3 layers of 3 x convolutions with a stride of 1.
A.3 TRAINING
The training is broken into iterations, where in each iteration one episode of data are collected, added to the replay buffer, and a number of gradient steps are taken using mini-batches sampled from the replay buffer.
Data Collection For the first 1000 steps, data are collected by taking actions sampled from a uniform distribution in 11, After that, data are collected by taking actions sampled from the learned actor policy, which requires rolling out the RSSM. At each training iteration, an entire episode of data are collected, where the length of the episode is 1000 steps of the underlying MuJoCo simulator. The actual number of steps is 1000 divided by the number times the same action is repeated, which is recurrent model
standard based on the task (Table We rollout entire episode because and it seemed natural to use the same model through time within an episode. Typically, model-free RL algorithms collect data one environment step time. We did not explore per-step data collection
13

---

Under review as a conference paper at ICLR 2022
Parameter
Value
Component
Architecture
Replay Buffer
100,000
Ohservation Encoder
CNN: 3x3, 2/. [32, 3x
Initial steps
1000
3, x3 conneccc.
Model Learning rate
3.e-4
layer
Critic Learning rate
1.e-3
PriorMLP
400,400, 400
Actor Learning rate
1.e-3
PosteriorMLP
400, 400]
GRUCell
200 hidden units
Target entropy
adims
GRU_MLP
400.4000 256, 256]. layer norm
Actor update frequency
ObsMLP
Target critic update frequency
Reward prediction
128, 128, 1]
Target critic update
0.005
Inverse dynamics pre-
128, 128, Ogims)
Action log std range
[-10, 2]
diction
KL-weight 3
0.01
Observation represen-
[256, 256, 50), layer-norm
Reward prediction weight o,
1.0
lation decoder
Inverse Dynamics weight aa
1.0
Actor
1024, 1024, 1024, 20,imml
Critic
1024, 1024, 1024, 1)
Weight decay
Gating (optional)
CNN :[16,5 5, [1, x 1, sigmoid.
Critic max grad norm clip
100
Actor max grad norm clip
10
Model max grad norm clip
10
Table 4: Architecture of model components.
Table 5: List of hyper pyper-parameters.
in this work, although that could work just as well. The data from the Distracting Control Suite is rendered 320 240 resolution which is resized to 100 100. Pixels normalized to [0, 1] by dividing by 255. All actions are normalized to lie in [-1, 1] and are modeled using tanh-squashed Gaussian distributions. For robosuite, visual observations at 256 x 256 resolution and a 128 x 128 crop near the center of the image was fed as input to the model.
rendered from the front-view camera
Model updates Training is done with mini-batches of N = 32 sequences of length T = 32 each sampled from the replay buffer. Data augmentation done by taking random 84 x 84 crops. The same crop position is used across the entire sequence. Each mini-batch is used to do three updates sequentially corresponding to the three losses: JM. JQ and Jn. The number of updates is chosen to be 0.5 times the number of steps in episode. All training hyper-parameters are listed in Table 5
Implementation The model is implemented using PyTorch (Paszke et al. 2019). The environment is based on MuJoCo Todorov et al. 2012). All training was done on single Nvidia A100 GPUs. Training time for 500K updates varies from 12-24 depending on the task. Training times differ due to different values of action repeat. Our implementation will be made public.
A.4 TRAINING CURVES FOR INDIVIDUAL LOSS TERMS
200 40 08
U0e
(a) Total loss Jar
(b) Contrastive loss Le
Dynamics LxI.
(d) Inverse dynamics loss
30
030
10c
200
400
(c) Reward loss (r. )
(1) Actor loss Jz.
(g) Critic JQ
(h) Episode reward
Figure 6: Training curves for typical run of CoRe training on cheetah, mynamic-medium setting.
14

---

Under review as a conference paper at ICLR 2022
Figure 7: Left: Examples from the Distracting Control Suite Stone et al. 2021). The top row shows the clean versions of each of the six tasks. The subsequent rows show examples from the easy. medium, and hard distraction settings. Right: Examples from the Robosuite door opening task (Zhu al. 2020p. The top row shows the clean versions of the tasks with a Panda and Jaco arm. The middle row shows observations with color and lighting distractions. The bottom row adds on camera distraction.
There are a number of loss terms which are linearly combined to create the total model loss Jas In this section, present training curves that show how these individual terms optimize with training. Figure 6 shows these loss terms, along with the RI. training losses (Jo and J+) and the episode reward. We see that the total loss Jar goes down expected, along with the individual terms. Note that the reward loss is low in the beginning because the agent gets zero rewards which is easy for the model to predict. As the agent starts receiving better rewards, the reward error increases but then eventually starts to come down. The KL-term has similar behavior. At the very beginning, the prior and posterior latent state distributions match each other (making their KL divergence small) but both are bad at modeling the data. As more training data is encountered, the KL rises sharply, then eventually reduces with more training. The actor loss (which dominated by the negative the Q-value of the actor's 's chosen action) goes down expected, indicating that the actor outputs actions that have high vavuuss. The critic loss appears to increase because the magnitude of the Q-values increases with training, causing the magnitude of the Bellman residual to go up well.
B COMPLETE RESULTS ON DISTRACTING CONTROL SUITE
In this section, report the performance of CoRe on all the distraction settings in the Distracting Control Suite. There are three difficulty levels: casy, medium, and hard as shown in Figure 7 The diffieulty is set by increasing the scale of the camera pose and color change and the number of background distractions that are used (Table 77). For each difficulty level, there are two settings: static, in which the distraction (a background image, is fixed throughout the episode, and dynamic in which the distraction changes smoothly. In the case of background distractions, the video plays back-and-forth, ensuring no sharp cuts. In the of the camera distractions, the camera moves along smooth trajectory.
particular choice of random camera pose)
Table 8 compares the performance of CoRe with model-free RL basclines such as SAC (Haarnoja et al. 2018) and QT-Opt (Kalashnikov et al.) 2018) combined with data augmentation techniques RAD (Laskin al. 2020b) and DrQ Yaratt et al. (02T) as reported in (Stone al. 2021 We also compare with recurrent SAC+RAD baseline (RSAC) which uses the same recurrent architecture as CoRe, but does not contrastively predict the next observation, or model the dynamics and reward. Comparisons are made at 500K environment steps, though we report our results at 1M environment
'The 'hard' level is descrihed the code but results reported in the main paper Stone al.
2021
15

---

Under review as a conference paper at ICLR 2022
Table 6: Action repeat
Task
Action repeat
Table 7: Distraction settings in the Distracting Control Suite.
Ball in cup catch
4
Cartpole swingup
Difticulty
Train videos
Val videos
Camera and Color change scale
Easy
0.1
Cheetah Run
4
Medium
8
0.2
Finger spin
Hard
60
30
0.3
Reacher easy
4
Walker walk
2
steps as well to show that our model continues to improve with more steps. Performance is averaged over 10 random seeds, and 100 validation episodes at each checkpoint. When distractions are present (Table[aa) all methods perform well, though CURL and CoRe are slightly better than the rest in terms of mean scores. On the easy benchmark (Table 8b), CoRe outperforms other methods in all tasks, except reacher. As discussed in Section 3.3, this points to a key limitation of contrastive learning-based methods, which is that they tend to remove constant information (such the goal location for reacher). However, at IM steps, the performance on reacher is much better, showing that the model is able eventually solve the task.
On the medium benchmark (Table CoRe outperforms other methods across all tasks, showing that it can deal with the presence of more distractions. The performance on the reacher tasks improves slightly over the easy setting, which shows that having more variation in the distractions actually helps training, whereas hurts the baseline methods. We also report results on the hard setting (Table [d).. which is documented in the DCS codebase. [Stone et al. (2021) do not report model- -feee baselines for this setting, presumably because the baseline models fail to train reasonable policies at all. However, CoRe is able to get off the ground and get reasonable performance even in this setting. Training to 1M steps continues to improve results.
Figure 8 shows the progression of validation reward for CoRe over 1M environment steps for all tasks and distractions settings. We can see that the hard-dynamic setting (green curves in the bottom two rows) is the hardest to fit because the performance increases slowly. However, in most cases the performance for that setting is still improving at 1M steps, and is likely to get better with more training. Our model struggles even 1M steps, but we can see from these plots that the model is likely to continue improving if trained beyond 1M steps in both static and dynamic settings. In our experiments, we did not tune hyper-parameee specifically for each task, so it is possible that some tasks can benefit from further tuning. In particular, for the reacher task, bigger batch-size can help since that is the only way to get access to diverse target positions.
the reacher task in terms of performance reported at 500K and
C ADDITIONAL ABLATIONS
C.I IMPORTANCE OF INVERSE DYNAMICS AND REWARD PREDICTION
In addition to forward dynamics prediction, the proposed CoRe model includes reward prediction and inverse dynamics prediction as auxiliary tasks. In this section, we present comparisons to ablated versions of our model that remove one or both of these tasks. In Figure 92 we can see that removing both tasks (Fwd only) is significantly worse than having both (Fwd inv + reward). Having any one of these alone is a big improvement on all tasks except reacher, where adding reward is much more important than inverse dynamics. It is interesting to see that in the absence of reward, adding inverse dynamics prediction (Fwd + inv) improves over having forward dynamics only. Asking the model to predict the action that takes the agent from one state to the next is different way of expressing the model dynamics, compared to predicting the next state given the current state and action. The fact that asking the model to do both simultaneously gives boost in performance indicates that inverse dynamics prediction shapes the latent state in ways that are complementary to forward dynamics.
C.2 IMPORTANCE OF UPDATING θ USING CRITIC LOSS
In our model we optimize the parameters θ of the world model (observation encoder and RSSM) using the critic loss Jq- This is similar the choice made in SLAA (Lee al. 2020) and DBC (Zhang
16

---

Under review as a conference paper at ICLR 2022
Table 8: Results on Distraction Control Suite benchmark. All results reported 500K steps, unless mentioned otherwise. Mean reward Standard Error. Bold numbers indicate the best performing models at 500K steps. indicates baselines reported in Stone et al. (2021.
(a) No distraction
Method
Mean
BiC-Catch
C-wwinuup
C-run
F-spin
R-easy
W-walk
SAC+RADT
836 1 29
962 =
843 = 10
515 1 13
976 =
962 1 8
762 1 184
QT-Opt+RADD
820 +
968
843 14
538 +
953
969 +
648
QT-Opt+DrQ'
801 5
962 =
851
534
952 =
974
532 : 29
RSAC
815 + 32
873 + 80
629 + 108
527 + 28
977 + -
952 +
934 +
CURL
859 + 21
963 -
850 3
534 1 37
922 = 29
949 - 13
937 1
CoRe
853 + 19
953 =
796 19
628 + 16
950 =
846 62
942 +
CoRe (IM steps)
879 + 20
811 + 97
852 7
730 +
966 13
963 7
954 + 2
(b) Benchmark casy
Method
Mean
BiC-Catch
C-swingup
C-run
F-spin
R-easy
W-walk
SAC+RADÍ
182 24
129 + 20
360
25
370 114
102
60 + 31
Static
OT-OPt+RADi
317
218 +: 44
446
220 ::
128
Qr-OOttDrr''
299 6
217 35
416 1:61:0 20
199 1
695 1:33
171 7
93 13:19
RSAC
2:
205 + 44
491 3:36
283
92
113 =
459
CURL
418
165 35
430
357 +37+-1 13
759
142
657
CoRe
634 +
854
562
459
319
742
CoRe (1M steps)
15
666
SAC+RAD
270
366 +: 59
297 21
198 + 39
338 + 59
173 11
249 + 138
Dynumic
QT-OPttRADD
343
490 64
467 + 12
170 =
393 + 91
428 68
109 12
QT-Opt+Dre
265
395 39
431
126 +
203 +
343
+
RSAC
275 1: 24
181 3:
465 21
292 1: 10
55
145
482 20
CURL
391 3:
102
1: 20
432
233 1:
648 32
253 40
678
35
CoRe
586 + 30
798
+ 30
499
22
423 + 22
713
+ 81
340 60
744
CoRe (IM steps)
722 +: 28
909 + 10
590
17
569
1 19
823 75
552 83
889
26
(c) Benchmark medium
Method
Mean
BiC-Catch
C-swingup
C-run
F-spin
R-easy
W-walk
SAC+RAD'
113 + 12
96
272 = 11
21 1 15
169 92
93 1
25 +
Static
QT-OPt+RADi
+
172 + 12
13
234 67
14100 16
63 3
QT-Opt+DrQ
170 : 11
169 : 25
283 5
124
266
112 + 16
64 4
RSAC
199 + 20
164 23 23
428 = 17
164 11
9
80 +
350 15
CURL
300 +: 28
124 - 33
304 - 20
277
621
+1
74
402 + 78
CoRe
561 +: 29
762 + 27
509 + 14
402 + 15
880 +
219 25
593 26
CoRe (IM steps)
690 1: 26
743 + 88
634 + 11
526 1: 17
924 4
543 4: 56
766 30
SAC+RAD
89 +
139
192
14+2
63 + 24
+
31 +
Dynumic
QT-OP++RADt
103+3
132 + 20
241 +7
52+33
+
105 10
64 42
QT-Opt+DrQ"
102+
114 22
243 +
54 +-
26
108
+5
65 +1
RSAC
198 + 20
257 + 89
277 + 19
244
+
90
47
112
1
17
211
111+1
15
CURL
223 +: 14
136 1 24
320 - 11
170
+ 10
163 3:
222
+ 37
328
+ 19
CoRe
480 + 23
706
+ 39
354
+ 26
354
+
10
540
1 73
445
1 48
479
1 31
CoRe steps)
684 1 24
832 + 22
483
+ 29
490
1
13
810
+ 60
801
4: 32
686
1 42
(d) Benchmark hard
Method
Mean
BIC-Catch
C-swingup
C-run
F-spin
R-easy
W-walk
RSAC
145 +: 14
113 + 12
295 + 40
140
42 10
74 +
204 + 24
Statie
CURL
202 + 16
163 + 36
199 = 30
239
1: 12
244 13
121
4: 16
247 + 53
CoRe
499 + 28
710 + 37
447 = 10
339 + 13
809 1
197
4 21
490 + 15
CoRe (IM steps)
638 + 27
875 9
554 - 11
469 + 21
898 + 12
386 - 39
645
+ 29
RSAC
138 + 10
165 + 19
213
151
56
86
158
+ 14
Dyn
CURL
95 +
103 - 18
192 1
78 + 13
5
73
119
25
CoRe
307 + 22
436 1 48
257 - 18
200
364 1 83
234 61
353 25
CoRe (IM steps)
467 + 29
562 + 65
350
26
345 +
620
419
505
al. 2021). However, reasonable alternative could be to optimize θ only using JM and keep the critic training separate. This would amount to separating the world model from the controller. As shown in Figure 9b] when excluding & from the critic including exclusion performs comparably on two tasks (ball in cup. cartpole), better on one (finger) and worse on three (cheetah, reacher, and walker). Overall, the inclusion regime works better. Therefore, chose to include O in the
17

---

Under review as a conference paper at ICLR 2022
ball Jn cup
corgole
1030
3,00 Easy Medium
3500 Easy Medium
Vol. elLeppoo. roward
200
Statioklard
Enger
wilker
1000
Studi Eusy
Etalio Eccy
Vol. cpocee roward
200
Stari--aayy
Stanko-aad
0101
11:2
11:4
0.55
D3
10
00
0.
114
0.15
D35
1.0
00
0.2
11.
0.55
08
10
dlopss
966
stops
965
steps
1e6
ball in cup
corgole
cheetah
Dynami
Dynam Faay
Dycamo-teedumm
Dynamio
Dy
vol. optoeds roward
400
200
Dynamio-Hard
tnggrr
maccher
walker
Dynamic
Dyramic
Dynamic
Cyrumii
vol. opioodo reward
400
200
0.0
0.4
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
00
0.2
0.4
0.6
08
stapa
166
stopa
TwB
Stapa
Tel
Figure Progression validation reward with training steps on distraction settings in the Distracting Control Suite. Top two rows are for the static setting and bottom two for dynamic. Each plot shows easy, medium, and hard difficulty levels.
critic. In future work, important direction to explore is the exclusion regime, especially in the multi-task setting, because separating the controller from the world model enables separating general understanding of the world from task-specific control policies, which is key generalization.
18

---

Under review as a conference paper at ICLR 2022
1000
Find ony
Crllo undades mode coaats
Find Romadd
300
Crdo tebadce soadd pesr
Feid
Podd Rievend
400
200
bell
carbcce
deodah
Enger
noadhar
bx
curaolo
choolal
Ergan
vulla
(a) Inverse dynamics and reward prediction.
(h) Optimizing (e, (Q.).)
Figure 9: Additional ablation results. Left In addition to forward dynamics, inverse dynamics prediction and reward prediction improve performance. Right: Optimizing model parameters O using the critic loss is beneficial.
tean3
pntarng D307
tear2 2311
betardd.1
belamo.1
1.
betar 100001
bea--3 2001
Performance of CoRe for different values of the KL-loss weight Bottom Performance a variant in which contrastive prediction is done from the posterior latent state, rather than the prior. We can see higher performance variance across values of 33 when the posterior latent state is used.
Figure 10: Top:
C.3 ROBUSTNESS TO WHEN RECONSTRUCTING FROM PRIOR VS POSTERIOR
CoRe predicts the next observation's 's feature from the prior latent state E, making it different from sequential VAE-based models like SLAC (Lee et al., 2020) and Dreamer (Hafner et al. 20200) where the posterior latent state is used to reconstruct the observation. We argued in the paper that doing so makes the model more robust to the choice of 3, the weight applied to the KL-term during optimization. In this experiment, we verify this argument by comparing CoRe with variant where the posterior latent state is decoded to contrastively match the true observation's representation. We train both models with five values of 3 and 5 different seeds each. In Figure 10 we see that the posterior version (bottom row) has more variance in performance across different values of B. The proposed CoRe model (top row) is more stable, and hence, easier train.
19

---

Under review as a conference paper at ICLR 2022
Parda Naw
Parde Lastingging
Panda Pante
Color
stepe
mape
staps
SACHERD
Piacon
Cafte
Figure 11: Results on the Robosuite Door Opening Task. agents receive only RGB inputs.
C.4 COMPARISON WITH MORE BASELINES FOR THE ROBOSUITE DOOR OPENING TASK
In this section, we present a comparison of CoRe with SAC+RAD, PSE, Recon, and CURL on the door opening task in Robosuite. The set of baselines is the same as that used for the Distracting Control Suite in Figure Figure TT shows the results. In this setting, the observations only come from the camera, i.e. proprioception is not provided separately. CoRe is able to perform significantly better than the baselines all settings that involve distractions.
20


Reviews:
[{"review_id": "u2n1VTkOARM", "reviewer": "Reviewer_skrV", "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": "# Strengths\n* The clarity of the paper is good and the approach has been described well.\n* The experiments have been well designed and indicate strong improvements in the model robustness (particularly Figure 2 and Table 3). Table 2 was helpful to understand the key design choices and their impact on final performance. \n* The mask visualizations in Figure 5 provide good clarity on how the robustness is achieved by the proposed method.\n* The ablation studies in the main paper and supplementary clearly highlight the benefits of individual components of the method.\n\n\n# Weaknesses\n\n## Novelty\n* The novelty of the approach isn't quite clear. In Sec. 1 (page 2), the authors highlight that one of the key findings is that \"contrastive learning can in fact lead to surprisingly strong robustness to severe distractions, provided that a recurrent state-space model is used\". This finding itself is valuable in my opinion, but the method \"CoRe\" does not seem novel to me. Relative to Dreamer and SLAC, the key novelty appears to be the contrastive loss term from Eqn (1) which predicts future observations instead of auto-encoding. However, the idea of using a recurrent model to predict encodings of future observations (i.e., not an auto-encoder) has been studied in CPC [1], action-conditioned CPC (aka CPC|A) [2], Predictions of Bootstrapped Latents (aka PBL) [3], etc. CPC does not use action conditioning, but uses contrastive learning. CPC|A uses action conditioning + contrastive learning (very similar to CoRE). PBL uses action conditioning + reconstruction (similar to \"Recon\" baseline, but has an additional loss to predict state conditioned on observation representation). \n* Given the above, the methodological novelty of CoRe is not clear to me (particularly, relative to CPC|A). Is there something different about the contrastive loss in CoRe (relative to CPC|A) without which the performance degrades severely? Or is the novelty only in the observation that recurrent models are needed for contrastive loss to work well? Note, while some recent work like DBC does not use recurrent models, prior approaches already use them with contrastive loss. \n\n\n## Intuition behind why recurrent state models are needed for contrastive learning\n*  [End of section 1] The authors suggest that when recurrent models are used along with contrastive learning, the smoothness of the state-space ensures the presence of informative hard negatives in the same mini-batch of training. This isn't clear to me. Positives and negatives are obtained from the next-step real observation's encodings (just a single frame) and does not use the recurrent model. How is the smoothness of the state-space related to hard negatives? While the performance degrades without the recurrent model in Figure 3 (top), isn't this more likely due to a poorer state representation caused by the lower capacity (no RNN) and lack of information aggregation over time? This would affect all models, not just the CoRe.\n\n\n## Missing baselines / related works\nA few related methods have not been compared with or discussed. \n\n* Augmented Temporal Contrast (aka ATC) [4] does not use recurrence, but uses future observation prediction and has been shown to achieve good improvements over CURL. \n* CPC|A [2] uses recurrence, future observation prediction (contrastive), and action conditioning. This is very similar to the proposed method and should be compared with.\n* PBL [3] uses recurrence, future observation prediction (reconstruction), and action conditioning. This method introduces an additional P(state | observation) term that results in bootstrapped learning (might improve over the Recon baseline).\n\n\n## Other concerns\n\n* Why is CoRe worse than PSE on 4 / 6 tasks for DAVIS (2 videos) in Table 1? Why is CoRe better with 60 videos?\n* Why are only 2 baselines used in Figure 4? Can the authors include the complete set? Do we observe similar trends as Figure 2?\n* Are the findings from Figure 5 specific to CoRe? Or do other methods (like PSE) also learn similar masking functions? \n\n\n\n[1] CPC - Oord, Aaron van den, Yazhe Li, and Oriol Vinyals. \"Representation learning with contrastive predictive coding.\" arXiv preprint arXiv:1807.03748 (2018).  \n[2] CPC|A - Guo, Zhaohan Daniel, et al. \"Neural predictive belief representations.\" arXiv preprint arXiv:1811.06407 (2018).  \n[3] PBL -  Guo, Zhaohan Daniel, et al. \"Bootstrap latent-predictive representations for multitask reinforcement learning.\" International Conference on Machine Learning. PMLR, 2020.\n[4] ATC - Stooke, Adam, et al. \"Decoupling representation learning from reinforcement learning.\" International Conference on Machine  Learning. PMLR, 2021. ", "summary_of_the_paper": "The paper proposes a recurrent state-space model that learns robust representations for robotic control. The proposed method builds on top of prior works on world-models which learn a latent dynamics model of the agent which can be used for planning and action selection. Different from prior work such as Dreamer and SLAC which rely on pixel-based observation reconstruction, this paper highlights that a simpler contrastive loss for the next-observation prediction achieves better results **if** a recurrent state-space model is used for the latent space. Results are presented on the Distracting Control Suite benchmark and show strong improvements over prior approaches.", "summary_of_the_review": "I am concerned about the lack of clear novelty in the paper, and other experimental issues highlighted in \"Weaknesses\". I will update my rating based on the author's responses.", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "bAsyaEtLsO", "reviewer": "Reviewer_FBRd", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "## Strong points\n\nThe paper is well written and easy to follow. I particularly\nliked the more straight-to-the-point and honest approach on writing.\nThe authors could have easily try to mask their similarity with some\nother methods in the literature but they were upfront with that.\nThis made the understanding of the paper way easier and also\neasier to understand the main contributions.\n\nThere is a substantial comparison with relevant models in the literature that already\ndeal with the DCS. The results obtained are really impressive\nand are a new state-of-the-art  on interactive environments\nwith distractive backgrounds.\n\nThe ablations specifically the ones concerning the usability of the recurrent\ncontrastive method are very useful. Specially the intuition on the negative\nsamples per mini batch.\nFurther, the appendix provides even more ablations, making very\nclear that a recurrent state has a major impact on the general performance.\n\n\nThe visualizations are also good standard to be used on this field which performs\ncontrol directly from pixels. They also show that their model is capable of directly\ndetecting the object of interest.\n\n##  Weak points\n\nThe proposed model is indeed very similar to Dreamer[1]. The authors address specifically\none point from dreamer: the contrastive learning strategy, which failed to produce better\nresults as reported by [1]. However, I do believe that this closer look into this\nspecific point can be useful for the community and insightful in general.\n\n## Questions\n\nI would be curious to find more references on the idea that hard negatives needs to be found\non the mini batch in order to contrastive learning to have a good performance. \n\nI wonder how the performance of this method would be in a more open task than mujoco\nstyle robotic control. Tasks with a different initial state or with other elements in the scene \nthat the agent need to interact might have a negative impact on this method results.\n\n\n[1] Hafner, Danijar, et al. \"Dream to control: Learning behaviors by latent imagination.\" arXiv preprint arXiv:1912.01603 (2019).", "summary_of_the_paper": "The paper shows a latent-planning RL model  to\nlearn control policies directly from pixels. To learn a better\nrepresentation it uses recurrent model contrastive learning approach\nwhich enhances the representation learning performance of single\nframe based contrastive methods.\nThis was tested robotic control suites with challenging distracting\nbackgrounds.\n\nThe main contribution of the work is the addition of recurrence to the\nmodel and the extensive testing and explanation on the reason it tends\nto work better.", "summary_of_the_review": "I think this is a very useful and well written paper.\nEven though the scope is small, the results are convincing\nand it shows a very clear way on how to effectively use contrastive\nrepresentation learning methods  while learning to control directly from\npixels.\n\n\n## After Rebuttal\n\nAfter reading the other reviewers comments and the rebuttal \nI see that i missed some literature that needed further comparison.\nI think it is a good and well written paper but I would lean for rejection \ngiven this new data.", "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": {}}, {"review_id": "SECOp22JoY", "reviewer": "Reviewer_cJX3", "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": "Weaknesses:\n\nThe major issue is that the proposed idea and the experiment setup is not novel. They highly overlap with a prior CoRL 2020 paper, Contrastive Variational Reinforcement Learning for Complex Observations (CVRL) [1], which, however, has not been cited in the submitted manuscript. They overlap in the following aspects:\n\n1/ The idea is the same. CVRL also extends RSSM using contrastive learning and aims to improve the robustness of the learned model against real-world observations with high frequency noise. Both of them use InfoNCE for contrastive learning and the same RSSM structure.  In addition, both CVRL and the proposed method use the policy loss from dreamer for learning the policy network. The resulting equations and loss functions of the two algorithms are almost the same.\n\n2/ CVRL also experimented on natural Mujoco games, which introduces moving backgrounds into the standard dm-control suites. This is exactly the same with the distracted dm-control suites used in the submitted manuscript.\n\nBeyond what has been discussed in the paper, CVRL has mathematically shown that by replacing the generative observation likelihood with a contrastive objective, we can lower bound the original ELBO.\n\nThere are some other weaknesses, but I believe the issues discussed above are sufficient to make it a clear rejection.", "summary_of_the_paper": "This paper has presented CoRe, Contrastive Recurrent state space model, for model-based robust model-based reinforcement learning for robotic control. Standard reconstruction-based state space models are less robust in the unstructured real-world scenarios because of the high-frequency details. Instead, CoRe learns the state space model with contrastive learning, which greatly improves robustness. In addition to this, a policy is being learned with SAC. Experiments on distracting control suites and several robotic control tasks demonstrate the better robustness of CoRe.", "summary_of_the_review": "The paper is highly similar to a prior work as mentioned above, and as a result, the contribution of the paper is very limited. I would vote for a rejection.", "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": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Meta Review of Robust Robotic Control from Pixels using Contrastive Recurrent State-Space Models

This work investigates a recurrent latent space planning model for robotic control from pixels, but unlike some previous work such as Dreamer and RNN+VAE-based World Models, they use a simpler contrastive loss for next-observation prediction. They presented results on the DM-control suite (from pixels) with distracting background settings. All reviewers (including myself) agree that this is a well-written paper, with clear explanation of their approach. The main weaknesses of the approach are on the experimental side (see review responses to author’s rebuttal by skrV and cjX3). Another recommendation from me is to strengthen the related work section to clearly position the work to previous work - there is clear novelty in this work, but this should be done to avoid confusion. The positive sign is that in the discussion phase, even the very critical cjX3, had increased their score and acknowledged the novelty from previous related work. In the current state, I cannot recommend acceptance, but I’m confident that with more compelling experiments recommended by the reviewers, and better positioning of the paper to previous work, I believe that this paper will surely be accepted at a future ML conference or journal. I’m looking forward to seeing a revised version of this paper for publication in the future.
Decision: Reject