DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representations
Authors: Fei Deng
Ingook Jang
Sungjin Ahn
Venue: ICLR 2022 Submitted (2022)
Keywords: model-based reinforcement learning
representation learning

Abstract: In model-based reinforcement learning (MBRL) such as Dreamer, the approaches based on observation reconstruction
often fail to discard task-irrelevant details, thus struggling to handle visual distractions or generalize to unseen distractions. To address this issue, previous work has proposed to contrastively learn the latent representations and its temporal dynamics, but showed inconsistent performance, often worse than Dreamer. Although, in computer vision, an alternative prototypical approach has often shown to be more accurate and robust, it is elusive how this approach can be combined best with the temporal dynamics learning in MBRL. In this work, we propose a reconstruction-free MBRL agent, called DreamerPro, to achieve this goal. Similar to SwAV, by encouraging uniform cluster assignment across the batch, we implicitly push apart the embeddings of different observations. Additionally, we let the temporal latent state to 'reconstruct' the cluster assignment of the observation, thereby relieving the world model from modeling low-level details. We evaluate our model on the standard setting of DeepMind Control Suite, and also on a natural background setting, where the background is replaced by natural videos irrelevant to the task. The results show that the proposed model is consistently better than the previous models.

Under review as a conference paper at ICLR 2022
DREAMERPRO: RECONSTRUCTION-FRER MODEL-BASED REINFORCEMENT LEARNING WITH PROTOTYPICAL REPRESENTATIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Top-performing Model-Based Reinforcement Learning (MBRL) agents, such as DREAMER, learn the world model by reconstructing the image observations. Hence, they often fail to discard task-irrelevant details and struggle to handle visual distractions. To address this issue, previous work has proposed to con- trastively learn the world model, but the performance tends to be inferior the absence of distractions. In this paper, seek to enhance robustness to distractions for MBRL agents. Specifically, we consider incorporating prototypical represen- tations, which have yielded more accurate and robust results than contrastive ap- proaches in computer vision. However, it remains elusive how prototypical repre- sentations can benefit temporal dynamics learning in MBRL, since they treat each image independently without capturing temporal structures. To this end, we pro- pose to learn the prototypes from the recurrent states of the world model. thereby distilling temporal structures from past observations and actions into the proto- types. The resulting model, DREAMER PRO, successfully combines DREAMER with prototypes, making large performance gains on the DeepMind Control suite when there are complex background distractions, while maintaining similar per- formance DREAMER in the standard setting.
INTRODUCTION
Model-Based Reinforcement Learning (MBRI., Sutton & Barto 2018 Sutton 1991) provides a so- lution to many problems in contemporary reinforcement learning. It improves sample efficiency by training a policy through simulations of a learned world model. Learning a world model also pro- vides way to efficiently represent experience data as general knowledge simulatable and reusable in arbitrary downstream tasks. In addition, it allows accurate and safe decisions via planning.
Among recent advances in image- mmsee-based MBRL, DREAMER is particularly notable as the first MBRL model outperforming popular model-free RL algorithms with better sample efficiency in both con- tinuous control (Hafner al.) 2020) and discrete control (Hafner et al. 2021) Unlike some previous model-based RL methods (Kaiser et al., 2019). it learns a world model that can be rolled out a compact latent representation space instead of the high-dimensional observation space. Also, policy learning can be done efficiently via backpropagation through the differentiable dynamics model.
In image-based RL, the key problem to learn low-dimensional state representation and, in the model- case, also its forward model. Although we can learn such representation directly by usually very slow do this due to the maximizing the rewards (Schrittwieser al., 2020), auxiliary tasks providing richer learning reward sparsity. Instead, it is more practical to introduce signal to facilitate representation learning without reward (or with sparse reward) (Sutton et al., 2011 Jaderberg et al. 2016). DREAMER achieves this by learning the representation and the dy- namics model in a way to reduce the reconstruction error of the observed sequences. However, acconstruction-based representation learning has limitations. First, it is computationally expensive to reconstruct the high-dimensional inputs, especially in models like DREAMER that needs to re- construct long-range videos. Second, it wastes the representation capacity to learn even the visual signals that are irrelevant to the task or umpredictable such as noisy background (Burda et al.) 2018). Thus, in MBRL it is of particular interest realize version of DREAMER without reconstruction.

---

Under review as a conference paper at ICLR 2022
Recently, there have been remarkable advances in reconstruction free representation learning in inforcement learning (Laskin et al. 2020a b Yarats et al. 2021c). The currently dominant approach is via contrastive learning. This approach requires pair wise comparisons to push apart different instances while pulling close an instance and its augmentation. Therefore, this method usually re- quires large batch size (so computationally expensive) to perform accurately and robustly. An alternative is the clustering-based or prototype-basee approach (Caron et al.) (020). By learning set of clusters represented by prototypes, it replaces the instance-wise comparison by a comparison to the clusters and thereby avoids the problems of contrastive learning. This approach is shown to perform more accurately and robustly many applications (Caron et al.) 2020 2021 Yarats al.)
2b2 than the contrastive method while also alleviating the need for maintaining large batch size. The prototype structure can also be used to implement an exploration method (Yarats et al. 2021b).
However, for reconstruction-feee free MBRL only the contrastive approach like Temporal Predictive Coding (TPC, Nguyen al., 202T) has been proposed so far. While TPC consistently outperforms DREAMER in the noisy background settings, for standard DeepMind Control suite (Tassa et al.) 2018 it showed quite inconsistent results by performing severely worse than DREAMER on some tasks. Therefore,
hypothesize that this inconsistent behavior may be fixed if the robustness and accuracy of the prototypical representations can be realized in MBRL and further improved with the support of temporal information.
In this paper, we propose a reconstruction-fret MBRL agent, called DREAMERPRO, by combining the prototypical representation learning with temporal dynamics learning. Similar to SWAV Caron itt al. 2020), by encouraging uniform cluster assignment across the batch, we implicitly pull apart the embeddings of different observations. Additionally, we let the temporal latent state to 'reconstruet' the cluster assignment of the observation, thereby relieving the world model from modeling low- level details. We evaluate our model on the standard setting of DeepMind Control suite, and also on a natural background setting, where the background is replaced by natural videos irrelevant to the task. The results show that the proposed model consistently outperforms previous methods.
The contributions of the paper are (1) the first reconstruction-fie MBRL agent based on the proto- typical representation and temporal dynamics and (2) the demonstration of the consistently im- proved accuracy and robustness of the proposed model in comparison to a contrastive reconstruction- free MBRL agent and Dreamer for both standard and natural background DMC tasks.
PRELIMINARIES
In this section, we briefly introduce the world model and learning algorithms used in DREAMERV2 (Hafner et al. 2021) which our model builds upon. To indicate the general DREAMER framework (Hafner et al.) 2020 2021), we omit its version number in the rest of the paper.
2.1 RECONSTRUCTION-BAN WORLD MODEL LEARNING
DREAMER learns a recurrent state-space model (RSSM, Hafner et al. 2019 to predict forward dynamics and rewards in partially observable environments. At each time step the agent receives an image observation og and a scalar reward тr (obtained by previous actions acc). The agent then chooses an action a, based on its policy. The RSSM models the observations, rewards, and transitions through a probabilistic generative process:
/ II
T
p(OLTT, FIT (I.T)
p(o;
I 8gg, aee) p(sx) I see, ace) da
(1)
t=1
II p(o: he 8t) p(re) he 8t) p(s: he) dsLr
where the latent variables $1.7 are the agent states, and h GRU(h st-1 a=-1) is a deter- ministic encoding of S<t and a<r. To infer the agent states from past observations and actions, variational encoder is introduced:
T
((s 01:7 01:))
II a(s set, 021,0)) t=1
II g(s h ht)
(3)
2

---

Under review as a conference paper at ICLR 2022
The training objective is to maximize the evidence lower bound (ELBO):
JDREAMER
Σ
By Dog
hes st) + log he, 8c)
DKL(a(s: (t) ) p(s ht))) (4)
38.
JR
Tki
2.2 POLICY LEARNING BY LATENT IMAGINATION
DREAMER interleaves policy learning with world model learning. During policy learning, the world model is fixed, and an actor and critic are trained cooperatively from the latent trajectories imagined by the world model. Specifically, the imagination starts at each non-terminal state & [h, s,) encountered during world model learning. Then, at each imagination step 1' > t, an action is sampled from the actor's stochastic policy: āy ~ π (ā ). The corresponding reward f+ and next state 21111 are predicted by the learned world model. Given the imagined trajectories, the actor improves its policy by maximizing the X-return (Sutton & Barto 20188) Schulman et al., al., 2018)) plus entropy regularizer that encourages exploration, while the eritic trained to approximate the X-return through a squared loss.
3 DREAMERPRO
To compute the DREAMER training objective, more specifically To in Equation a decoder is required to reconstruct the image observation o, from the state 2t [e). Because this recon- struction loss operates in pixel space where all pixels are weighted equally, DREAMER tends to allocate most of its capacity modeling complex visual patterns that cover large pixel area (e.g.. backgrounds). This leads to poor task performance when those visual patterns are task irrelevant, as 2021I.
shown in previous work (Nguyen
Fortunately, during policy learning, what we need is accurate reward and next state prediction, which are respectively encouraged by JR. and Jkr. In other words, the decoder not required for pol- icy learning. The main purpose of having the decoder and the associated loss To. as shown in DREAMER, learn meaningful representations that cannot be obtained by Th and JKL alone.
The above observations motivate us to improve robustness to visual distractions by replacing the reconstruction-basee representation learning DREAMER with reconstruction-fiee methods. For this, we take inspiration from recent developments in self-supervised image representation learning, which can be divided into contrastive (van den Oord et al. 20199] Chen al., 2020 He et al. 2020) and non-contrastive (Grill et al. 2020 Caron et al. 2020) methods. We prefer non-contrastive methods as they can be applied to small batch sizes.
This can speed up both world model learning propose to combine DREAMER with the
and policy learning (in wall clock time). Therefore,
prototypical representations used in SWAV (Caron et al.] 2020). a top-performing non-contrastive representation learning method. We name the resulting model DREAMERPRO, and provide the model description in the following.
DREAMER PRO uses the same policy learning algorithm DREAMER, but learns the world model without reconstructing the observations. This is achieved by clustering the observation into a set of K trainable prototypes {c1,.... ...,,... CK) ). and then predicting the cluster assignment from the state as well as an augmented view of the observation. See Figure [ an illustration.
Concretely, given a sequence of observations 01:7 sampled from the replay buffer, we obtain two
augmented views oir, (2) by applying random shifts (Laskin et al. 2020b Yarats et al. 2021c) with bilinear interpolation (Yarats al. 2021a). We ensure that the augmentation is consistent
(() across time steps. Each view E [1,2} is fed to the RSSM to obtain the states =1:7-- To predict
the cluster assignment from =t we first apply linear projection followed by éz-normalization to
obtain a vector xt of the same dimension as the prototypes, and then take a softmax over the dot products of 2 and all the prototypes:
(1)
(8) (.k)
softmax
CK
(5)
T
3

---

Under review as a conference paper at ICLR 2022
Projections
Probabilities
RSSM State
zt
I t
Softmax
Ut
T
Sinkhorn-Knopp (over batch)
Temp
fō
aug(o:)
Yt
Wt H1
JEMA
JSWAV
fe
Softmax
aug(or)
Yt
Ut
Prototypes
...
Stop Gradient
Figure 1: DREAMERPRO learns the world model through online clustering, eliminating the need for reconstruction. At each time step f, it first compares the observation "ci.....ck] to obtain the target cluster assignment W/-- Then, it predicts this target from both the world model state žt and another augmented view of the observation (each aug(or) denotes an independent application of data augmentation). The predictions are improved by optimizing the two objective terms, Fomm and Jww., respectively, where the first term crucially distills temporal structures from 2 into the prototypes.
of trainable prototypes
(i) Here, Ui is the predicted probability that state 2/ maps to cluster k. 7 is temperature parameter, and the prototypes ck) are also no-omrmalieed
Analogously, to predict the cluster assignment from an augmented observation of we feed to a convolutional encoder (shared with the RSSM), apply linear projection followed by lg-
normalization, and obtain vector ( We summarize this process as: y, fe(o,"). where collectively denotes the parameters of the convolutional encoder and the linear projection layer. The
prediction probabilities are again given by a softmax:
( yp
yi
CK
K = softmax
(6)
(8) where og is the predicted probability that observation 0g maps to cluster
To obtain the targets for the above predictions (i.e., Equations 5 [an we apply the Sinkhorn-
Knopp algorithm (Cuturi) 2013) to the cluster assignment scores computed from the output a
momentum encoder fø (He et al. 2020 Grill al. 2020 Caron al. 2021). whose parameters 0
are updated using the exponential moving average of 0: 4 17)0 + n0.. For each observation
of the scores are given by the dot products (g(() Ci....,
(6)
CK). where v.())
(o() is
the momentum encoder output. The -inkhorn-Knopp algorithm is applied to the two augmented
batches (o ) {o(2)- separately to encourage uniform cluster assignment within each augmented batch and avoid trivial solutions. We specifically choose the number of prototypes K = B x T, where B is the batch size, that the observation embeddings are implicitly pushed apart from each other. The outcome of the Korkhorm-Knoph algorithm is a set of cluster assignment targets
(wt,1 (6) for each observation o.
4

---

(1) Under review as a conference paper at ICLR 2022
Now that we have the cluster assignment predictions and targets, the representation learning objec- tive is simply to maximize the prediction accuracies:
K
JSWAV
(1) t,k
(2) log Vekk + wif log ui.k
(2)
k=1
(1) (1) Wiih log ui + log u
(2)
Temp
(8)
Here, Tswy improves prediction from an augmented view. This is the same loss as used in SWAV (Caron al., 2020), and shown to induce useful features for static images. However, it ignores the temporal structure which is crucial in reinforcement learning, Hence, we add a second term, Femp- that improves prediction from the state of the same view. This has the effect of making the prototypes close to the states that summarize the past observations and actions, thereby distilling temporal structure into the prototypes. From another perspective, Temp is similar to T in the sense that we now 'reconstructing' the cluster assignment of the observation instead of the observation itself. This frees the world model from modeling complex visual details, allowing more capacity to be devoted to task-relevant features.
The overall world model learning objective for DREAMERPRO can be obtained by replacing To in Equation 4 with JwwV + I TTmp'
PRREAMERPRO
Σ
Eg + Reum + + Tkul
Eg
(9)
where Th and Thr are now averaged over the two augmented views.
4 EXPERIMENTS
Environments. We evaluate our model and the baselines on six image-based continuous control tasks from the DeepMind Control (DMC) suite TTssa al. 2018). We choose the set of tasks based on those considered in PLANET (Hafner et al. 2019). Specifically, we replace Cartpole Swingup and Walker Walk with their more challenging counterparts, Cartpole Swingup Sparse and Walker Run, and keep the remaining tasks. In addition to the standard setting, we also consider a natural background setting (Zhang et al.] [20TT] Nguyen al., [02I). where the background is replaced by task-irrelevant natural videos randomly sampled from the 'driving car' class in the Kinetics 400 dataset (Kay et al., 2017). Following TPC (Nguyen et al., 2021), we use two separate sets of background videos for training and evaluation. Hence, the natural background setting tests gener- alization to unseen distractions. We note that the recently released Distracting Control Suite (DCS, Stone et al. 202 serves similar purpose. However, the background distractions in DCS seem less challenging, as there are fewer videos and the ground plane is made visible for most tasks. In our preliminary experiments, Swingup Sparse in the natural background setting. We therefore switch back to Cartpole Swingup in this setting.
model and all the baselines achieved close to zero returns Cartpole
Baselines. Our main baselines are DREAMER (Hafner et al., 2021), DREAMING (Okada & Taniguchi 2021). and TPC (Nguyen al.) 202I). the cate-offthhe-s art for reconstruction-batt and reconstruction- free MBRL. In particular, TPC has shown better performance than CVRL (Ma et al.) 2020 DBC Zhang et al. (02T), and CURL (Laskin et al.. 2020a) on the same datasets. The re- cently proposed PSE (Agarwal et al. (021) has demonstrated impressive results on DCS. However, is only shown to work in the model- free setting and requires pretrained policy, while our model learns both the world model and the policy from scratch.
Implementation details. We implement our model and DREAMING based on a newer version of DREAMER while the official implementation of TPC is based on an older version. For fair
https: http:://s....t.c.c..a//////////a
119edabbebb

---

Under review as a conference paper at ICLR 2022
Cartpole Swingup Sparse
Cheetah Run
Cup Catch
1000
1000
1000
800
M
800
800
600
600
400
400
400
200
200
200
0.0
0.6
0.8
1.0
0.0
0.4
0.6
0.2
0.4
Finger Spin
Reacher Easy
Walker Run
1000
1000
1000
800
800
800
600
600
400
400
400
200
200
200
0.0
0.4
0.6
0.8
0.0
0.2
0.4
0.6
1.0
0.2
0.4
0.6
1.0
1e6
1e6
1e6
DreamerPro
Dreamer
Dreaming
TPC
Figure 2: Performance curves in standard DMC. DREAMERPRO is the only model that is comparable or better than DREAMER on all tasks. In particular, DREAMERPRO greatly outperforms DREAMER on Finger Spin and Reacher Easy, and achieves better data efficiency on Cup Catch.
Table I: Final performance in standard DMC.
Task
DREAMER
DREAMING
TPC
DREAMER PRO
Cartpole Swingup Sparse
820 + 23
830 + 12
770 9
813 + 32
Cheetah Run
840 + 74
18
782 2 82
897 + 8
Cup Catch
967 + 3
965 13
948-7
961 + 10
Finger Spin
559 + 54
722
524
811 232
Reacher Easy
721 + 51
975 2
503 - 185
873 + 127
Walker Run
737 26
422 25
222 + 29
784 + 28
comparison, we re-implement TPC based on the newer version. We adopt the default values for the DREAMER hyperparameters, except that we use continuous latents and tanh normal as the distri- bution output by the actor. We find these changes improve DREAMER'S 's performance in the standard DMC, and therefore these values for all models in both the standard and the natural background setting. Following TPC, we increase the weight of the reward loss Jh to 1000 for all models in the natural background setting to further encourage extraction of task-relevant information. While in the original TPC, this weight is chosen separately for each task from {100, 1000). we find the weight of 1000 works consistently better in our re- immmmmmmttaiiinn, which also obtains better results than reported in the original paper. We use the default batch size of 50 for DREAMER, DREAMING, and DREAMER The batch size for TPC is chosen be 150, that it has similar wall clock training time as DREAMERPRO.
Evaluation protocol. For each task, we train each model for 1M environment steps (equivalent to 500K actor steps, as the action repeat is set to 2). The evaluation return is computed every 10K steps, and averaged over 10 episodes. In all figures and tables, the mean and standard deviation are computed from 3 independent runs.
4.1 PERFORMANCE STANDARD DMC
We show the performance curves in Figure 2 and the final performance in Table I for the standard setting. DREAMERPRO is the only model that achieves comparable or even better performance than DREAMER all tasks. Notably, DREAMERPRO outperforms DREAMER by large margin

---

Under review as a conference paper at ICLR 2022
Cartpole Swingup
Cheetah Run
Cup Catch
800
600
600
400
400
200
200
200
0.0
0.2
0.0
0.2
0.4
0.6
0.8
1.0
Finger Spin
Reacher Easy
Walker Run
1000
400
750
600
300
500
400
200
250
200
100
0.0
0.2
0.4
0.0
0.2
0.4
0.6
0.8
1.0
0.6
1.0
1e6
1e6
1e6
DreamerPro
Dreamer
Dreaming
TPC
TPC-Batch-300
Figure 3: Performance curves in natural background DMC. DREAMERPRO significantly outper- forms all baselines on Cartpole Swingup, Finger Spin, and Walker Run, while DREAMER completely fails on all tasks.
Table 2: Final performance in natural background DMC.
Task
DREAMER
DREAMING
TPC
TPC-Batch-33 300
DREAMERPRO
Cartpole Swingup
126 + 16
332 + 66
521 + 80
479 + 45
671 + 42
Cheetah Run
20
334
444 = 35
477 +
349 61
Cup Catch
88 + 73
553 1 60
477 + 175
550 1
493 + 109
Finger Spin
10
209
655 155:133
:11=115
826 162
Reacher Easy
+2
-00-290
462 +-2--1300
614 + 164
641 + 123
Walker Run
35
219
161 +
136 + 17
394 33
Finger Spin and Reacher Easy, and demonstrates better data efficiency Cup Catch. We notice large variance in DREAMERPRO's performance on Finger Spin. Further investigation reveals that DREAMERPRO learned close optimal behavior (with average episode returns above 950) on two of the seeds, while converged to a suboptimal behavior (with average episode returns around 500) on the other seed. The low variance of DREAMER indicates that it hardly achieved close optimal behavior. Our results suggest for the first time that prototypical representations (and reconstruction- free representation learning in general) can be beneficial to MBRL even in the absence of strong visual distractions.
4.2 PERFORMANCE NATURAL BACKGROUND DMC
Figure 3 and Table 2 respectively show the performance curves and final evaluation returns obtained by all models in the natural background setting. DREAMER completely fails on all tasks, showing the inability of reconstruction-bunt bassed representation learning to deal with complex visual distractions. In contrast, PREAMERPEE
achieves the best performance on out tasks, with large performance gains from baselines on Cartpole Swingup, Finger Spin, and Walker Run. We additionally train TPC with a batch size of 300 (denoted TPC-Batch- Batch--00h, and DREAMER is still able to outperform it 4 out of 6 tasks. These results indicate that the advantage of prototypical representations over contrastive learning in computer vision can indeed be transferred to MBRL for better robustness to visual distractions.

---

Under review as a conference paper at ICLR 2022
Imput
DREAMERPRO
DREAMER
TPC
Visualization of learned latent states through reconstruction from auxiliary decoder. The first row shows the input images, and the remaining rows show the reconstruction from the RSSM state 2, for each model.
Figure 4:
4.3 VISUALIZATION AND ANALYSIS
To better understand how the model works and explain the performance gaps, we visualize the learned latent states through reconstruction from an auxiliary decoder (Figure 4)
Figure (Left) shows the reconstructions for Cup Catch after 100K environment steps. Note that the reconstruction of the ball is only possible from the states learned by DREAMERPRO, explaining its better data efficiency. DREAMER fails to reconstruct the ball this early stage, probably because the ball takes only few pixels, lowering its priority in the reconstruction loss.
Figure (Right) shows the reconstructions for Walker Run after 1M environment steps. We see that both DREAMER and DREAMERPRO capture some information of the background, but only DREAMERPRO able to recover the posture of the Walker. The reconstruction from TPC does not seem relevant to the Walker the background. This indicates that while TPC may be better at discarding distractors, DREAMER PRO better at retaining task-relevant information.
4.4 ABLATION STUDY
We now show the individual effect of the two loss terms, JSWAY and Tiemp in Figure 6 Here, each of the ablated versions, DreamerPro-No-SWAY and DreamerPro-N--Tenge removes one of the loss terms. We did not investigate removing both terms, as its failure has been shown in DREAMER (Hafner al., 2020). We train the ablated versions natural background DMC, and observe that both terms are necessary for achieving good performance. In particular, naively combining SWAV with DREAMER (i.c., DreamerPro-No-Temn, leads inferior performance, as it ignores the tem- poral structure. On the other hand, Tempp alone is not sufficient to provide meaningful cluster assignment targets and learning signals for the convolutional encoder.
5 RELATED WORK
Self-supervised representation learning for static images. Recent works in self-supervised learning have shown effectiveness in learning representations from high-dimensional data. CPC (van den Oord et al.) 2019) learns representations by maximizing the mutual information between the encoded representations and its future prediction using noise-contrastive estimation. SimCLR (Chen et al., 2020) shows that the contrastive data can be generated using the data in the training mini-batch by applying random augmentations. MoCo (He et al. 2020), on the other hand, improves the contrastive training by generating the representations from momentum encoder in- stead the trained network. Despite the success in some tasks, one weakness of the contrastive approaches is that it require the model compare a larger amount of samples, which demands large batch sizes or memory banks. To address this problem, some works propose to learn the image representations without discriminating between samples. Particularly, BYOL (Grill et al.
8

---

Under review as a conference paper at ICLR 2022
2020) introduces momentum encoder to provide target representations for the training network. SwAV (Caron et al. 2020) proposes to learn the embeddings by matching them to a set of learned clusters. DINO (Caron al. 202T) replaces the clusters in SwAV with categorical heads and uses the centering and sharpening technique to prevent representations collapsing. Unlike our model, these works treat each image independently and ignore the temporal structure of the environment, which is crucial in learning the forward dynamics and policy in MBRL.
Representation learning for model-free reinforcement learning. It has been shown that adopt- ing data augmentation techniques like random shifts in the observation space enables robust learn- ing from pixel input in any model-free reinforcement learning algorithm (Laskin et al., 2020bb Yarats et al.] 2021c a). Recent works have also shown that self-supervised representation learn- ing techniques can bring significant improvement to reinforcement learning methods. For example, CURL Laskiin et al., 2020a) performs contrastive learning along with off-policy RI. algorithms and shows that it significantly improves sample-efficiency and model performance over pixel-based methods. Other works aim to improve the representation learning quality by combining temporal prediction models in the representation learning process (Schwarzer et al. 2021a b Stooke et al. 2021 Yarats et al., 2021b Guo et al., 2020 Gregor et al.. 2019). However, the main purpose of the temporal prediction models in these works is mainly to obtain the abstract representations of the observations, and they are not shown to support long-horizon imagination.
Model-based reinforcement learning with reconstruction. Model based reinforcement learning from raw pixel data learn the representation space by minimizing the observation reconstruction loss. World Models (Ha & Schmidhuber (018)) learn the latent dynamics of the environment in two-stage process to evolve their linear controllers in imagination. SOLAR (Zhang et al., 2019) models the dynamics as time-varying laussr-Gaussian and solves robotic tasks via guided policy search. Dreamer (Hafner et al., 2020) jointly learns the RSSM and latent state space from obser- vation reconstruction loss. DeepMDP (Gelada et al., 2019) also propose a latent dynamics model- based method that uses bisimulation metrics and reconstruction loss in Atari However, reconstruc- tion based methods are susceptible to noise and objects irrelevant to the task in the environment (guuyen et al.j 202I). Furthermore, in a few cases, the latent representation fails to reconstruct small task objects in the environment (Okada & Taniguchi 2021).
Reinforcement learning under visual distractions. A large body of works on robust representation learning focuses on contrastive objectives. For example. CVRL (Ma et al. 2020) proposes to learn representations from complex observations by maximizing the mutual information between an im- age and its corresponding embedding using contrastive objectives. However, the learning objective of CVRI. encourages the representation model to learn as much information as possible, includ- ing task-irrelevant information. Dreaming (Okada & Taniguchi 2021) and TPC (Nguyen al.] 2021] tackle this problem by incorporating dynamic model and applying contrastive learning in the temporal dimension, which encourages the model to capture controllable and predictable infor- mation in the latent space. Bisimulation metrics method such as DBC (Zhang et al. 2021) and PSE (Agarwal al., (02T) is another type of representation learning robust visual distractions. Using the bisimulation metrics that quantify the behavioral similarity between states, these methods make the mode robust to task-irrelevant information. However, DBC cannot generalize to unseen back- grounds (Nguyen et al. (002). and PSE is only shown to work in the model-free setting and requires a pre- policy to compute the similarity metries, while our model learns both the world model and the policy from scratch.
CONCLUSION
In this work, we presented the first reconstruction-free MBRL agent based on the prototypical rep- resentation and its temporal dynamics. In experiments, we demonstrated the consistently improved accuracy and robustness of the proposed model in comparison to the Temporal Predictive Coding (TPC) agent and the Dreamer agent for both standard and natural background DMC tasks. Our results suggest that there are unexplored broad areas in reconstruction free MBRL. Interesting fu- ture directions are to apply this model on Atari games and to investigate the possibility of learning hierarchical structures such as skills without reconstruction.

---

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 in reinforcement learning. In International Conference on Learning Representations, 2021.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv: 12894, 2018.
Mathilde Caron, Ishan Mísra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Pro- cessing Systems, volume 33, pp. 9912-9924. Curran Associates, Inc., 2020.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers, 2021.
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Hal Daumé III and Aarti Singh (eds.), Proceed- ings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 13-18 Jul 2020.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In C. J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger (eds.), Advances in Neural Infor- mation Processing Systems, volume 26. Curran Associates, Inc., 2013.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G. Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2170-2179. PMLR, 09-15 Jun 2019.
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, F. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Mann-Batteen 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 H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 21271-21284. Curran Associates, Inc., 2020.
Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, Jean- ManttBatienn Grill, Florent Altché, Remi Munos, and Mohammad Gheshlaghi Azar. Bootstrap predit-pr representations for multi- task reinforcement learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research. pp. 3875-3886. PMLR, 13-18 Jul 2020.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Gameet (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2555-2566 PMLR, 09-15 Jun 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.
10

---

Under review as a conference paper at ICLR 2022
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:161..0533 2016.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint wrXiv:1903.00350 2019.
Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijaya- narasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, Mustafa Suleyman, and Andrew Zisserman. The kinetics human action video dataset, 2017.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representa- tions for reinforcement learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5639-56666 1639- PMLR, 13-18 Jul 2020a.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp- 19884-19895. Curran Associates, Inc., 2020b.
Xiao Ma, Siwei Chen, David Hsu, and Wee Sun Lee. Contrastive variational reinforcement learning for complex observations, 2020.
Tung D Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8130-8139 139. PMLR, 18-24 Jul 2021.
Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-bases reinforcement learning by latent imagination without reconstruction, 2021.
Julian Schrittweser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-609 2020.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation, 2018.
Max Schwarzer, Ankesh Anand, Rishab Goel, Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self- predicceeee representations. In International Conference on Learning Representations, 2021a.
Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, De- von Hjelm, Philip Bachman, and Aaron Courville. Pretraining representations for data-efficient reinforcement learning, 2021b.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite challenging benchmark for reinforcement learning from pixels, 2021.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 9870-9879. PMLR, 18- Jul 2021.

---

Under review as a conference paper at ICLR 2022
Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting, ACM Sigart Bulletin, 2(4):160-163, 1991.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S Sutton, Joseph Modayil, Michael Delp. Thomas Degris, Patrick M Pilarski, Adam White, and Doina Precup. Horde: A scalable real-time architecture for learning knowledge from unsu- pervised sensorimotor interaction. In The 10th International Conference Autonomous Agents and Multiagent Destems-Volumme 2, pp. 761-768, 2011.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap, and Martin Ried- miller. Deepmind control suite, 2018.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding, 2019.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous con- trol: Improved data-augmented reinforcement learning, 2021a.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with pro- totypical representations. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th In- ternational Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 11920-11931. PMLR, 18-24 Jul b021
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, 2021c.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learn- ing invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representattons, 2021.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine.
SOLAR:
Deep structured representations for model-based reinforcement learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 7444-7453. PMLR, 09-15 Jun 2019.
A HYPERPARAMETERS
For hyperparameters that are shared with DREAMER, we use the default values suggested in the config file in the official implementation of DREAMER, with the following two exceptions. We set rssm. discrete
False and actor dist tanh-normal, as we find these changes improve performance over the default setting. The additional hyperparameters introduced DREAMER PRO are listed in Table We find it helpful to freeze the prototypes for the first 10K gradient updates. In the natural background setting, we add squared loss that encourages the norm of projections (before (z-oomaaizzatio) to be close to 1. This helps stabilize the model.
Table Additional hyperparameters in DREAMERPRO.
Hyperparameter
Value
Number of prototypes K
2500
Prototype dimension
32
Softmax temperature т
0.1
Sinkhorn iterations
3
Sinkhorn epsilon
0.0125
Momentum update fraction r)
0.05
12

---

Under review as a conference paper at ICLR 2022
B NEAREST NEIGHBOR QUERIES IN LATENT SPACE
Query
DREAMERPRO
DREAMER
TPC
Figure 5: Visualization of learned latent states through nearest neighbor queries.
We sample a batch of trajectories from the training replay buffer, and obtain the latent state for each image. Then, given a query image, we show the three images in the batch whose latent states are the closest the query image. We use the same batch and same query images for all models.
For TPC. the nearest neighbors tend to contain different backgrounds, but the agent's 's states can also be very different (See Cup Catch and Walker Run). On the other hand, the nearest neighbors for DREAMER PRO tend to have similar backgrounds and also similar agent states.
Our results suggest that DREAMERPRO and TPC work in quite different ways. DREAMERPRO tries to retain task-relevant information at the cost of also including some distractors, while TPC focuses more on discarding distractors. An interesting future direction is to simultaneously consider these two factors and achieve better balance.
13

---

Under review as a conference paper at ICLR 2022
C
ABLATION RESULTS
Cartpole Swingup
Cheetah Run
Cup Catch
800
400
600
600
400
400
200
200
0.0
0.2
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.8
1.0
Finger Spin
Reacher Easy
Walker Run
1000
400
750
300
500
200
250
100
0.0
0.2
0.4
0.6
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
1e6
DreamerPro
DreamerPrroNo-SWAN
DreamerPro-No-Temp
Figure Ablation study. Both JSwAv and Tiemp are necessary for achieving good performance.
D ATARI RESULTS
Bank Heist
Enduro
Name This Game
15000
1000
2000
10000
500
1000
5000
0.0
0.5
1.0
1.5
2.0
2.5
0
3
4
2
4
8
1e7
1e7
1e7
Skiing
Seaquest
Robotank
-10000
15000
60
10000
-20000
40
20
-30000
0
1e7
1e7
1e7
DreamerPro
Dreamer
IQN@200M
Rainbow@200M
Figure 7: Performance curves in six Atari games.
To show the potential of DREAMERPRO to leverage the benefits of world models and discrete latents in complex environments. we train DREAMERPRO and 80M environment steps depending on the convergence speed of DREAMER on these games. We freeze the prototypes for 30K gradient updates. The weights for reward and KL losses are 100 and 1, respectively. We additionally replace the linear projection layer from 2, to 2, by an MLP with one hidden layer of size 600. All other hyperparameters are kept default.
subset of six Atari games for 25M, 50M.
14

---

Under review as a conference paper at ICLR 2022
On 4 out of the 6 games, DREAMERPRO obtains comparable performance DREAMER, matching or surpassing model-free baselines that are trained for 200M environment steps. We note that this is the first time a reconstruction-free MBRL agent shows promising results Atari with discrete latents.
E PROTOTYPE VISUALIZATIONS
We visualize the first 5 prototypes learned for each task using nearest neighbor queries, shown in each row of Figures 8 13 To do so, sample batch of trajectories from the training replay buffer, and obtain the projection It from latent state & for each image. Then, given prototype cu query, show the images in the batch whose projections are the closest the prototype.
TIU B
RE 55
Figure 8: Prototype visualization for Cartpole Swingup.
R
M
Techs
Figure 9: Prototype visualization for Cheetah Run.
15

---

Under review as a conference paper at ICLR 2022
Figure 10: Prototype visualization for Cup Catch.
Drie
Figure I1: Prototype visualization for Finger Spin.
Figure 12: Prototype visualization for Reacher Easy.
16

---

Under review as a conference paper at ICLR 2022
FO
FO
STOR
FO
FO CE
Figure 13: Prototype visualization for Walker Run.
17


Reviews:
[{"review_id": "z_zxf1Kv50g", "reviewer": "Reviewer_VAKJ", "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\n1. The paper is well-written and easy to follow. \n2. The idea is straightforward and the experimental results are strong in DMC suite. \n3. Introducing $\\mathcal{L}_\\text{temp}$ is a clean solution to integrating the temporal notion to the original SWAV objective. \n\nWeaknesses\n1. The major concern is the complexity of the environments. The results are evaluated in DMC, and in the 6 environments the proposed method outperforms DREAMER in terms of converging performance only in ReacherEasy. Further evaluation on Atari will make the results much more convincing.\n2. The paper is not particularly strong in terms of novelty as swapping the reconstruction objective with prototypical objective is already shown to be effective in model-free RL (Yarats et al. 2021). This won't be a big problem if the paper can showcase effectiveness in more complex environments. \n\n\nMinor\n1. The batch size $B$ in all experiments seems not specified in the paper. \n", "summary_of_the_paper": "This work proposes a reconstruction free method for pixel-based MBRL, utilizing cluster assignment adapted from SWAV with an additional temporal term. The method learns a latent space that discards task-irrelevant visual details and is more efficient compared to contrastive methods. ", "summary_of_the_review": "I think it is a very reasonable approach to combine cluster assignment to MBRL and the paper does a good job addressing the temporal consistency, a challenge that is not presented in model-free RL. However, I am not confident how the method could scale to higher complexity environments. I would be happy to increase my score if the above concern is addressed. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Zv3-2aSU1Oe", "reviewer": "Reviewer_xtkN", "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": "###  Strengths\n\n- The approach is clear and well motivated, and a natural extension of representation learning for MBRL\n- Related work section is thorough and well written\n- Experimental results show an increased robustness to visual distractions compared to dreamer\n\n###  Areas to improve\n\nThe experiments are not as thorough as they could be.\n- Several baselines are cited but not included in experiments (ie PSE, CURL, many of the model-free temporal prediction models).  I understand that the paper is not attempting to contribute a model-free vs model-based comparison, but it makes the results seem thin.\n- Three random seeds seems too few\n- One of the main stated motivations of using a non-contrastive representation approach was to reduce the required batch size, but there is little evidence here with which to evaluate that claim.\n- I like the ablation study comparing the different terms of the additional loss (Figure 4).  The paper would be stronger with some additional exploration of the design space of the presented idea.  For example, is the approach sensitive to the hyperparameters in Table 3?  Another example, for the temporal consistency loss, could you also include terms for enforcing the consistency across timesteps and augmentations simultaneously, eg w^(1)log u^(2)?\n\nI found the diagram in Figure 1 a bit unclear because it ignores the augmentation index i, and has AUG(o_t) repeated twice.  What do the two augmentations in the diagram represent, if both sets of augmentations are actually passed through all paths of the diagram?\n\n\n## Update after rebuttal\n\nThank you for addressing my concerns.  I will maintain my score that this paper is above the acceptance threshold.", "summary_of_the_paper": "The paper presents DreamerPro, a method for learning a latent dynamics model from image observations without a pixel-based reconstruction loss.  The model combines the clustering approach of SwAV with an additional temporal consistency loss.  Experimental results show improved performance over baselines in distraction-free and natural-video-background settings on DeepMind Control Suite tasks.\n", "summary_of_the_review": "Overall, I found the paper to be a novel exploration of a natural combination of two methods from MBRL and SSL.  The paper makes modest claims, and supports them with good evidence and a well-written explanation.  Although the results show empirical improvements, the paper could do more to contribute to a more general understanding of these improvements and the approach itself.", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "KT_zI1_r_yM", "reviewer": "Reviewer_pipU", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The core novelty of the paper is to combine Dreamer approach to model-based reinforcement learning in pixel space with prototypical based method for learning representations from SwAV. This led to a method that can learn a good enough representation for planning, without reconstructing the observed images. Removing reconstruction is generally well motivated, however, the paper fails to support some of its key claims.\n\n1. Unclear motivation and focus. At the very high level, the paper mentions two limitations for reconstruction-based methods (in intro). First, they are computationally expensive, and second, they learn signals irrelevant to the task. In the results, however, the paper entirely ignores one and focuses on the second i.e. there is no talk or analysis on the computational aspect of the proposed method and it seems the authors' main focus was to get better results on environment with random video backgrounds. \n\n2. Weak baselines, particularly for contrastive methods. The paper is quite vocal about the limitations of contrastive representation learning and why the authors chose to use prototypical based method instead. However, the only contrastive baseline is TPC which is not representative of this group of models. e.g. the Dreamer paper itself also proposes a contrastive loss for training the latent model although it is not temporal. There is also Dreaming (https://arxiv.org/abs/2007.14535) and CVRL (https://arxiv.org/abs/2008.02430). Moreover, contrastive and prototypical representation learning are not the only solutions to this problem e.g. there are already existing bisimulation methods which address this exact problem (https://arxiv.org/abs/2006.10742). \n\n3. Lack of analysis. Unfortunately, the paper does not provide any analysis on *how* the model outperforms Dreamer. More strangely, the paper claims better sample efficiency vs Dreamer even in absence of distraction. This is strange because reconstruction of a basic static background does not \"waste the representation capacity\" as argued by the author. Intuitively, it can be assumed that it is mainly because the model learnt a more robust representation, however, there is no explicit theory or experiment to support this claim. For examples of such analysis please refer to Fig 5a in TPC paper. \n\n===== Other issues \n1. The ablation study (Fig 4) is poorly described. It's not even clear which set of environments (distracting or not distracting) the model is tested on and why a particular subset of original six tasks is visualized. \n2. The first reference to SwAV (3rd paragraph of page 2) is not cited.\n3. In page 4, it is mentioned that the number of prototypes K is set to B x T. but it is not clear what T is. Is T the time step? So the number of prototypes grow by the time and the model is approximating millions of prototypes at the end? \n4. The reported numbers for TPC does not match the reported numbers in the original paper. For example, on Walker Run, TPC reported ~400 while this paper reported ~200. Is this because of the re-implementation? ", "summary_of_the_paper": "The paper proposes DreamerPro, a clustering based (i.e. prototype-based) version of Dreamer, which learns the latent representations by prototypical representation learning instead of image reconstruction. In other words, the proposed method is Dreamer without reconstruction replaced by SwAV template prototypical method. The authors test their method on a subset of six DeepMind Control tasks from pixels, with an without distracting backgrounds, to demonstrate that it is capable of outperforming Dreamer and TPC in almost all cases. ", "summary_of_the_review": "The proposed method is a combination of two previously established methods, Dreamer and SwAV. This combination and its applications in MBRL is novel, however, I don't find it to be particularly significant. Reconstruction free MBRL from pixels is not a new problem and the paper does not provide enough theoretical (nor empirical) insights as why a prototype-based method is a better fit for such setting vs other methods. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "79pMk0jFc0Q", "reviewer": "Reviewer_FvUb", "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": "=== Strengths ===\n\n+ The paper is well written and motivated, and the presentation is easy to follow. \n+ The presented approach shows convincing experimental results on the selected benchmark.\n\n=== Weaknesses ===\n\n- I would like to see a more thorough experimental analysis. The latest Dreamer implementation also evaluates all tasks in the Atari benchmark. I would like to see how DreamerPro compares to the basic Dreamer on the Atari benchmark as well.\n- I would like to see some visualizations on the learned features using the prototype learning loss. For example, a side-by-side comparison of t-SNE or nearest neighbor plots of the image observations with their representations learned from 1) the original reconstruction loss and 2) the proposed prototype learning loss. I would also like to visualize what the learned prototypes correspond to in the observation space.\n- The authors mention that they use continuous latent space for both Dreamer and DreamerPro as they found it to perform better. However, according to Dreamerv2, discrete latent space is one of the major improvements they found and it outperformed continuous latent space. I would appreciate it if the authors can further clarify this in the response.\n- I am concerned about the novelty of this paper, because it does seem like a straightforward combination of Dreamer and SWAV, with one additional novelty which is the proposed L_temp loss. This feels somewhat incremental to me.\n\n### Post rebuttal updates\nThe authors did a good job in rebuttal and with the additional experiments and visualizations, the paper is now in a much better shape in my honest opinion. I am updating my score from 3 to 5 to reflect the changes. While the paper proposes a valid combination of existing techniques (SwAV + Dreamer) with reasonable performance on the selected task of \"RL with random visual distractions\", to me this paper is not significant or novel enough to make a strong case for acceptance.", "summary_of_the_paper": "This paper proposes DreamerPro, an extension to Dreamer where it swaps the reconstruction loss with a prototype learning loss based on SWAV. At each optimization step, DreamerPro augments each batch of observations into two using different augmentations. It uses the Sinkhorn-Knopp algorithm to compute the target clusters for each batch, with the goal of   encouraging uniform clustering assignment within each augmented batch. Intuitively, the prototype learning loss rewards prediction accuracy both from across the views, which adapts from the original SWAV approach, and within the same view, which the authors mention help build temporal information into the learned prototypes. DreamerPro outperforms the original Dreamer on most of the tasks on the Deepmind Control benchmark, and drastically improves on the Deepmind Control benchmark with noisy backgrounds.\n", "summary_of_the_review": "The paper is well-written and the presented approach shows promising results on the selected benchmark. However, it seems that the paper lacks novelty at the moment, and I would like to see a more thorough experimental section with more visualizations. Based on these, I do not recommend acceptance at this moment.", "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": {}}]

Author rebuttal:

Meta-review: The authors propose an alteration to Dreamer that incorporates a swav-like objective. The reviewers raised a number of issues with the paper, overall arguing for rejection. In particular, the reviewers felt that the work was not well motivated, weak performance, that a number of baselines were missing, and a lack of analysis of the results, a lack of novelty. While the authors addressed many of these concerns during the rebuttal, the majority of reviewers still felt this was not enough and that the paper did not meet the bar for acceptance. Therefore, I recommend rejection at this stage so that these concerns can be addressed.
Decision: Reject