DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representations

1 INTRODUCTION . 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 . 2 PRELIMINARIES . In this section , we briefly introduce the world model and learning algorithms used in DREAMERV2 ( Hafner et al. , 2021 ) which our model builds upon . To indicate the general DREAMER framework ( Hafner et al. , 2020 ; 2021 ) , we omit its version number in the rest of the paper . 2.1 RECONSTRUCTION-BASED WORLD MODEL LEARNING . DREAMER learns a recurrent state-space model ( RSSM , Hafner et al. , 2019 ) to predict forward dynamics and rewards in partially observable environments . At each time step 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 ) 2.2 POLICY LEARNING BY LATENT IMAGINATION . DREAMER interleaves policy learning with world model learning . During policy learning , the world model is fixed , and an actor and 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 . 3 DREAMERPRO . 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 can not 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 .

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