Original: D REAMER P RO : R ECONSTRUCTION -F REE M ODEL -B ASED R EINFORCEMENT L EARNING WITH P ROTOTYPICAL R EPRESENTATIONS
Revised: D REAMER P RO : R ECONSTRUCTION -F REE M ODEL -B ASED R EINFORCEMENT L EARNING WITH P ROTOTYPICAL R EPRESENTATIONS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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 visualdistractions or generalize to unseen distractions.
Revised: Top-performing Model-Based Reinforcement Learning (MBRL) agents, such as D REAMER , learn the world model by reconstructing the image observations. Hence, they often fail to discard task-irrelevant details and struggle to handle visual distractions.

Original: 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.
Revised: 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.

Original: 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 D REAMER P RO , to achieve this goal. Similar to S W AV, by encouraging uniform cluster assignment across the batch, we implicitly push apart the embeddings of different observations.
Revised: Speciﬁcally, 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 beneﬁt temporal dynamics learning in MBRL, since they treat each image independently without capturing temporal structures.

Original: 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.
Revised: 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.

Original: We evaluate our model on thestandard 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.
Revised: The resulting model, D REAMER P RO , successfully combines D REAMER with prototypes, making large performance gains on the DeepMind Control suite when there are complex background distractions, while maintaining similar performance as D REAMER in the standard setting.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Model-Based Reinforcement Learning (MBRL) (Sutton & Barto, 2018; Sutton, 1991) provides a solution to many problems in contemporary reinforcement learning.
Revised: Model-Based Reinforcement Learning (MBRL, Sutton & Barto, 2018; Sutton, 1991) provides a solution to many problems in contemporary reinforcement learning.

Original: It improves sample efﬁciency by training a policy through simulations of a learned world model.
Revised: It improves sample efﬁciency by training a policy through simulations of a learned world model.

Original: Learning a world model also provides a way to efﬁciently represent experience data as general knowledge simulatable and reusable in arbitrary downstream tasks.
Revised: Learning a world model also provides a way to efﬁciently represent experience data as general knowledge simulatable and reusable in arbitrary downstream tasks.

Original: In addition, it allows accurate and safe decisions via planning.
Revised: In addition, it allows accurate and safe decisions via planning.

Original: Among recent advances in image-based MBRL, D REAMER is particularly notable as the ﬁrst MBRL model outperforming popular model-free RL algorithms with better sample efﬁciency in both continuous control (Hafner et al., 2020) and discrete control (Hafner et al., 2021).
Revised: Among recent advances in image-based MBRL, D REAMER is particularly notable as the ﬁrst MBRL model outperforming popular model-free RL algorithms with better sample efﬁciency in both continuous control (Hafner et al., 2020) and discrete control (Hafner et al., 2021).

Original: 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.
Revised: 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.

Original: Also, policy learning can be done efﬁciently via backpropagation through the differentiable dynamics model.
Revised: Also, policy learning can be done efﬁciently via backpropagation through the differentiable dynamics model.

Original: In image-based RL, the key problem is to learn low-dimensional state representation and, in the model-based case, also its forward model.
Revised: In image-based RL, the key problem is to learn low-dimensional state representation and, in the model-based case, also its forward model.

Original: 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.
Revised: 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.

Original: Instead, it is more practical to introduce auxiliary tasks providing richer learning signal to facilitate representation learning without reward (or with sparse reward)
Revised: Instead, it is more practical to introduce auxiliary tasks providing richer learning signal to facilitate representation learning without reward (or with sparse reward)

Original: (Sutton et al., 2011; Jaderberg et al., 2016).
Revised: (Sutton et al., 2011; Jaderberg et al., 2016).

Original: D REAMER achieves this by learning the representation and the dynamics model in a way to reduce the reconstruction error of the observed sequences.
Revised: D REAMER achieves this by learning the representation and the dynamics model in a way to reduce the reconstruction error of the observed sequences.

Original: However, reconstruction-based representation learning has limitations.
Revised: However, reconstruction-based representation learning has limitations.

Original: First, it is computationally expensive to reconstruct the high-dimensional inputs, especially in models like D REAMER that needs to reconstruct long-range videos.
Revised: First, it is computationally expensive to reconstruct the high-dimensional inputs, especially in models like D REAMER that needs to reconstruct long-range videos.

Original: 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).
Revised: 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).

Original: Thus, in MBRL it is of particular interest to realize a version of D REAMER without reconstruction.
Revised: Thus, in MBRL it is of particular interest to realize a version of D REAMER without reconstruction.

Original: Recently, there have been remarkable advances in reconstruction-free representation learning in reinforcement learning (Laskin et al., 2020a;b; Yarats et al., 2021c).
Revised: Recently, there have been remarkable advances in reconstruction-free representation learning in reinforcement learning (Laskin et al., 2020a;b; Yarats et al., 2021c).

Original: The currently dominant approach is via contrastive learning.
Revised: The currently dominant approach is via contrastive learning.

Original: This approach requires pair-wise comparisons to push apart different instances while pulling close an instance and its augmentation.
Revised: This approach requires pair-wise comparisons to push apart different instances while pulling close an instance and its augmentation.

Original: Therefore, this method usually requires a large batch size (so computationally expensive) to perform accurately and robustly.
Revised: Therefore, this method usually requires a large batch size (so computationally expensive) to perform accurately and robustly.

Original: An alternative is the clustering-based or prototype-based approach (Caron et al., 2020).
Revised: An alternative is the clustering-based or prototype-based approach (Caron et al., 2020).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: The prototype structure can also be used to implement an exploration method (Yarats et al., 2021b).
Revised: The prototype structure can also be used to implement an exploration method (Yarats et al., 2021b).

Original: However, for reconstruction-free MBRL only the contrastive approach like Temporal Predictive
Revised: However, for reconstruction-free MBRL only the contrastive approach like Temporal Predictive

Original: Coding (TPC) (Nguyen et al., 2021) has been proposed so far.
Revised: Coding (TPC, Nguyen et al., 2021) has been proposed so far.

Original: Thus, it is important to investigate (1) how, in the MBRL setting, we can realize the improved robustness and accuracy of the prototypical representation as shown in computer vision, (2) whether and how temporal information available in MBRL can be combined with the prototypes to improve the representation further, and ﬁnally (3) what the challenges and beneﬁts are in doing so.
Revised: While TPC consistently outperforms

Original: For instance, although TPC has shown to consistently outperform D REAMER in the noisy background settings, for standard DeepMind Control Suite (DMC) (Tassa et al., 2018)
Revised: D REAMER in the noisy background settings, for standard DeepMind Control suite (Tassa et al., 2018)

Original: it showed quite inconsistent results by performing severely worse than D REAMER for some tasks.
Revised: it showed quite inconsistent results by performing severely worse than D REAMER on some tasks.

Original: In this work, wetherefore hypothesize that this inconsistent behavior may be ﬁxedin the MBRL setting if the robustness and accuracy of the prototypical representation can be realized  and further improved with the support of temporal information.
Revised: Therefore, we hypothesize that this inconsistent behavior may be ﬁxed if the robustness and accuracy of the prototypical representations can be realized in MBRL and further improved with the support of temporal information.

Original: In this paper, we propose a reconstruction-free MBRL agent, called D REAMER P RO , by combining the prototypical representation learning with temporal dynamics learning.
Revised: In this paper, we propose a reconstruction-free MBRL agent, called D REAMER P RO , by combining the prototypical representation learning with temporal dynamics learning.

Original: Similar to SwAV, by encouraging uniform cluster assignment across the batch, we implicitly pull apart the embeddings of different observations.
Revised: Similar to SwAV (Caron et al., 2020), by encouraging uniform cluster assignment across the batch, we implicitly pull apart the embeddings of different observations.

Original: 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.
Revised: Additionally, we let the temporal latent state to ‘reconstruct’ the cluster assignment of the observation, thereby relieving the world model from modeling lowlevel details.

Original: We evaluate our model on the standard setting of DeepMind Control Suite, and also on a naturalbackground setting, where the background is replaced by natural videos irrelevant to the task.
Revised: 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.

Original: The results show that the proposed modelis consistently more accurate than the previous models.
Revised: The results show that the proposed model consistently outperforms previous methods.

Original: The contributions of the paper are (1) the ﬁrst 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.
Revised: The contributions of the paper are (1) the ﬁrst 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.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: In this section, we brieﬂy introduce the world model and learning algorithms used in
Revised: In this section, we brieﬂy introduce the world model and learning algorithms used in

Original: D REAMER V2 (Hafner et al., 2021) which our model builds upon.
Revised: D REAMER V2 (Hafner et al., 2021) which our model builds upon.

Original: To indicate the general D REAMER framework (Hafner et al., 2020; 2021), we omit its version number in the rest of the paper.
Revised: To indicate the general D REAMER framework (Hafner et al., 2020; 2021), we omit its version number in the rest of the paper.

Original: 2. R ECONSTRUCTION - BASED WORLD MODEL LEARNING
Revised: 2. R ECONSTRUCTION - BASED WORLD MODEL LEARNING

Original: D REAMER learns a recurrent state-space model (RSSM, Hafner et al. to predict forward dynamics and rewards in partially observable environments.
Revised: D REAMER learns a recurrent state-space model (RSSM, Hafner et al., 2019) to predict forward dynamics and rewards in partially observable environments.

Original: At each time step t , the agent receives an image observation o t and a scalar reward r t (obtained by previous actions a <t ).
Revised: At each time step t , the agent receives an image observation o t and a scalar reward r t (obtained by previous actions a <t ).

Original: The agent then chooses an action a t based on its policy.
Revised: The agent then chooses an action a t based on its policy.

Original: The RSSM models the observations, rewards, and transitions through a probabilistic generative process:
Revised: The RSSM models the observations, rewards, and transitions through a probabilistic generative process:

Original: [Equation]
Revised: [Equation]

Original: where the latent variables s 1: T are the agent states, and h t =
Revised: where the latent variables s 1: T are the agent states, and h t =

Original: GRU ( h t − 1 , s t − 1 , a t − 1 ) is a deterministic encoding of s <t and a <t .
Revised: GRU ( h t − 1 , s t − 1 , a t − 1 ) is a deterministic encoding of s <t and a <t .

Original: To infer the agent states from past observations and actions, a variational encoder is introduced:
Revised: To infer the agent states from past observations and actions, a variational encoder is introduced:

Original: [Equation]
Revised: [Equation]

Original: The training objective is to maximize the evidence lower bound (ELBO):
Revised: The training objective is to maximize the evidence lower bound (ELBO):

Original: [Equation]
Revised: [Equation]

Original: 2. P OLICY LEARNING BY LATENT IMAGINATION
Revised: 2. P OLICY LEARNING BY LATENT IMAGINATION

Original: D REAMER interleaves policy learning with world model learning.
Revised: D REAMER interleaves policy learning with world model learning.

Original: During policy learning, the world model is ﬁxed, and an actor and a critic are trained cooperatively from the latent trajectories imagined by the world model.
Revised: During policy learning, the world model is ﬁxed, and an actor and a critic are trained cooperatively from the latent trajectories imagined by the world model.

Original: Speciﬁcally, the imagination starts at each non-terminal state ˆ z t =
Revised: Speciﬁcally, the imagination starts at each non-terminal state ˆ z t =

Original: [ h t , s t ] encountered during world model learning.
Revised: [ h t , s t ] encountered during world model learning.

Original: Then, at each imagination step t (cid:48) ≥ t , an action is sampled from the actor’s stochastic policy: ˆ a t (cid:48) ∼ π (ˆ a t (cid:48) | ˆ z t (cid:48) ) .
Revised: Then, at each imagination step t (cid:48) ≥ t , an action is sampled from the actor’s stochastic policy: ˆ a t (cid:48) ∼ π (ˆ a t (cid:48) | ˆ z t (cid:48) ) .

Original: The corresponding reward ˆ r t (cid:48) + and next state ˆ z t (cid:48) +1 are predicted by the learned world model.
Revised: The corresponding reward ˆ r t (cid:48) +1 and next state ˆ z t (cid:48) +1 are predicted by the learned world model.

Original: 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.
Revised: 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.

Original:  D REAMER P RO
Revised:  D REAMER P RO

Original: To compute the D REAMER training objective, more speciﬁcally J t O in Equation 4, a decoder is required to reconstruct the image observation
Revised: To compute the D REAMER training objective, more speciﬁcally J t O in Equation 4, a decoder is required to reconstruct the image observation

Original: o t from the state z t =
Revised: o t from the state z t =

Original: [ h t , s t ] .
Revised: [ h t , s t ] .

Original: Because this reconstruction loss operates in pixel space where all pixels are weighted equally, D REAMER tends to allocate most of its capacity to modeling complex visual patterns that cover a large pixel area (e.g., backgrounds).
Revised: Because this reconstruction loss operates in pixel space where all pixels are weighted equally, D REAMER tends to allocate most of its capacity to modeling complex visual patterns that cover a large pixel area (e.g., backgrounds).

Original: This leads to poor task performance when those visual patterns are task irrelevant, as shown in previous work (Nguyen et al., 2021).
Revised: This leads to poor task performance when those visual patterns are task irrelevant, as shown in previous work (Nguyen et al., 2021).

Original: Fortunately, during policy learning, what we need is accurate reward and next state prediction, which are respectively encouraged by J t R and J t KL .
Revised: Fortunately, during policy learning, what we need is accurate reward and next state prediction, which are respectively encouraged by J t R and J t KL .

Original: In other words, the decoder is not required for policy learning.
Revised: In other words, the decoder is not required for policy learning.

Original: The main purpose of having the decoder and the associated loss J t O , as shown in
Revised: The main purpose of having the decoder and the associated loss J t O , as shown in

Original: D REAMER , is to learn meaningful representations that cannot be obtained by J t R and J t KL alone.
Revised: D REAMER , is to learn meaningful representations that cannot be obtained by J t R and J t KL alone.

Original: The above observations motivate us to improve robustness to visual distractions by replacing the reconstruction-based representation learning in D REAMER with reconstruction-free methods.
Revised: The above observations motivate us to improve robustness to visual distractions by replacing the reconstruction-based representation learning in D REAMER with reconstruction-free methods.

Original: 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;
Revised: 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;

Original: Chen et al., 2020; He et al., 2020) and non-contrastive (Grill et al., 2020; Caron et al., 2020) methods.
Revised: Chen et al., 2020; He et al., 2020) and non-contrastive (Grill et al., 2020; Caron et al., 2020) methods.

Original: We prefer non-contrastive methods as they can be applied to small batch sizes.
Revised: We prefer non-contrastive methods as they can be applied to small batch sizes.

Original: This can speed up both world model learning and policy learning (in wall clock time).
Revised: This can speed up both world model learning and policy learning (in wall clock time).

Original: Therefore, we propose to combine D REAMER with the prototypical representations used in S W AV (Caron et al., 2020), a top-performing non-contrastive representation learning method.
Revised: Therefore, we propose to combine D REAMER with the prototypical representations used in S W AV (Caron et al., 2020), a top-performing non-contrastive representation learning method.

Original: We name the resulting model D REAMER P RO , and provide the model description in the following.D REAMER
Revised: We name the resulting model D REAMER P RO , and provide the model description in the following.

Original:  P RO uses the same policy learning algorithm as D REAMER , but learns the world model without reconstructing the observations.
Revised: D REAMER P RO uses the same policy learning algorithm as D REAMER , but learns the world model without reconstructing the observations.

Original: This is achieved by clustering the observation into a set of
Revised: This is achieved by clustering the observation into a set of

Original: K trainable prototypes { c 1 , . .
Revised: K trainable prototypes { c 1 , . .

Original: . , c K } , and then predicting the cluster assignment from the state as well as an augmented view of the observation.
Revised: . , c K } , and then predicting the cluster assignment from the state as well as an augmented view of the observation.

Original: See Figure 1 for an illustration.
Revised: See Figure 1 for an illustration.

Original: Concretely, given a sequence of observations o 1: 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).
Revised: Concretely, given a sequence of observations o 1: 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).

Original: We ensure that the augmentation is consistent across time steps.
Revised: We ensure that the augmentation is consistent across time steps.

Original: Each view i ∈ { 1 , 2 } is fed to the RSSM to obtain the states z ( i ) 1: T .
Revised: Each view i ∈ { 1 , 2 } is fed to the RSSM to obtain the states z ( i ) 1: T .

Original: To predict the cluster assignment from z ( i ) t , we ﬁrst apply a linear projection followed by (cid:96) 2 -normalization to obtain a vector x ( i ) t of the same dimension as the prototypes, and then take a softmax over the dot [Equation]
Revised: To predict the cluster assignment from z ( i ) t , we ﬁrst apply a linear projection followed by (cid:96) 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:

Original: ! 𝑐 " ⊙ ⊙ ⊙ 𝑥 ! "

Original: [Equation]
Revised: [Equation]

Original: 𝑦 !

Original: 𝑢 !

Original: 𝑣 ! 𝑤 ! ℐ "#$% ℐ &’() 𝑧 ! // AUG 𝑜 !
Revised: [Table]

Original: AUG 𝑜 ! 𝑓 * 𝑓 +* EMA Prototypes
Revised:

Original: State Projections Probabilities // - Stop Gradient RSSM
Revised: [Table]

Original: Softmax Softmax Sinkhornoverbatch products of x ( i ) t and all the prototypes: [Equation]
Revised: /hxoUibv0ad/6NmXYW2nogcDjnXnLPCRPOtHHdb6e0sbm1vVPereztHxweVY9PulqmitAOkVyqfog15UzQjmGG036iKI5DTnvh9Db3e09UaSbFg5klNIjxWLCIEWys5LeVfKQkp3pYrbl1dwG0TryC1KBAe1j9GowkSWMqDOFYa99zExNkWBlGOJ1XBqmmCSZTPKa+pQLHVAfZ4uQ5urDKCEVS2ScMWqi/NzIcaz2LQzsZYzPRq14u/uf5qYmug4yJJDVUkOVHUcqRkSjPj0ZM2cB8ZgkmitlbEZlghYmxLVVsCd5q5HXSbdS9q7p336i1boo6ynAG53AJHjShBXfQhg4QkPAMr/DmGOfFeXc+lqMlp9g5hT9wPn8AopuReA==</latexit> Projections

Original: 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 { c 1 , . .
Revised: 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 { c 1 , . .

Original: , c K } are also (cid:96) 2 -normalized.
Revised: , c K } are also (cid:96) 2 -normalized.

Original: 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 (cid:96) 2normalization, and obtain a vector y ( i ) t .
Revised: 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 (cid:96) 2normalization, and obtain a vector y ( i ) t .

Original: We summarize this process as: y ( i ) t
Revised: We summarize this process as: y ( i ) t

Original: = f
Revised: = f

Original: θ ( o ( i ) t ) , where θ collectively denotes the parameters of the convolutional encoder and the linear projection layer.
Revised: θ ( o ( i ) t ) , where θ collectively denotes the parameters of the convolutional encoder and the linear projection layer.

Original: The prediction probabilities are again given by a softmax:
Revised: The prediction probabilities are again given by a softmax:

Original: [Equation]
Revised: [Equation]

Original:  where v ( i ) t,k is the predicted probability that observation o ( i ) t maps to cluster k .
Revised: 6) where v ( i ) t,k is the predicted probability that observation o ( i ) t maps to cluster k .

Original: 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
Revised: 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

Original: f ¯ θ
Revised: f ¯ θ

Original: (He et al., 2020; Grill et al., 2020; Caron et al., 2021), whose parameters ¯ θ are updated using the exponential moving average of θ : ¯ θ ← (1 − η ) ¯ θ + ηθ .
Revised: (He et al., 2020; Grill et al., 2020; Caron et al., 2021), whose parameters ¯ θ are updated using the exponential moving average of θ : ¯ θ ← (1 − η )¯ θ + ηθ .

Original: For each observation o ( i ) t , the scores are given by the dot products (¯ y ( i ) t · c 1 , . . .
Revised: For each observation o ( i ) t , the scores are given by the dot products (¯ y ( i ) t · c 1 , . . .

Original: , ¯ y ( i ) t · c K ) , where ¯ y ( i ) t
Revised: , ¯ y ( i ) t · c K ) , where ¯ y ( i ) t

Original: = f ¯ θ ( o ( i ) t ) is the momentum encoder output.
Revised: = f ¯ θ ( o ( i ) t ) is the momentum encoder output.

Original: 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.
Revised: 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.

Original: We speciﬁcally 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.
Revised: We speciﬁcally 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.

Original: The outcome of the Sinkhorn-Knopp algorithm is a set of cluster assignment targets ( w ( i ) t, 1 , . . .
Revised: The outcome of the Sinkhorn-Knopp algorithm is a set of cluster assignment targets ( w ( i ) t, 1 , . . .

Original: , w ( i ) t,K ) for each observation o ( i ) t .
Revised: , w ( i ) t,K ) for each observation o ( i ) t .

Original: Now that we have the cluster assignment predictions and targets, the representation learning objective is simply to maximize the prediction accuracies:
Revised: Now that we have the cluster assignment predictions and targets, the representation learning objective is simply to maximize the prediction accuracies:

Original: [Equation]
Revised: [Equation]

Original: Here, J t S W AV improves prediction from an augmented view.
Revised: Here, J t S W AV improves prediction from an augmented view.

Original: This is the same loss as used in S W AV (Caron et al., 2020), and is shown to induce useful features for static images.
Revised: This is the same loss as used in S W AV (Caron et al., 2020), and is shown to induce useful features for static images.

Original: However, it ignores the temporal structure which is crucial in reinforcement learning.
Revised: However, it ignores the temporal structure which is crucial in reinforcement learning.

Original: Hence, we add a second term,
Revised: Hence, we add a second term,

Original: J t Temp , that improves prediction from the state of the same view.
Revised: J t Temp , that improves prediction from the state of the same view.

Original: 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.
Revised: 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.

Original: From another perspective, J t Temp is similar to J t O in the sense that we are now ‘reconstructing’ the cluster assignment of the observation instead of the observation itself.
Revised: From another perspective, J t Temp is similar to J t O in the sense that we are now ‘reconstructing’ the cluster assignment of the observation instead of the observation itself.

Original: This frees the world model from modeling complex visual details, allowing more capacity to be devoted to task-relevant features.
Revised: This frees the world model from modeling complex visual details, allowing more capacity to be devoted to task-relevant features.

Original: The overall world model learning objective for D REAMER P RO can be obtained by replacing J t O in
Revised: The overall world model learning objective for D REAMER P RO can be obtained by replacing J t O in

Original: Equation 4 with J t S WAV + J t
Revised: Equation 4 with J t S W

Original: Temp :

Original: [Equation]
Revised: [Equation]

Original: where J t R and J t KL are now averaged over the two augmented views.
Revised: where J t R and J t KL are now averaged over the two augmented views.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Environments.
Revised: Environments.

Original: We evaluate our model and the baselines on six image-based continuous control tasks from the DeepMind Control (DMC) suite (Tassa et al., 2018).
Revised: We evaluate our model and the baselines on six image-based continuous control tasks from the DeepMind Control (DMC) suite (Tassa et al., 2018).

Original: We choose the set of tasks based on those considered in P LA N ET (Hafner et al., 2019).
Revised: We choose the set of tasks based on those considered in P LA N ET (Hafner et al., 2019).

Original: Speciﬁcally, we replace Cartpole Swingup and Walker Walk with their more challenging counterparts, Cartpole Swingup Sparse and 
Revised: Speciﬁcally, we replace Cartpole Swingup and Walker Walk with their more challenging counterparts, Cartpole Swingup Sparse and Walker

Original: Walker Run, and keep the remaining tasks.
Revised: Run, and keep the remaining tasks.

Original: 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).
Revised: 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 Kineticsdataset (Kay et al., 2017).

Original: Following TPC (Nguyen et al., 2021), we use two separate sets of background videos for training and evaluation.
Revised: Following TPC (Nguyen et al., 2021), we use two separate sets of background videos for training and evaluation.

Original: Hence, the natural background setting tests generalization to unseen distractions.
Revised: Hence, the natural background setting tests generalization to unseen distractions.

Original: We note that the recently released Distracting Control Suite (DCS,Stone et al.
Revised: We note that the recently released Distracting Control Suite (DCS,

Original: (2021)) serves a similar purpose.
Revised: Stone et al., 2021) serves a similar purpose.

Original: However, the background distractions in DCS seem less challenging, as there are fewer videos and the ground plane is made visible for most tasks.
Revised: However, the background distractions in DCS seem less challenging, as there are fewer videos and the ground plane is made visible for most tasks.

Original: In our preliminary experiments, our model and all the baselines achieved close to zero returns on 
Revised: In our preliminary experiments, our model and all the baselines achieved close to zero returns on Cartpole

Original: Cartpole Swingup Sparse in the natural background setting.
Revised: Swingup Sparse in the natural background setting.

Original: We therefore switch back to Cartpole Swingup in this setting.
Revised: We therefore switch back to Cartpole Swingup in this setting.

Original: Baselines.
Revised: Baselines.

Original: Our main baselines are D REAMER (Hafner et al., 2021) and TPC (Nguyen et al., 2021), the state-of-the-art for reconstruction-based and reconstruction-free model-based reinforcement learning, respectively.
Revised: Our main baselines are D REAMER (Hafner et al., 2021), D REAMING (Okada & Taniguchi, 2021), and TPC (Nguyen et al., 2021), the state-of-the-art for reconstruction-based and reconstruction-free MBRL.

Original: In particular, TPC has shown better performance than CVRL (Ma et al., 2020) and DBC (Zhang et al., 2021) on the same datasets.
Revised: 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.

Original: The recently proposed PSE (Agarwal et al., 2021) has demonstrated impressive results on DCS.
Revised: The recently proposed PSE (Agarwal et al., 2021) has demonstrated impressive results on DCS.

Original: 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.
Revised: 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.

Original: Implementation details.
Revised: Implementation details.

Original: We implement our model  based on a newer version of D REAMER 1 , while the ofﬁcial implementation of TPC 2 is based on an older version.
Revised: We implement our model and D REAMING based on a newer version of D REAMER 1 , while the ofﬁcial implementation of [Figure]

Original: For fair comparison, we reimplement TPC based on the newer version.
Revised: [Table] comparison, we re-implement TPC based on the newer version.

Original: We adopt the default values for the D REAMER hyperparameters, except that we use continuous latents and tanh normal as the distribution output by the actor.
Revised: We adopt the default values for the D REAMER hyperparameters, except that we use continuous latents and tanh normal as the distribution output by the actor.

Original: We ﬁnd these changes improve D REAMER ’s performance in the standard DMC, and therefore use these values for all models in both the standard and the natural background setting.
Revised: We ﬁnd these changes improve D REAMER ’s performance in the standard DMC, and therefore use these values for all models in both the standard and the natural background setting.

Original: Following TPC, we increase the weight of the reward loss J t R to 1000 for all models in the natural background setting to further encourage extraction of task-relevant information.
Revised: Following TPC, we increase the weight of the reward loss J t R to 1000 for all models in the natural background setting to further encourage extraction of task-relevant information.

Original: While in the original TPC, this weight is chosen separately for each task from { 100 , 1000 } , we ﬁnd the weight of 1000 works consistently better in our re-implementation, which also obtains better results than reported in the original paper
Revised: While in the original TPC, this weight is chosen separately for each task from { 100 , 1000 } , we ﬁnd the weight of 1000 works consistently better in our re-implementation, which also obtains better results than reported in the original paper.

Original: [Figure]
Revised: We use the default batch size of 50 for D REAMER , D REAMING , and D REAMER P RO .

Original: [Table]
Revised: The batch size for TPC is chosen to be 150, so that it has similar wall clock training time as D REAMER P RO .

Original: Evaluation protocol.
Revised: Evaluation protocol.

Original: For each task, we train each model for 1 M environment steps (equivalent to 500 K actor steps, as the action repeat is set to 2 ).
Revised: For each task, we train each model for 1 M environment steps (equivalent to 500 K actor steps, as the action repeat is set to 2 ).

Original: The evaluation return is computed every 10 K steps, and averaged over 10 episodes.
Revised: The evaluation return is computed every 10 K steps, and averaged over 10 episodes.

Original: In all ﬁgures and tables, the mean and standard deviation are computed from 3 independent runs.
Revised: In all ﬁgures and tables, the mean and standard deviation are computed from 3 independent runs.

Original: In addition to the episode returns, we also report the Normalized

Original: Distance to the Best (NDB).

Original: Given the performance x τ,α of an algorithm α on a task τ , the NDB β τ,α is deﬁned as:

Original: [Equation]

Original: We report the mean and variance of β τ,α for each algorithm α over all tasks τ , which indicate the bestness and consistency of algorithm α , respectively.

Original: An optimal algorithm that consistently achieves the best performance across all tasks will have zero mean and zero variance.

Original: 4. P ERFORMANCE IN STANDARD DMC
Revised: 4. P ERFORMANCE IN STANDARD DMC

Original: We show the performance curves in Figure 2 and the ﬁnal performance in Table 1 for the standard setting.
Revised: We show the performance curves in Figure 2 and the ﬁnal performance in Table 1 for the standard setting.

Original: In contrast to TPC which underperforms D REAMER on most tasks (most severely on Reacher Easy and Walker Run), D REAMER P RO  achieves comparable or even better performance [Equation]
Revised: D REAMER P RO is the only model that achieves comparable or even better performance than D REAMER on all tasks. Notably, D REAMER P RO outperforms D REAMER by a large margin on

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: Mean and STD of NDB ( ↓ )

Original: [Table] [Table] than D REAMER on all tasks. Notably, D REAMER P RO outperforms D REAMER by a large margin on Finger Spin and Reacher Easy, and demonstrates better data efﬁciency on Cup Catch.
Revised: Finger Spin and Reacher Easy, and demonstrates better data efﬁciency on Cup Catch.

Original: We notice a large variance in D REAMER P RO ’s performance on Finger Spin.
Revised: We notice a large variance in D REAMER P RO ’s performance on Finger Spin.

Original: Further investigation reveals that D REAMER P RO 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. 
Revised: Further investigation reveals that D REAMER P RO 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 D REAMER indicates that it hardly achieved close to optimal behavior.

Original: Our results suggest for the ﬁrst time that prototypical representations (and reconstruction-free representation learning in general) can be beneﬁcial to MBRL even in the absence of strong visual distractions.
Revised: Our results suggest for the ﬁrst time that prototypical representations (and reconstructionfree representation learning in general) can be beneﬁcial to MBRL even in the absence of strong visual distractions.

Original: 4. P ERFORMANCE IN NATURAL BACKGROUND DMC
Revised: 4. P ERFORMANCE IN NATURAL BACKGROUND DMC

Original: Figure 3 and Table 2 respectively show the performance curves and ﬁnal evaluation returns obtained by all models in the natural background setting.
Revised: Figure 3 and Table 2 respectively show the performance curves and ﬁnal evaluation returns obtained by all models in the natural background setting.

Original: D REAMER completely fails on all tasks, showing the inability of reconstruction-based representation learning to deal with complex visual distractions.
Revised: D REAMER completely fails on all tasks, showing the inability of reconstruction-based representation learning to deal with complex visual distractions.

Original: In contrast, D REAMER P RO achieves the best performance on 5 out of 6 tasks, with large performance gains from TPC on Cartpole Swingup, Finger Spin,Reacher Easy, and Walker Run.
Revised: In contrast, D REAMER P RO achieves the best performance on 4 out of 6 tasks, with large performance gains from baselines on Cartpole Swingup, Finger Spin, and Walker Run.

Revised: We additionally train TPC with a batch size of 300 (denoted TPC-Batch-300), and D REAMER P RO is still able to outperform it on 4 out of 6 tasks.

Original: 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. 
Revised: 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. [Table]

Original: [Figure]
Revised: [Figure]

Revised: , and the remaining rows show the reconstruction from the

Revised: RSSM state z t for each model.

Revised: 4. V ISUALIZATION AND ANALYSIS

Revised: 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).

Revised: Figure 4 (Left) shows the reconstructions for Cup Catch after 100K environment steps.

Revised: Note that the reconstruction of the ball is only possible from the states learned by D REAMER P RO , explaining its better data efﬁciency.

Revised: D REAMER 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.

Revised: Figure 4 (Right) shows the reconstructions for Walker Run after 1M environment steps.

Revised: We see that both D REAMER and D REAMER P RO capture some information of the background, but only

Revised: D REAMER P RO is able to recover the posture of the Walker.

Revised: The reconstruction from TPC does not seem relevant to the Walker or the background.

Revised: This indicates that while TPC may be better at discarding distractors, D REAMER P RO is better at retaining task-relevant information.

Original: 4. A BLATION STUDY
Revised: 4. A BLATION STUDY

Original: We now show the individual effect of the two loss terms, J t S W AV and J t Temp , in Figure 4.
Revised: We now show the individual effect of the two loss terms, J t S W AV and J t Temp , in Figure 6.

Original: Here, each of the ablated versions, DreamerPro-No-SwAV and DreamerPro-No-Temp, removes one of the loss terms.
Revised: Here, each of the ablated versions, DreamerPro-No-SwAV and DreamerPro-No-Temp, removes one of the loss terms.

Original: We did not investigate removing both terms, as its failure has been shown in D REAMER (Hafner et al., 2020).
Revised: We did not investigate removing both terms, as its failure has been shown in D REAMER (Hafner et al., 2020).

Original: We  observe that both terms are necessary for achieving good performance.
Revised: We train the ablated versions in natural background DMC, and observe that both terms are necessary for achieving good performance.

Original: In particular, naively combining S W AV with D REAMER (i.e., DreamerPro-No-Temp) leads to inferior performance, as it ignores the temporal structure.
Revised: In particular, naively combining S W AV with D REAMER (i.e., DreamerPro-No-Temp) leads to inferior performance, as it ignores the temporal structure.

Original: On the other hand, J t Temp alone is not sufﬁcient to provide meaningful cluster assignment targets and learning signals for the convolutional encoder.
Revised: On the other hand, J t Temp alone is not sufﬁcient to provide meaningful cluster assignment targets and learning signals for the convolutional encoder.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Self-supervised representation learning for static images.
Revised: Self-supervised representation learning for static images.

Original: Recent works in self-supervised learning have shown its effectiveness in learning representations from high-dimensional data.
Revised: Recent works in self-supervised learning have shown its effectiveness in learning representations from high-dimensional data.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: To address this problem, some works propose to learn the image representations without discriminating between samples.
Revised: To address this problem, some works propose to learn the image representations without discriminating between samples.

Original: Particularly, BYOL (Grill et al., 2020) introduces a momentum encoder to provide target representations for the training network.
Revised: Particularly, BYOL (Grill et al., 2020) introduces a momentum encoder to provide target representations for the training network.

Original: SwAV (Caron et al., 2020) proposes to learn the embeddings by matching them to a set of learned clusters.
Revised: SwAV (Caron et al., 2020) proposes to learn the embeddings by matching them to a set of learned clusters.

Original: And DINO (Caron et al., 2021) replace the clusters in SwAV with categorical heads and uses the centering and sharpening technique to prevent representations collapsing.
Revised: DINO (Caron et al., 2021) replaces the clusters in SwAV with categorical heads and uses the centering and sharpening technique to prevent representations collapsing.

Original: Unlike our model, these works treat each image independently and ignore the temporal structure of the environment, with is crucial in learning the  policy in RL tasks.
Revised: 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.

Original: Representation learning for model-free reinforcement learning.
Revised: Representation learning for model-free reinforcement learning.

Original: 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
Revised: 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

Original: (Laskin et al., 2020b; Yarats et al., 2021c;a).
Revised: (Laskin et al., 2020b; Yarats et al., 2021c;a).

Original: Recent works have also shown that self-supervised representation learning techniques can bring signiﬁcant improvement to reinforcement learning methods.
Revised: Recent works have also shown that self-supervised representation learning techniques can bring signiﬁcant improvement to reinforcement learning methods.

Original: For example, CURL (Laskin et al., 2020a) performs contrastive learning along with off-policy RL algorithms and shows that it signiﬁcantly improves sample-efﬁciency and model performance over pixel-based methods.
Revised: For example, CURL (Laskin et al., 2020a) performs contrastive learning along with off-policy RL algorithms and shows that it signiﬁcantly improves sample-efﬁciency and model performance over pixel-based methods.

Original: 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;
Revised: 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;

Original: Yarats et al., 2021b;
Revised: Yarats et al., 2021b;

Original: Guo et al., 2020; Gregor et al., 2019).
Revised: Guo et al., 2020; Gregor et al., 2019).

Original: 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.
Revised: 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.

Original: Model-based reinforcement learning with reconstruction.
Revised: Model-based reinforcement learning with reconstruction.

Original: Model based reinforcement learning from raw pixel data can learn the representation space by minimizing the observation reconstruction loss.
Revised: Model based reinforcement learning from raw pixel data can learn the representation space by minimizing the observation reconstruction loss.

Original: World Models (Ha & Schmidhuber) learn the latent dynamics of the environment in a two-stage process to evolve their linear controllers in imagination.
Revised: World Models (Ha & Schmidhuber, 2018) learn the latent dynamics of the environment in a two-stage process to evolve their linear controllers in imagination.

Original: SOLAR (M. Zhang & Levine., 2019) models the dynamics as time-varying linear-Gaussian and solves robotic tasks via guided policy search.
Revised: SOLAR (Zhang et al., 2019) models the dynamics as time-varying linear-Gaussian and solves robotic tasks via guided policy search.

Original:  (Hafner et al., 2020) jointly learns the RSSM and latent state space from observation reconstruction loss.
Revised: Dreamer (Hafner et al., 2020) jointly learns the RSSM and latent state space from observation reconstruction loss.

Original:  (Gelada et al., 2019) also propose a latent dynamics model-based method that uses bisimulation metrics and reconstruction loss in Atari.
Revised: DeepMDP (Gelada et al., 2019) also propose a latent dynamics modelbased method that uses bisimulation metrics and reconstruction loss in Atari.

Original: However, reconstruction based methods are susceptible to noise and objects irrelevant to the task in the environment (Nguyen et al., 2021;
Revised: However, reconstruction based methods are susceptible to noise and objects irrelevant to the task in the environment (Nguyen et al., 2021).

Original: Okada & Taniguchi, 2021).

Original: Furthermore, in a few cases, the latent representation fails to reconstruct small task-relevant objects in the environment(Okada & Taniguchi, 2021).
Revised: Furthermore, in a few cases, the latent representation fails to reconstruct small task-relevant objects in the environment (Okada & Taniguchi, 2021).

Original: Reinforcement learning under visual distractions.
Revised: Reinforcement learning under visual distractions.

Original: A large body of works on robust representation learning focuses on contrastive objectives.
Revised: A large body of works on robust representation learning focuses on contrastive objectives.

Original: 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.
Revised: 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.

Original: However, the learning objective of CVRL encourages the representation model to learn as much information as possible, including task-irrelevant information.
Revised: However, the learning objective of CVRL encourages the representation model to learn as much information as possible, including task-irrelevant information.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Using the bisimulation metrics that quantify the behavioral similarity between states, these methods make the mode robust to task-irrelevant information.
Revised: Using the bisimulation metrics that quantify the behavioral similarity between states, these methods make the mode robust to task-irrelevant information.

Original: However, DBC cannot generalize to unseen backgrounds (Nguyen et al., 2021), and PSE  requires a pre-trained policy to compute the similarity metrics while our model learns  the policy from scratch.
Revised: 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.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we presented the ﬁrst reconstruction-free MBRL agent based on the prototypical representation and its temporal dynamics.
Revised: In this work, we presented the ﬁrst reconstruction-free MBRL agent based on the prototypical representation and its temporal dynamics.

Original: In experiments, we demonstrated the consistently improved accuracy and robustness of the proposed model in comparison to the Temporal Predictive Control (TPC) agent and the Dreamer agent for both standard and natural background DMC tasks.
Revised: 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.

Original: Our results suggest that there are unexplored broad areas in reconstruction-free MBRL.
Revised: Our results suggest that there are unexplored broad areas in reconstruction-free MBRL.

Original: 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.
Revised: 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.

Original: A H YPERPARAMETERS
Revised: A H YPERPARAMETERS

Original: For hyperparameters that are shared with D REAMER , we use the default values suggested in the conﬁg ﬁle in the ofﬁcial implementation of D REAMER , with the following two exceptions.
Revised: For hyperparameters that are shared with D REAMER , we use the default values suggested in the conﬁg ﬁle in the ofﬁcial implementation of D REAMER , with the following two exceptions.

Original: We set discrete =
Revised: We set rssm.discrete =

Original: False and actor.dist = tanh normal , as we ﬁnd these changes improve performance over the default setting.
Revised: False and actor.dist = tanh normal , as we ﬁnd these changes improve performance over the default setting.

Original: The additional hyperparameters introduced in D REAMER P RO are listed in Table 3.
Revised: The additional hyperparameters introduced in D REAMER P RO are listed in Table 3.

Revised: We ﬁnd it helpful to freeze the prototypes for the ﬁrst 10 K gradient updates.

Revised: In the natural background setting, we add a squared loss that encourages the (cid:96) 2 -norm of projections (before (cid:96) 2 -normalization) to be close to 1 .

Revised: This helps stabilize the model.

Original: [Table]
Revised: [Table]

Revised: B N EAREST NEIGHBOR QUERIES IN LATENT SPACE

Revised: [Table]

Revised: APr8wcxT5G1</latexit

Revised: TPC

Revised: We sample a batch of trajectories from the training replay buffer, and obtain the latent state for each image.

Revised: Then, given a query image, we show the three images in the batch whose latent states are the closest to the query image.

Revised: We use the same batch and same query images for all models.

Revised: 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).

Revised: On the other hand, the nearest neighbors for

Revised: D REAMER P RO tend to have similar backgrounds and also similar agent states.

Revised: Our results suggest that D REAMER P RO and TPC work in quite different ways.

Revised: D REAMER P RO tries to retain task-relevant information at the cost of also including some distractors, while TPC focuses more on discarding distractors.

Revised: An interesting future direction is to simultaneously consider these two factors and achieve a better balance.

Revised: C A BLATION RESULTS

Revised: [Figure]

Revised: D A TARI RESULTS

Revised: [Figure]

Revised: To show the potential of D REAMER P RO to leverage the beneﬁts of world models and discrete latents in complex environments

Revised: [Figure]

Revised: we train D REAMER P RO on a subset of six Atari games for 25M, 50M, and 80M environment steps depending on the convergence speed of D REAMER on these games.

Revised: We freeze the prototypes for 30K gradient updates.

Revised: The weights for reward and KL losses are 100 and 1, respectively.

Revised: We additionally replace the linear projection layer from z t to x t by an MLP with one hidden layer of size 600.

Revised: All other hyperparameters are kept as default.

Revised: On 4 out of the 6 games, D REAMER P RO obtains comparable performance to D REAMER , matching or surpassing model-free baselines that are trained for 200M environment steps.

Revised: We note that this is the ﬁrst time a reconstruction-free MBRL agent shows promising results on Atari with discrete latents.

Revised: E P ROTOTYPE VISUALIZATIONS

Revised: We visualize the ﬁrst 5 prototypes learned for each task using nearest neighbor queries, shown in each row of Figures 8 - 13.

Revised: To do so, we sample a batch of trajectories from the training replay buffer, and obtain the projection x t from latent state z t for each image.

Revised: Then, given a prototype c k as query, we show the ten images in the batch whose projections are the closest to the prototype.