Mask-based Latent Reconstruction for Reinforcement Learning
Anonymous Author(s)
AffiliationAddress
Abstract
For deep reinforcement learning (RL) from pixels, learning effective state representations is crucial for achieving high performance.
However, in practice, limitedexperience and high-dimensional input prevent effective representation learning.
To address this, motivated by the success of masked modeling in other researchfields, we introduce mask-based reconstruction to promote state representationlearning in RL.
Specifically, we propose a simple yet effective self-supervisedmethod, Mask-based Latent Reconstruction (MLR), to predict the complete staterepresentations in the latent space from the observations with spatially and temporally masked pixels.
MLR enables the better use of context information whenlearning state representations to make them more informative, which facilitates RLagent training.
Extensive experiments show that our MLR significantly improvesthe sample efficiency in RL and outperforms the state-of-the-art sample-efficient
RL methods on multiple continuous and discrete control benchmarks.
Introduction
Learning effective state representations is crucial for reinforcement learning (RL) from visual signals(where a sequence of images is usually the input to an RL network), such as in DeepMind Control
Suite [43], Atari games [5] etc .
Inspired by the success of masked pretraining in the fields of naturallanguage processing (NLP) [11, 37, 38, 7] and computer vision (CV) [4, 21, 47], we make the firstendeavor to explore the idea of mask-based reconstruction in RL.
Masked pretraining exploits the reconstruction of masked word embeddings or pixels to promotefeature learning in NLP or CV.
This is, in fact, not straightforwardly applicable for RL due tothe following two reasons.
First, RL agents learn from interactions with environments, where theexperienced states vary as the policy network is updated.
Intuitively, collecting additional rollouts forpretraining is often costly in real-world applications.
Besides, it is challenging to learn effective staterepresentations without the awareness of the learned policy.
Second, visual signals usually have highinformation density, which may contain distractions and redundancies for policy learning.
Thus, for
RL, performing reconstruction in the original (pixel) space is not as necessary as it is in the CV or
NLP domain.
Based on the analysis above, we study the masked modeling tailored to vision-based RL.
We present
Mask-based Latent Reconstruction (MLR), a simple yet effective self-supervised method, to betterlearn state representations in RL.
Contrary to treating masked modeling as a pretraining task in thefields of CV and NLP, our proposed MLR is an auxiliary objective optimized together with the policylearning objectives.
In this way, the coordination between representation learning and policy learningis considered within a joint training framework.
Apart from this, another key difference comparedto vision/language research is that we reconstruct masked pixels in the latent space instead of theinput space.
We take the state representations ( i .
e ., features) inferred from original unmasked frames as the reconstruction targets.
This effectively reduces unnecessary reconstruction relative to thepixel-level one and further facilitates the coordination between representation learning and policylearning because the state representations are directly optimized.
Consecutive frames are highly correlated.
In MLR, we exploit this property to enable the learnedstate representations to be more informative, predictive and consistent over both spatial and temporaldimensions.
Specifically, we randomly mask a portion of space-time cubes in the input observation( i .
e ., video clip) sequence and reconstruct the information in the missing contents in the latent space.
In this way, similar to the spatial reconstruction for images in [21, 47], MLR enhances the awarenessof the agents of the global context information of the entire input observations and promotes thestate representations to be predictive in both spatial and temporal dimensions.
The global predictiveinformation is encouraged to be encoded into each frame-level state representation, achieving betterrepresentation learning and further facilitating policy learning.
Not only is an effective masked modeling method proposed, but we also conduct a systematicalempirical study for the practices of masking and reconstruction that are as applicable to RL aspossible.
First, we study the influence of masking strategies by comparing spatial masking, temporalmasking and spatial-temporal masking.
Second, we investigate the differences between masking andreconstructing in the pixel space and in the latent space.
Finally, we study how to effectively addreconstruction supervisions in the latent space.
Our contributions are summarized below:
• We introduce the idea of enhancing representation learning by mask-based reconstruction to RLto improve the sample efficiency.
We integrate the mask-based reconstruction into RL trainingwith an auxiliary objective, obviating the need for collecting additional rollouts for pretraining andhelping the coordination between representation learning and policy learning in RL.• We propose Mask-based Latent Reconstruction (MLR), a self-supervised masked modeling methodto improve the state representations for RL.
Tailored to RL, we propose to randomly mask spacetime cubes in the pixel space and reconstruct the information in the missing content in the latentspace .
This is shown to be effective for improving the sample efficiency on multiple continuousand discrete control benchmarks ( e
g ., DeepMind Control Suite [43] and Atari games [5]).• A systematical empirical study is conducted to investigate the good practices of masking andreconstructing operations in MLR for RL.
This demonstrates the effectiveness of our proposeddesigns in MLR.
Related Work2.1 Representation Learning for RL
Reinforcement learning from visual signals is of high practical value in real-world applicationssuch as robotics, video game AI, etc .
However, such high-dimensional observations may containdistractions or redundant information, imposing considerable challenges for RL agents to learneffective representations [41].
Many prior works address this challenge by taking advantage of selfsupervised learning to promote the representation learning of the states in RL.
A popular approachis to jointly learn policy learning objectives and auxiliary objectives such as pixel reconstruction[41, 50], reward prediction [23, 41], bisimulation [53], dynamics prediction [41, 13, 28, 29, 39, 52],and contrastive learning of instance discrimination [27] or (spatial -) temporal discrimination [35,2, 42, 54, 33].
In this line, BYOL [12] style auxiliary objectives, which often cooperate with dataaugmentation, show promising performance [39, 52, 48, 19, 14].
Extended introduction to BYOLand the BYOL-style objectives can be found in Appendix A.3.
Another feasible way of acquiringgood representations is to pretrain the state encoder to learn effective representations for the originalobservations before policy learning.
It requires additional offline sample collection or early accessto the environments [20, 42, 32, 31, 40], which is not fully consistent with the principle of sampleefficiency in practice.
This work aims to design a more effective auxiliary task to improve the learnedrepresentations toward sample-efficient RL.
2.2 Sample-Efficient Reinforcement Learning
Collecting rollouts from the interaction with the environment is commonly costly, especially inthe real world, leaving the sample efficiency of RL algorithms concerned.
To improve the sample
𝑀𝑀𝑀𝑀𝑀𝑀𝑀𝑀 ( � ) 𝒐𝒐 𝑡𝑡
𝒐𝒐 𝑡𝑡+
𝒐𝒐 𝑡𝑡+2 𝒐𝒐 𝑡𝑡+3 𝒐𝒐 𝑡𝑡+𝐾𝐾−1 �𝒐𝒐 𝑡𝑡 �𝒐𝒐 𝑡𝑡+1 �𝒐𝒐 𝑡𝑡+2 �𝒐𝒐 𝑡𝑡+3 �𝒐𝒐 𝑡𝑡+𝐾𝐾− 𝐴𝐴𝐴𝐴𝐴𝐴 ( � ) 𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝐸𝐸𝑂𝑂𝐸𝐸𝐸𝐸𝐸𝐸𝑂𝑂𝐸𝐸 �𝒔𝒔 𝑡𝑡 �𝒔𝒔 𝑡𝑡+1 �𝒔𝒔 𝑡𝑡+2 �𝒔𝒔 𝑡𝑡+
�𝒔𝒔 𝑡𝑡+𝐾𝐾− 𝑀𝑀𝐸𝐸𝑀𝑀𝑂𝑂𝑂𝑂𝑀𝑀𝐴𝐴𝑀𝑀𝐸𝐸𝑂𝑂𝐸𝐸𝐸𝐸𝐸𝐸𝑂𝑂𝐸𝐸 �𝒔𝒔 𝑡𝑡 �𝒔𝒔 𝑡𝑡+1 �
𝒔𝒔 𝑡𝑡+2 �𝒔𝒔 𝑡𝑡+3 �𝒔𝒔 𝑡𝑡+𝐾𝐾− 𝒂𝒂 𝑡𝑡 : 𝑡𝑡+𝐾𝐾−
Original Observation Seq.
Action Seq.
Postion Emb.
𝑃𝑃𝐸𝐸𝑂𝑂𝐸𝐸𝑂𝑂𝐸𝐸𝑀𝑀𝑂𝑂𝑃𝑃𝑂𝑂 𝐿𝐿𝑀𝑀𝑀𝑀𝑂𝑂𝑂𝑂𝑀𝑀 𝐷𝐷𝑂𝑂𝐸𝐸𝐸𝐸𝐸𝐸𝑂𝑂𝐸𝐸 𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂 𝑃𝑃𝐸𝐸𝐸𝐸𝑃𝑃 .
𝐻𝐻𝑂𝑂𝑀𝑀𝐸𝐸 �𝒚𝒚 𝑡𝑡+1 �𝒚𝒚 𝑡𝑡+2 �𝒚𝒚 𝑡𝑡+3 �𝒚𝒚 𝑡𝑡+𝐾𝐾−1 �𝒚𝒚 𝑡𝑡 𝑪𝑪𝒐𝒐𝒔𝒔𝑪𝑪𝑪𝑪𝑪𝑪𝑺𝑺𝑪𝑪𝑺𝑺𝑪𝑪𝑺𝑺𝒂𝒂𝑺𝑺𝑪𝑪𝑺𝑺𝒚𝒚
MovingAverage �𝒔𝒔 𝑡𝑡 �𝒔𝒔 𝑡𝑡+1 �𝒔𝒔 𝑡𝑡+2 �𝒔𝒔 𝑡𝑡+3 �𝒔𝒔 𝑡𝑡+𝐾𝐾− 𝒑𝒑
Momentum 𝑃𝑃𝐸𝐸𝐸𝐸𝑃𝑃 .
𝐻𝐻𝑂𝑂𝑀𝑀𝐸𝐸 𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂𝑂 𝑃𝑃𝐸𝐸𝑂𝑂𝐸𝐸 .
𝐻𝐻𝑂𝑂𝑀𝑀𝐸𝐸 �𝒚𝒚 𝑡𝑡+1 �𝒚𝒚 𝑡𝑡+2 �𝒚𝒚 𝑡𝑡+3 �𝒚𝒚 𝑡𝑡+𝐾𝐾−1 �𝒚𝒚 𝑡𝑡
Stop 𝐺𝐺𝐸𝐸𝑀𝑀𝐸𝐸𝑂𝑂𝑂𝑂𝑂𝑂𝑀𝑀 𝐴𝐴𝐴𝐴𝐴𝐴 ( � )
Masked Observation Seq.
MovingAverage 𝑡 : 𝑡+𝐾− efficiency of vision-based RL ( i .
e ., RL from pixel observations), recent works design auxiliarytasks to explicitly improve the learned representations [50, 27–29, 54, 30, 39, 51, 52] or adopt dataaugmentation techniques, such as random crop/shift, to improve the diversity of data used for trainingbased on collected samples [49, 26].
Besides, there are some model-based methods that learn (world)models in the pixel [55] or latent space [16–18, 51], and perform planning, imagination or policylearning based on the learned models.
We focus on the auxiliary task line in this work.
2.3 Masked Language/Image Modeling
Masked language modeling (MLM) [11] and its autoregressive variants [37, 38, 7] achieve significantsuccess in the NLP field and produce impacts in other domains.
MLM masks a portion of wordtokens from the input sentence and trains the model to predict the masked tokens, which has beendemonstrated to be generally effective in learning language representations for various downstreamtasks.
For computer vision (CV) tasks, similar to MLM, masked image modeling (MIM) learnsrepresentations for images/videos by pretraining the neural networks to reconstruct masked pixelsfrom visible ones.
As an early exploration, Context Encoder [36] apply this idea to Convolutional
Neural Network (CNN) model to train a CNN model with a masked region inpainting task.
Withthe recent popularity of the Transfomer-based architectures, a series of works [9, 4, 21, 47, 46] dustoff the idea of MIM and show impressive performance on learning representations for vision tasks.
Inspired by MLM and MIM, we explore the masked modeling for RL to exploit the high correlation invision data to improve agents’ awareness of global-scope dynamics in learning state representations.
Most importantly, we propose to predict the masked information in the latent space , instead of thepixel space like the aforementioned MIM works, which better coordinates the representation learningand the policy learning in RL.
Approach
3.1 Background
Vision-based RL aims to learn policies from pixel observations by interacting with the environment.
The learning process corresponds to a partially observable Markov decision process (POMDP) [6, 24],formulated as ( O , A , p, r, γ ) , where O , A , p , r and γ are the observation space ( e .
g ., pixels), theaction space, the transition dynamics p
=
Pr ( o t +1 | o ≤ t , a t ) , the reward function O×A → R and thediscount factor, respectively.
Following common practice [34], several consecutive observations arestacked to construct a state s t = { o t , o t − 1 , · · · } to convert the POMDP to an MDP [6].
The transitiondynamics and the reward function can be written as p t =
Pr ( s t +1 | s t , a t ) and r t = r
( s t , a t ) ,respectively.
The objective of RL is to learn a policy π ( a t | s t ) that maximizes the cumulativediscounted return E π (cid:80) ∞ t =
γ t r t , where γ ∈ [0 , 1] .
3.2 Mask-based Latent Reconstruction
Mask-based Latent Reconstruction (MLR) is an auxiliary objective to promote representation learningin vision-based RL and is generally applicable to different RL algorithms, e .
g ., Soft Actor-Critic(SAC) and Rainbow [22].
The core idea of MLR is to facilitate state representation learning byreconstructing spatially and temporally masked pixels in the latent space.
This mechanism enablesthe better use of context information when learning state representations, further enhancing theunderstanding of RL agents for visual signals.
We illustrate the overall framework of MLR in Figure1 and elaborate on it below.
Framework.
In MLR, as shown in Figure 1, we mask a portion of pixels in the input observationsequence along its spatial and temporal dimensions.
We encode the masked sequence and the originalsequence from observations to latent states with an encoder and a momentum encoder, respectively.
We perform predictive reconstruction from the states corresponding to the masked sequence by takingthe states encoded from the original sequence as the target.
We add reconstruction supervisionsbetween the prediction results and the targets in the decoded latent space.
The processes of masking ,encoding , decoding and reconstruction are introduced in detail below.𝒐𝒐 𝑡𝑡+3 𝒐𝒐 𝑡𝑡+ 𝒐𝒐 𝑡𝑡+ 𝒐𝒐 𝑡𝑡 ℎ 𝑤𝑤 𝑘𝑘 𝑊𝑊 𝐻𝐻 𝐾𝐾 a slight abuse of notation, for simplicitywe define the input to the deep RL networksat each timestep as o t ) and each observationcontaining d RGB frames, all the observationsin the sequence are stacked to be
the of K × H × W where H × W isthe h w .
We then randomly mask of o
+ K − 1 }
(ii) Encoding.
We adopt two encoders to learn state representations from masked observationsand original observations respectively.
A CNN-based encoder f is used to encode each maskedobservation ˜
o t
+ i into its corresponding latent state ˜s t
+ i ∈ R d .
After the encoding, we obtain asequence of the masked latent states ˜ τ
sK =
{ ˜s t , ˜s t +1 , · · · , ˜s t + K − 1 } for masked observations.
Theparameters of this encoder are updated based on gradient back-propagation in an end-to-end way.
We thus call it “online” encoder.
The state representations inferred from original observations aretaken as the targets of subsequently described reconstruction.
To make them more robust, inspired by [27, 39, 52], we exploit another encoder for the encoding of original observations.
This encoder,called “momentum” encoder as in Figure 1, has the same architecture as the online encoder, and itsparameters are updated by an exponential moving average (EMA) of the online encoder weights θ fwith the momentum coefficient m ∈ [0 , 1) , as formulated below:
[Equation]
(iii) Decoding.
Similar to the masked modeling in the CV field, e .
g ., [21, 47], the online encoder inour MLR predicts the information of the masked contents from the unmasked contents.
As opposedto pixel-space reconstruction in [21, 47], MLR performs the reconstruction in the latent space sincepixels usually have high information density [21] which may contain distractions and redundanciesfor the policy learning in RL.
Further, we consider the fact in RL that the next state is determinedby the current state as well as the action and propose to utilize the actions to reduce the possibleambiguity during prediction as an RL-tailored masked modeling design.
To this end, we adopt atransformer-based latent decoder to further infer the reconstruction results via a global messagepassing where the actions and temporal information are exploited.
Through this process, the predictedinformation is “passed” to its corresponding state representations.
As shown in Figure 3, the input tokens of the latent decoder consist of both the masked state sequence˜ τ sK ( i .
e ., state tokens) and the corresponding action sequence τ aK = { a t , a t +1 , · · · , a t + K − 1 } ( i .
e .,action tokens).
Each action token is embedded as a feature vector with the same dimension as thestated token using an embedding layer, through an embedding layer Emb ( · ) .
Following the commonpractices in transformer-based models [45], we adopt the relative positional embeddings to encoderelative temporal positional information into both state and action tokens via an element-wise additiondenoted by " + " in the following equation (see Appendix B.1 for more details).
Notably, the state andaction token at the same timestep t
+ i share the same positional embedding p t + i ∈ R d .
Thus, theinputs of latent decoder can be mathematically represented as follows:
[Equation]
The input token sequence is passed through a Transformer encoder [45] consisting of L attentionlayers.
Each layer is composed of a Multi-Headed Self-Attention (MSA) layer [45], a layer normalisation (LN) [3], and multilayer perceptron (MLP) blocks.
The process can be described as follows:
[Equation] [Equation]
The output tokens of the latent decoder, i .
e ., ˆ τ sK = { ˆs t , ˆs t +1 , · · · , ˆs t + K − 1 } , are the predictivereconstruction results for the latent representations inferred from the original observations.
Weelaborate on the reconstruction loss between the prediction results and the corresponding targets inthe following.
[Equation]
𝐾𝐾−
[Equation]
𝐸𝐸𝑇𝑇𝐸𝐸𝑇𝑇𝐸𝐸𝑇𝑇𝑇𝑇
PostionEmb
[Equation]
𝒑𝒑 �𝒔𝒔 𝑡𝑡
[Equation]
+1 �𝒔𝒔 𝑡𝑡+2 �𝒔𝒔 𝑡𝑡+3 �𝒔𝒔 𝑡𝑡+𝐾𝐾−1 𝒂𝒂 𝑡𝑡 : 𝑡𝑡+𝐾𝐾−1 ActionSeq.
𝐸𝐸𝑇𝑇𝐸𝐸 ( � )
StateTokens �𝒔𝒔 𝑡𝑡 �𝒔𝒔 𝑡𝑡+ �𝒔𝒔 𝑡𝑡+ �𝒔𝒔 𝑡𝑡+ �
𝒔𝒔 𝑡𝑡+𝐾𝐾−
(iv) Reconstruction.
Motivated by the successof BYOL [12] in self-supervised learning, weuse an asymmetric architecture for calculatingthe distance between the predicted/reconstructedlatent states and the target states, similar to[39, 52].
For the outputs of the latent decoder, we use a projection head g and a prediction head q to get the final prediction result ˆy t
+ i =
q
( g ( ˆs t + i ))
corresponding to s t
+ i .
For the encoded results of original observations,we use a momentum-updated projection head¯ g whose weights are updated with an EMA ofthe weights of the online projection head.
Thesetwo projection heads have the same architectures.
The outputs of the momentum projection head ¯ g ,i .
e .,
¯y t
+ i = ¯ g ( ¯s t + i ) , are the final reconstruction targets.
Here, we apply a stop-gradient operation asillustrated in Figure 1 to avoid model collapse, following [12].
The objective of MLR is to enforce the final prediction result ˆy t
+ i to be as close as possible to itscorresponding target ¯y t
+ i .
To achieve this, we design the reconstruction loss in our proposed MLRby calculating the cosine similarity between ˆy t
+ i and ¯y t
+ i , which can be formulated below:
[Equation]
The loss L mlr is used to update the parameters of the online networks, including encoder f , predictivelatent decoder ϕ , projection head g and prediction head q .
Through our proposed self-supervisedauxiliary objective MLR, the learned state representations by the encoder will be more informativeand thus can further facilitate policy learning.
Objective.
The proposed MLR is an auxiliary task, which is optimized together with the policylearning.
Thus, the overall loss function L total for RL agent training is:
[Equation]
(6) where L rl and L mlr are the loss functions of the base RL agent ( e .
g ., SAC [15] and Rainbow [22])and the proposed mask-based latent reconstruction, respectively.
λ is a hyperparameter for balancingthe two terms.
Notably, the agent of vision-based RL commonly consists of two parts, i .
e ., the (state)representation network ( i .
e ., encoder) and the policy learning network.
The encoder of MLR is takenas the representation network to encode observations into the state representations for RL training.
The latent decoder is not adopted because the unmasked observations are taken as the inputs.
Moredetails can be found in Appendix B.
Experiment4.1 Setup
Environments and Evaluation.
We evaluate the sample efficiency of our MLR on both the continuous control benchmark DeepMind Control Suite (DMControl) [43] and the discrete controlbenchmark Atari [5].
On DMControl, following the previous works [27, 49, 26, 52], we choose sixcommonly used environments from DMControl, i .
e ., Finger, spin ; Cartpole, swingup ; Reacher, easy ;
Cheetah, run ; Walker, walk and Ball in cup, catch for the evaluation.
We test the performance of
RL agents with the mean and median scores over 10 episodes at 100k and 500k environment steps,referred to as DMControl-100k and DMControl-500k benchmarks, respectively.
The score of eachenvironment ranges from 0 to 1000 [43].
For discrete control, we test agents on the Atari-100k benchmark [55, 44, 25, 27] which contains 26 Atari games and allows the agents 100k interaction steps( i .
e ., 400k environment steps with action repeat of 4) for training.
Human-normalized score (HNS) 1is used to measure the performance on each game.
Considering the high variance of the scores onthis benchmark, we test each run over 100 episodes [39].
To achieve a more rigorous evaluation onhigh-variance benchmarks such as Atari-100k with a few runs, recent work Rliable [1] systematicallystudies the evaluation bias issue in deep RL and recommends robust and efficient aggregate metricsto evaluate the overall performance (across all tasks and runs), e .
g ., interquartile-mean (IQM) andoptimality gap (OG) 2 , with percentile confidence intervals (CIs, estimated by the percentile bootstrapwith stratified sampling).
We follow the recommendation and report the aggregate metrics on the
Atari-100k benchmark with 95% CIs.
Implementation.
SAC [15] and Rainbow [22] are taken as the base continuous-control agent anddiscrete-control agent, respectively (see Appendix A for the details of the two algorithms).
In ourexperiments, we denote the base agents trained only by RL loss L rl (as in Equation 6) as Baseline ,while denoting the models of applying our proposed MLR to the base agents as MLR for the brevity.
Note that compared to naive SAC or Rainbow on pixels, our Baseline additionally adopts dataaugmentation (random crop and random intensity).
This practice follows prior works [26, 49] whichuncovers that applying proper data augmentation can significantly improve the sample efficiency of
SAC or Rainbow. As shown in Equation 6, we set a weight λ to balance L rl and L mlr so that thegradients of these two loss items lie in a similar range and empirically find λ = 1 works well in mostenvironments.
In MLR, by default, we set the length of a sampled trajectory K to 16 and mask ratioη to 50%.
We set the size of the masked cube ( k × h × w ) to 8 × 10 × 10 on most DMControl tasksand 8 × 12 × 12 on the Atari games.
More implementation details can be found in Appendix B.
4.2 Comparison with State-of-the-Arts
DMControl.
We compare our proposed MLR with the state-of-the-art (SOTA) sample-efficient RLmethods proposed for continuous control, including PlaNet [16], Dreamer [17], SAC+AE [50], SLAC
Our method augments Baseline with the proposed MLR objective (denoted as MLR )
[Table]
±865 ±871 ±872 ± 5 Reacher, easy 210 ± 390 793 ±627 ±- 929 ±942 ±942 ±943 ±957 ± 41 Cheetah, run 305 ± 131 570 ±550 ±640 ±518 ±660 ±719 ±602 ±674 ± 37 Walker, walk 351 ±897 ±847 ±842 ±902 ±921 ±928 ±818 ±939 ± 10 Ball in cup, catch 460 ±879 ±794 ±852 ±959 ±963 ±967 ±960 ±964 ±
Mean 393.
896.5 Median 405.
[Figure]
Gap
[Figure]
[28], CURL [27], DrQ [49] and PlayVirtual [52].
The comparison results are shown in Table 1, andall results are averaged over 10 random seeds.
We can observe that: (i) MLR significantly improves
Baseline in both DMControl-100k and DMControl-500k, and achieves consistent gains relative to
Baseline across all environments.
It is worthy to mention that, for the DMControl-100k, our proposedmethod outperforms Baseline by 25.3% and 34.7% in mean and median scores, respectively.
Thisdemonstrates the superiority of MLR in improving the sample efficiency of RL algorithms.
(ii)
The
RL agent equipped with MLR outperforms most state-of-the-art methods on DMControl-100k and
DMControl-500k.
Specifically, our method surpasses the best previous method ( i .
e .,
PlayVirtual)by 43.5 and 35.5 respectively on mean and median scores on DMControl-100k.
Our method alsodelivers the best median score and reaches a comparable mean score with the strongest SOTA methodon DMControl-500k.
Atari-100k.
We further compare MLR with the SOTA model-free methods for discrete control,including DER [44], OTR [25], CURL [27], DrQ [49], DrQ( ϵ ) (DrQ using the ϵ -greedy parameters in[8]) , SPR [39] and PlayVirtual [52].
These methods and our MLR for Atari are all based on Rainbow[22].
The Atari-100k results are shown in Figure 4.
MLR achieves an interquartile-mean (IQM)score of 0.432, which is 28.2% higher than SPR (IQM: 0.337) and 15.5% higher than PlayVirtual(IQM: 0.374).
This indicates that MLR has the highest sample efficiency overall.
For the optimalitygap (OG) metric, MLR reaches an OG of 0.522 better than SPR (OG: 0.577) and PlayVirtual (OG:0.558), showing that our method performs closer to the desired target, i .
e ., human-level performance.
The full scores of MLR (over 3 random seeds) across the 26 Atari games and more comparisons andanalysis can be found in Appendix C.
4.3 Ablation Study
In this section, we investigate the effectiveness of our MLR auxiliary objective, the impact of maskingmethods and the model/architecture design.
Unless otherwise specified we conduct the ablation and
[Table]
Effectiveness evaluation.
Besides the comparison with SOTA methods, we demonstrate the effectiveness of our proposed MLR by studying its improvements compared to our Baseline.
The numericalresults are presented in Table 1, while the curves of test performance during the training process aregiven in Appendix C. Both numerical results and test performance curves can demonstrate that ourmethod obviously outperforms Baseline across different environments thanks to more informativerepresentations learned by MLR.
Masking strategy.
We compare three design choices of the masking operation: (i)
Spatial masking(denoted as MLR-S ): we randomly mask patches for each frame independently.
(ii) Temporal masking(denoted as MLR-T ): we divide the input observation sequence into multiple segments along thetemporal dimension and mask out a portion of segments randomly.
(Here, the segment length isset to be equal to the temporal length of cube, i .
e ., k .)
(iii) Spatial-temporal masking (also referredto as “cube masking”): as aforementioned and illustrated in Figure 2, we rasterize the observationsequence into non-overlapping cubes and randomly mask a portion of them.
Except for the differencesdescribed above, other configurations for masking remain the same as our proposed spatial-temporal( i .
e ., cube) masking.
From the results in Table 2, we have the following observations: (i) All threemasking strategies ( i .
e ., MLR-S , MLR-T and MLR ) achieve mean score improvements comparedto Baseline by 18.5% , 12.2% and 25.0% , respectively, and achieve median score improvementsby 23.4% , 25.0% and 35.9% , respectively.
This demonstrates the effectiveness of the core idea ofintroducing mask-based reconstruction to improve the representation learning of RL.
(ii)
Spatialtemporal masking is the most effective strategy over these three design choices.
This strategy matchesbetter with the nature of video data due to its spatial-temporal continuity in masking.
It encourages thestate representations to be more predictive and consistent along the spatial and temporal dimensions,thus facilitating the policy learning in RL.
Reconstruction target.
In masked language/image modeling, reconstruction/prediction is commonlyperformed in the original signal space, such as word embeddings or RGB pixels.
To study theinfluence of the reconstruction targets for the task of RL, we compare two different reconstructionspaces: (i) Pixel space reconstruction (denoted as MLR-Pixel ): we predict the masked contentsdirectly by reconstructing original pixels, like the practices in CV and NLP domains; (ii) Latentspace reconstruction ( i .
e ., MLR ): we reconstruct the state representations ( i .
e ., features) of originalobservations from masked observations, as we proposed in MLR.
Table 2 shows the comparisonresults.
The reconstruction in the latent space is superior to that in the pixel space in improvingsample efficiency.
As discussed in the preceding sections, vision data might contain distractions andredundancies for policy learning, making the pixel-level reconstruction unnecessary.
Besides, latentspace reconstruction is more conducive to the coordination between the representation learning andthe policy learning in RL, because the state representations are directly optimized.     
 
                     
     0HDQ0HGLDQ Mask ratio.
In recent works of masked imagemodeling [21, 47], the mask ratio is found to becrucial for the final performance.
We study theinfluences of different masking ratios on sample efficiency in Figure 5 (with 3 random seeds)and find that the ratio of 50% is an appropriatechoice for our proposed MLR.
An over-smallvalue of this ratio could not eliminate redundancy, making the objective easy to reach byextrapolation from neighboring contents that are free of capturing and understanding semantics from visual signals.
An over-large value leaves toofew contexts for achieving the reconstruction goal.
As discussed in [21, 47], the choice of this ratiovaries for different modalities and depends on the information density of the input signals.
Action token.
We study the contributions of action tokens used for the latent decoder in Figure 1 bydiscarding it from our proposed framework.
The results are given in Table 3.
Intuitively, predictiononly from visual signals is of more or less ambiguity.
Exploiting action tokens benefits by reducingsuch ambiguity so that the gradients of less uncertainty can be obtained for updating the encoder.
[Table]
Masking features.
We compare “masking pixels” and “masking features” in Table 3.
Maskingfeatures (denoted by MLR-F ) does not perform equally well compared with masking pixels asproposed in MLR, but it still achieves significant improvements to Baseline .
Why not use a latent decoder for targets?
We have tried to add a momentum-updated latent decoderfor obtaining the target states ¯s t
+ i as illustrated in Figure 1.
The results in Table 3 show that addingthe momentum decoder leads to performance drops.
Since there is no information prediction for thestate representation learning from the original observation sequence without masking, we do not needto refine the implicitly predicted information like that in the outputs ˆs t
+ i of the online encoder.
Decoder depth.
We analyze the influence of using Transformer-based latent decoders of differentdepths.
The experimental result is shown in Appendix C. Generally, deeper latent decoders lead toworse sample efficiency with lower mean and median scores.
Similarity loss.
We compare models using two kinds of similarity metrics to measure the distance inthe latent space and observe that using cosine similarity loss is better than using mean squared error(MSE) loss.
The results and analysis can be found in Appendix C.
Projection and prediction heads.
Previous works [12, 39, 10] have shown that in self-supervisedlearning, supervising the feature representations in the projected space via the projection/predictionheads is often better than in the original feature space.
We investigate the effect of the two heads andfind that both improve agent performance (see Appendix C).
Sequence length and cube size.
These two factors can be viewed as hyperparameters.
Theircorresponding experimental analysis and results are in Appendix C.
Conclusion
In this work, we make the first effort to introduce the mask-based reconstruction to RL for facilitatingpolicy learning by improving the learned state representations.
We propose MLR, a simple yeteffective self-supervised auxiliary objective to reconstruct the masked information in the latent space.
In this way, the learned state representations are encouraged to include richer and more informativefeatures.
Extensive experiments demonstrate the effectiveness of MLR and show that MLR achievesthe state-of-the-art performance on both DeepMind Control and Atari benchmarks.
We conduct adetailed ablation study for the proposed designs in MLR and analyze their differences from thatin NLP and CV domains.
We hope our proposed method can inspire further research for visionbased RL from the perspective of improving representation learning.
Moreover, the masked latentreconstruction concept is also worthy of being explored and extended in computer vision and naturallanguage processing.
We are looking forward to seeing more mutual promotion between differentresearch fields.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Appendix C.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] SeeAppendix D.
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main ex-
perimental results (either in the supplemental material or as a URL)?
[Yes] See the
supplemental material.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they were chosen)?
[Yes] See Appendix B.
(c) Did you report error bars (e
g., with respect to the random seed after running experi-
ments multiple times)? [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix B.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] See Appendix B.
(b) Did you mention the license of the assets?
[Yes] See Appendix B.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
See the supplemental material.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[Yes] See Appendix B.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[Yes] See Appendix B.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]