DREAMERPRO: RECONSTRUCTION-FREE MODEL-BASED REINFORCEMENT LEARNING
Year: 2021

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 contrastively learn the world model, but the performance tends to be inferior in the absence of distractions. In this paper, we seek to enhance robustness to distractions for MBRL agents. Specifically, we consider incorporating prototypical representations, which have yielded more accurate and robust results than contrastive approaches in computer vision. However, it remains elusive how prototypical representations can benefit temporal dynamics learning in MBRL, since they treat each image independently without capturing temporal structures. To this end, we propose to learn the prototypes from the recurrent states of the world model, thereby distilling temporal structures from past observations and actions into the prototypes. The resulting model, DREAMERPRO, successfully combines DREAMER with prototypes, making large performance gains on the DeepMind Control suite when there are complex background distractions, while maintaining similar performance as DREAMER in the standard setting.

heading: 1 INTRODUCTION
text: Model-Based Reinforcement Learning (MBRL, Sutton & Barto, 2018; Sutton, 1991) provides a solution 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 provides a 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-based MBRL, DREAMER is particularly notable as the first MBRL model outperforming popular model-free RL algorithms with better sample efficiency in both continuous control (Hafner et 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 in 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 is to learn low-dimensional state representation and, in the model-based case, also its forward model. Although we can learn such representation directly by maximizing the rewards (Schrittwieser et al., 2020), it is usually very slow to do this due to the reward sparsity. Instead, it is more practical to introduce auxiliary tasks providing richer learning 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 dynamics model in a way to reduce the reconstruction error of the observed sequences. However, reconstruction-based representation learning has limitations. First, it is computationally expensive to reconstruct the high-dimensional inputs, especially in models like DREAMER that needs to reconstruct long-range videos. Second, it wastes the representation capacity to learn even the visual signals that are irrelevant to the task or unpredictable such as noisy background (Burda et al., 2018). Thus, in MBRL it is of particular interest to realize a version of DREAMER without reconstruction.
Recently, there have been remarkable advances in reconstruction-free representation learning in reinforcement 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 requires a large batch size (so computationally expensive) to perform accurately and robustly. An alternative is the clustering-based or prototype-based approach (Caron et al., 2020). By learning a 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 in many applications (Caron et al., 2020; 2021; Yarats et al., 2021b) than the contrastive method while also alleviating the need for maintaining a large batch size. The prototype structure can also be used to implement an exploration method (Yarats et al., 2021b).
However, for reconstruction-free MBRL only the contrastive approach like Temporal Predictive Coding (TPC, Nguyen et al., 2021) 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, we 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-free MBRL agent, called DREAMERPRO, by combining the prototypical representation learning with temporal dynamics learning. Similar to SwAV (Caron et 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 ‘reconstruct’ the cluster assignment of the observation, thereby relieving the world model from modeling lowlevel 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-free MBRL agent based on the prototypical representation and its temporal dynamics and (2) the demonstration of the consistently improved accuracy and robustness of the proposed model in comparison to a contrastive reconstructionfree MBRL agent and Dreamer for both standard and natural background DMC tasks.

heading: 2 PRELIMINARIES
text: 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.

heading: 2.1 RECONSTRUCTION-BASED WORLD MODEL LEARNING
text: 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 t, the agent receives an image observation ot and a scalar reward rt (obtained by previous actions a<t). The agent then chooses an action at based on its policy. The RSSM models the observations, rewards, and transitions through a probabilistic generative process:
p(o1:T , r1:T | a1:T ) = ∫ T∏
t=1
p(ot | s≤t, a<t) p(rt | s≤t, a<t) p(st | s<t, a<t) ds1:T (1)
= ∫ T∏ t=1 p(ot | ht, st) p(rt | ht, st) p(st | ht) ds1:T , (2)
where the latent variables s1:T are the agent states, and ht = GRU(ht−1, st−1, at−1) is a deterministic encoding of s<t and a<t. To infer the agent states from past observations and actions, a variational encoder is introduced:
q(s1:T | o1:T , a1:T ) = T∏ t=1 q(st | s<t, a<t, ot) = T∏ t=1 q(st | ht, ot) . (3)
The training objective is to maximize the evidence lower bound (ELBO):
JDREAMER = T∑ t=1
Eq[log p(ot | ht, st)︸ ︷︷ ︸ J tO + log p(rt | ht, st)︸ ︷︷ ︸ J tR −DKL(q(st | ht, ot) ‖ p(st | ht))︸ ︷︷ ︸ J tKL ] . (4)

heading: 2.2 POLICY LEARNING BY LATENT IMAGINATION
text: DREAMER interleaves policy learning with world model learning. During policy learning, the world model is fixed, and an actor and a critic are trained cooperatively from the latent trajectories imagined by the world model. Specifically, the imagination starts at each non-terminal state ẑt = [ht, st] encountered during world model learning. Then, at each imagination step t′ ≥ t, an action is sampled from the actor’s stochastic policy: ât′ ∼ π(ât′ | ẑt′). The corresponding reward r̂t′+1 and next state ẑt′+1 are predicted by the learned world model. Given the imagined trajectories, the actor improves its policy by maximizing the λ-return (Sutton & Barto, 2018; Schulman et al., 2018) plus an entropy regularizer that encourages exploration, while the critic is trained to approximate the λ-return through a squared loss.

heading: 3 DREAMERPRO
text: To compute the DREAMER training objective, more specifically J tO in Equation 4, a decoder is required to reconstruct the image observation ot from the state zt = [ht, st]. Because this reconstruction loss operates in pixel space where all pixels are weighted equally, DREAMER tends to allocate most of its capacity to modeling complex visual patterns that cover a large pixel area (e.g., backgrounds). This leads to poor task performance when those visual patterns are task irrelevant, as shown in previous work (Nguyen et al., 2021).
Fortunately, during policy learning, what we need is accurate reward and next state prediction, which are respectively encouraged by J tR and J tKL. In other words, the decoder is not required for policy learning. The main purpose of having the decoder and the associated loss J tO, as shown in DREAMER, is to learn meaningful representations that cannot be obtained by J tR and J tKL alone. The above observations motivate us to improve robustness to visual distractions by replacing the reconstruction-based representation learning in DREAMER with reconstruction-free 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., 2019; Chen et 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 and policy learning (in wall clock time). Therefore, we propose to combine DREAMER with the 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.
DREAMERPRO uses the same policy learning algorithm as 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 1 for an illustration.
Concretely, given a sequence of observations o1:T sampled from the replay buffer, we obtain two augmented views o(1)1:T , o (2) 1:T by applying random shifts (Laskin et al., 2020b; Yarats et al., 2021c) with bilinear interpolation (Yarats et al., 2021a). We ensure that the augmentation is consistent across time steps. Each view i ∈ {1, 2} is fed to the RSSM to obtain the states z(i)1:T . To predict the cluster assignment from z(i)t , we first apply a linear projection followed by `2-normalization to obtain a vector x(i)t of the same dimension as the prototypes, and then take a softmax over the dot products of x(i)t and all the prototypes:
(u (i) t,1, . . . , u (i) t,K) = softmax
( x
(i) t · c1 τ , . . . , x (i) t · cK τ
) . (5)
<latexit sha1_base64="KrmYHDjyuqZmm4bzKsDv5hOoqUo=">AAAB9HicbVA9SwNBEN2LXzF+RS1tFoNgFe5SqJUEbCwjmA9IjrC3mUuW7O2du3OBcOR32FgoYuuPsfPfuEmu0MQHA4/3ZpiZFyRSGHTdb6ewsbm1vVPcLe3tHxwelY9PWiZONYcmj2WsOwEzIIWCJgqU0Ek0sCiQ0A7Gd3O/PQFtRKwecZqAH7GhEqHgDK3kN3QcsEBIgQJMv1xxq+4CdJ14OamQHI1++as3iHkagUIumTFdz03Qz5hGwSXMSr3UQML4mA2ha6liERg/Wxw9oxdWGdAw1rYU0oX6eyJjkTHTKLCdEcORWfXm4n9eN8Xwxs+ESlIExZeLwlRSjOk8AToQGjjKqSWMa/s5p3zENONocyrZELzVl9dJq1b1rqreQ61Sv83jKJIzck4uiUeuSZ3ckwZpEk6eyDN5JW/OxHlx3p2PZWvByWdOyR84nz8N2JJF</latexit>
Here, u(i)t,k is the predicted probability that state z (i) t maps to cluster k, τ is a temperature parameter, and the prototypes {c1, . . . , cK} are also `2-normalized.
Analogously, to predict the cluster assignment from an augmented observation o(i)t , we feed it to a convolutional encoder (shared with the RSSM), apply a linear projection followed by `2- normalization, and obtain a vector y(i)t . We summarize this process as: y (i) t = fθ(o (i) t ), where θ collectively denotes the parameters of the convolutional encoder and the linear projection layer. The prediction probabilities are again given by a softmax:
(v (i) t,1, . . . , v (i) t,K) = softmax
( y
(i) t · c1 τ , . . . , y (i) t · cK τ
) , (6)
where v(i)t,k is the predicted probability that observation o (i) t maps to cluster k.
To obtain the targets for the above two predictions (i.e., Equations 5 and 6), we apply the SinkhornKnopp algorithm (Cuturi, 2013) to the cluster assignment scores computed from the output of a momentum encoder fθ̄ (He et al., 2020; Grill et al., 2020; Caron et al., 2021), whose parameters θ̄ are updated using the exponential moving average of θ: θ̄ ← (1 − η)θ̄ + ηθ. For each observation o
(i) t , the scores are given by the dot products (ȳ (i) t · c1, . . . , ȳ(i)t · cK), where ȳ(i)t = fθ̄(o(i)t ) is the momentum encoder output. The Sinkhorn-Knopp algorithm is applied to the two augmented batches {o(1)1:T }, {o (2) 1:T } separately to encourage uniform cluster assignment within each augmented batch and avoid trivial solutions. We specifically choose the number of prototypes K = B × T , where B is the batch size, so that the observation embeddings are implicitly pushed apart from each other. The outcome of the Sinkhorn-Knopp algorithm is a set of cluster assignment targets (w
(i) t,1, . . . , w (i) t,K) for each observation o (i) t .
Now that we have the cluster assignment predictions and targets, the representation learning objective is simply to maximize the prediction accuracies:
J tSWAV = 1
2 K∑ k=1 ( w (1) t,k log v (2) t,k + w (2) t,k log v (1) t,k ) , (7)
J tTemp = 1
2 K∑ k=1 ( w (1) t,k log u (1) t,k + w (2) t,k log u (2) t,k ) . (8)
Here, J tSWAV improves prediction from an augmented view. This is the same loss as used in SWAV (Caron et al., 2020), and is 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, J tTemp, 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, J tTemp is similar toJ tO in the sense that we are 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 J tO in Equation 4 with J tSWAV + J tTemp:
JDREAMERPRO = T∑ t=1 Eq[J tSWAV + J tTemp + J tR − J tKL] , (9)
where J tR and J tKL are now averaged over the two augmented views.

heading: 4 EXPERIMENTS
text: Environments. We evaluate our model and the baselines on six image-based continuous control tasks from the DeepMind Control (DMC) suite (Tassa et 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., 2021; Nguyen et al., 2021), 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 generalization to unseen distractions. We note that the recently released Distracting Control Suite (DCS, Stone et al., 2021) serves a 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, our model and all the baselines achieved close to zero returns on Cartpole Swingup Sparse in the natural background setting. We therefore switch back to Cartpole Swingup in this setting.
Baselines. Our main baselines are DREAMER (Hafner et al., 2021), DREAMING (Okada & Taniguchi, 2021), and TPC (Nguyen et al., 2021), the state-of-the-art for reconstruction-based and reconstruction-free MBRL. In particular, TPC has shown better performance than CVRL (Ma et al., 2020), DBC (Zhang et al., 2021), and CURL (Laskin et al., 2020a) on the same datasets. The recently proposed PSE (Agarwal et al., 2021) has demonstrated impressive results on DCS. However, it is only shown to work in the model-free setting and requires a 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 DREAMER1, while the official implementation of TPC2 is based on an older version. For fair
1https://github.com/danijar/dreamerv2/tree/e783832f01b2c845c195587158c4 e129edabaebb
2https://github.com/VinAIResearch/TPC-tensorflow
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 distribution output by the actor. We find these changes improve DREAMER’s performance in the standard DMC, and therefore use these values for all models in both the standard and the natural background setting. Following TPC, we increase the weight of the reward loss J tR 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-implementation, which also obtains better results than reported in the original paper. We use the default batch size of 50 for DREAMER, DREAMING, and DREAMERPRO. The batch size for TPC is chosen to be 150, so 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.

heading: 4.1 PERFORMANCE IN STANDARD DMC
text: We show the performance curves in Figure 2 and the final performance in Table 1 for the standard setting. DREAMERPRO is the only model that achieves comparable or even better performance than DREAMER on all tasks. Notably, DREAMERPRO outperforms DREAMER by a large margin on
Finger Spin and Reacher Easy, and demonstrates better data efficiency on Cup Catch. We notice a large variance in DREAMERPRO’s performance on Finger Spin. Further investigation reveals that DREAMERPRO learned close to 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 to optimal behavior. Our results suggest for the first time that prototypical representations (and reconstructionfree representation learning in general) can be beneficial to MBRL even in the absence of strong visual distractions.

heading: 4.2 PERFORMANCE IN NATURAL BACKGROUND DMC
text: 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-based representation learning to deal with complex visual distractions. In contrast, DREAMERPRO achieves the best performance on 4 out of 6 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-300), and DREAMERPRO is still able to outperform it on 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.

heading: 4.3 VISUALIZATION AND ANALYSIS
text: 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 4 (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 at this early stage, probably because the ball takes only a few pixels, lowering its priority in the reconstruction loss.
Figure 4 (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 is able to recover the posture of the Walker. The reconstruction from TPC does not seem relevant to the Walker or the background. This indicates that while TPC may be better at discarding distractors, DREAMERPRO is better at retaining task-relevant information.

heading: 4.4 ABLATION STUDY
text: We now show the individual effect of the two loss terms, J tSWAV and J tTemp, in Figure 6. Here, each of the ablated versions, DreamerPro-No-SwAV and DreamerPro-No-Temp, removes one of the loss terms. We did not investigate removing both terms, as its failure has been shown in DREAMER (Hafner et al., 2020). We train the ablated versions in natural background DMC, and observe that both terms are necessary for achieving good performance. In particular, naively combining SWAV with DREAMER (i.e., DreamerPro-No-Temp) leads to inferior performance, as it ignores the temporal structure. On the other hand, J tTemp alone is not sufficient to provide meaningful cluster assignment targets and learning signals for the convolutional encoder.

heading: 5 RELATED WORK
text: Self-supervised representation learning for static images. Recent works in self-supervised learning have shown its 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 a momentum encoder instead of the trained network. Despite the success in some tasks, one weakness of the contrastive approaches is that it require the model to 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.,
2020) introduces a 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 et al., 2021) 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 adopting data augmentation techniques like random shifts in the observation space enables robust learning from pixel input in any model-free reinforcement learning algorithm (Laskin et al., 2020b; Yarats et al., 2021c;a). Recent works have also shown that self-supervised representation learning techniques can bring significant improvement to reinforcement learning methods. For example, CURL (Laskin et al., 2020a) performs contrastive learning along with off-policy RL 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 can learn the representation space by minimizing the observation reconstruction loss. World Models (Ha & Schmidhuber, 2018) learn the latent dynamics of the environment in a two-stage process to evolve their linear controllers in imagination. SOLAR (Zhang et al., 2019) models the dynamics as time-varying linear-Gaussian and solves robotic tasks via guided policy search. Dreamer (Hafner et al., 2020) jointly learns the RSSM and latent state space from observation reconstruction loss. DeepMDP (Gelada et al., 2019) also propose a latent dynamics modelbased method that uses bisimulation metrics and reconstruction loss in Atari. However, reconstruction based methods are susceptible to noise and objects irrelevant to the task in the environment (Nguyen et al., 2021). Furthermore, in a few cases, the latent representation fails to reconstruct small task-relevant 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 image and its corresponding embedding using contrastive objectives. However, the learning objective of CVRL encourages the representation model to learn as much information as possible, including task-irrelevant information. Dreaming (Okada & Taniguchi, 2021) and TPC (Nguyen et al., 2021) tackle this problem by incorporating a dynamic model and applying contrastive learning in the temporal dimension, which encourages the model to capture controllable and predictable information in the latent space. Bisimulation metrics method such as DBC (Zhang et al., 2021) and PSE (Agarwal et al., 2021) is another type of representation learning robust to 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 backgrounds (Nguyen et al., 2021), and PSE is only shown to work in the model-free setting and requires a pre-trained policy to compute the similarity metrics, while our model learns both the world model and the policy from scratch.

heading: 6 CONCLUSION
text: In this work, we presented the first reconstruction-free MBRL agent based on the prototypical representation 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 future directions are to apply this model on Atari games and to investigate the possibility of learning hierarchical structures such as skills without reconstruction.

heading: A HYPERPARAMETERS
text: 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 in DREAMERPRO are listed in Table 3. We find it helpful to freeze the prototypes for the first 10K gradient updates. In the natural background setting, we add a squared loss that encourages the `2-norm of projections (before `2-normalization) to be close to 1. This helps stabilize the model.
B NEAREST NEIGHBOR QUERIES IN LATENT SPACE
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 to 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 states can also be very different (See Cup Catch and Walker Run). On the other hand, the nearest neighbors for DREAMERPRO 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 a better balance.

heading: C ABLATION RESULTS

heading: D ATARI RESULTS
text: To show the potential of DREAMERPRO to leverage the benefits of world models and discrete latents in complex environments, we train DREAMERPRO on a subset of six Atari games for 25M, 50M, 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 zt to xt by an MLP with one hidden layer of size 600. All other hyperparameters are kept as default.
On 4 out of the 6 games, DREAMERPRO obtains comparable performance to 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 on Atari with discrete latents.

heading: E PROTOTYPE VISUALIZATIONS
text: 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, we sample a batch of trajectories from the training replay buffer, and obtain the projection xt from latent state zt for each image. Then, given a prototype ck as query, we show the ten images in the batch whose projections are the closest to the prototype.

Review summaries:
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. 

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. 

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.