Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 3: reject, not good enough
review_content: 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. 
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.

1. 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. 

2. 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). 

3. 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. 

===== Other issues 
1. 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. 
2. The first reference to SwAV (3rd paragraph of page 2) is not cited.
3. 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? 
4. 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_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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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.

main_review: === Strengths ===

+ The paper is well written and motivated, and the presentation is easy to follow. 
+ The presented approach shows convincing experimental results on the selected benchmark.

=== Weaknesses ===

- 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.
- 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.
- 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.
- 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.

### Post rebuttal updates
The 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_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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.

main_review: ###  Strengths

- The approach is clear and well motivated, and a natural extension of representation learning for MBRL
- Related work section is thorough and well written
- Experimental results show an increased robustness to visual distractions compared to dreamer

###  Areas to improve

The experiments are not as thorough as they could be.
- 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.
- Three random seeds seems too few
- 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.
- 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)?

I 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?


## Update after rebuttal

Thank you for addressing my concerns.  I will maintain my score that this paper is above the acceptance threshold.
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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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. 
main_review: Strengths
1. The paper is well-written and easy to follow. 
2. The idea is straightforward and the experimental results are strong in DMC suite. 
3. Introducing $\mathcal{L}_\text{temp}$ is a clean solution to integrating the temporal notion to the original SWAV objective. 

Weaknesses
1. 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.
2. 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. 


Minor
1. The batch size $B$ in all experiments seems not specified in the paper. 

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.
review_title: null

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