Mask-based Latent Reconstruction for Reinforcement Learning
Authors: Tao Yu
Zhizheng Zhang
Cuiling Lan
Yan Lu
Zhibo Chen
Venue: NeurIPS 2022 Accept (2022)
Keywords: Reinforcement learning
mask-based modeling
sample efficiency
representation learning

Abstract: For deep reinforcement learning (RL) from pixels, learning effective state representations is crucial for achieving high performance. However, in practice, limited experience and high-dimensional inputs prevent effective representation learning. To address this, motivated by the success of mask-based modeling in other research fields, we introduce mask-based reconstruction to promote state representation learning in RL. Specifically, we propose a simple yet effective self-supervised method, Mask-based Latent Reconstruction (MLR), to predict complete state representations in the latent space from the observations with spatially and temporally masked pixels. MLR enables better use of context information when learning state representations to make them more informative, which facilitates the training of RL agents. Extensive experiments show that our MLR significantly improves the sample efficiency in RL and outperforms the state-of-the-art sample-efficient RL methods on multiple continuous and discrete control benchmarks. Our code is available at https://github.com/microsoft/Mask-based-Latent-Reconstruction.

Latent Reconstruction for Reinforcement Learning
Tao Yu!+t Zhizheng Zhang Cuiling Lan² Yan Lu² Zhibo Chen'
University of Science and Technology of China "Microsoft Research Asia
yutao6660mail. ust.. edu, cn, (zhizzhang, culan, wwt.yyyllllmcccssssaal com chenzhibo@ustc. edu.
Abstract
For deep reinforcement learning (RL) from pixels, learning effective state repre- sentations crucial for achieving high performance. However, in practice, limited experience and high-dimensional inputs prevent effective representation learning. To address this, motivated by the success mask-basee based modeling in other research fields, we introduce mask-based reconstruction to promote state representation learning in RI.. Specifically, we propose a simple yet effective self-supervised method, Mask-based Latent Reconstruction (MLR), to predict complete state rep- resentations in the latent space from the observations with spatially and temporally masked pixels. MLR enables better use of context information when learning state representations to make them more informative, which facilitates the training of RL agents. Extensive experiments show that MLR significantly improves the sample efficiency in RL and outperforms the state-of setee--ffthheaaee art sample-efficient RL methods on multiple continuous and discrete control benchmarks. Our code available at
Introduction
Learning effective state representations is crucial for reinforcement learning (RL) from visual signals (where sequence of images is usually the input of an RL agent), such as in DeepMind Control Suite
43 Atari games [5] Inspired by the success of mask-hased pretraining in the fields of natural
language processing (NLP) [II 37 38 7 and computer vision (CV) 21 47
we make the first
endeavor to explore the idea of mask-based modeling in RL.
Mask-based pretraining exploits the reconstruction of masked word embeddings or pixels to promote feature learning in NLP or CV fields. This is, in fact, not straightforwardly applicable for RL due to the following two reasons. First, RL agents learn policies from the interactions with environments, where the experienced states vary the policy network updated. Intuitively, collecting additional rollouts for pretraining is often costly especially in real-world applications. Besides, it is challenging learn effective state representations without the awareness of the learned policy. Second, visual signals usually have high information densities, which may contain distractions and redundancies for policy learning. Thus, for RL, performing reconstruction in the original (pixel) space not as necessary as it is in the CV or NLP fields.
Based on the analysis above, study the mask-based modeling tailored vision-based RL. We present Mask-basse Latent Reconstruction (MLR), a simple yet effective self-supervised method, to better learn state representations in RL. Contrary to treating mask- based modeling pretraining task in the fields of CV and NLP, our proposed MLR is an auxiliary objective optimized together
Equual contribution.
This work done when Tao Yu was intern at Microsoft Research Asia.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

with the policy learning objectives. In this way, the coordination between representation learning and policy learning is considered within a joint training framework. Apart from this, another key difference compared to vision/language research is that we reconstruct masked pixels in the latent space instead of the input 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 the pixel-Ievel one and further facilitates the coordination between representation learning and policy learning because the state representations are directly optimized.
Consecutive frames are highly correlated. In MLR, we exploit this property to enable the learned state representations to be more informative, predictive and consistent over both spatial and temporal dimensions. Specifically,, we randomly mask a portion of space-time cubes in the input observation (i.e., video clip) sequence and reconstruct the feature representations of the missing contents in the latent space. In this way, similar to the spatial reconstruction for images in 21 the awareness of the agents on the global context information of the entire input observations and promotes the state representations to predictive in both spatial and temporal dimensions. The global predictive information is encouraged to be encoded into each frame-level state representation, achieving better representation learning and further facilitating policy learning.
MLR enhances
We not only propose an effective mask-based modeling method, but also conduct systematical empirical study for the practices of masking and reconstruction that are as applicable to RL as possible. First, we study the influence of masking strategies by comparing spatial masking. temporal masking and spatial-temporal masking. Second, we investigate the differences between masking and reconstructing in the pixel space and in the latent space. Finally, we study how to effectively add reconstruction supervisions in the latent space.
Our contributions are summarized below:
We introduce the idea of enhancing representation learning by mask-based modeling to RL to improve the sample efficiency. We integrate the mask-based reconstruction into the training of RL an auxiliary objective, obviating the need for collecting additional rollouts for pretraining and helping the coordination between representation learning and policy learning in RL.
We propose Mask-based Latent Reconstruction (MLR), self-supervised mask-based modeling method to improve the state representations for RI.. Tailored to RI., we propose to randomly mask space-time cubes in the pixel space and reconstruct the information of missing contents in the latent space. This is shown to be effective for improving the sample efficiency on multiple continuous and 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 and reconstructing operations in MLR for RL. This demonstrates the effectiveness of our proposed designs in the proposed MLR.
2 Related Work
2.1 Representation Learning for RL
Reinforcement learning from visual signals is of high practical value in real-world applications such as robotics, video game AI, etc. However, such nigh-dimensiooal observations may contain distractions or redundant information, imposing considerable challenges for RL agents learn effective representations [41] Many prior works address this challenge by taking advantage of self- supervised learning to promote the representation learning of the states in RL. A popular approach is to jointly learn policy learning objectives and auxiliary objectives such as pixel reconstruction
41 [5), reward prediction [23] 41) bisimulation (53), dynamics prediction
I3 28 29 39
52]
and contrastive learning of instance discrimination [27] or (spatial -) temporal discrimination [35]
2 42 54 33 In this line, BYOL [2] -syyle auxiliary objectives, which are often adopted with
data augmentation, show promising performance (39 52 48 19 14). More detailed introduction for BYOL and the BYOL-style objectives can be found in Appendix A.3. Another feasible way of acquiring good representations is to pretrain the state encoder to learn effective state representations for the original observations before policy learning. requires additional offline sample collection or
early access to the environments 20 of sample efficiency in practice. This work aims to design a more effective auxiliary task to improve the learned representations toward sample-efficien RL.
32 31 40) which is not fully consistent with the principle

---

2.2 Sample-efficient Reinforcement Learning
Collecting rollouts from the interaction with the environment is commonly costly, especially in the real world, leaving the sample efficiency of RL algorithms concerned. To improve the sample efficiency of vision-based RL (i.e., RL from pixel observations), recent works design auxiliary tasks to explicitly improve the learned representations [50] (2 [54 30,
(39 [1 52 or adopt data
augmentation techniques, such as random crop/shift, to improve the diversity of data used for training
49 (6]. 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 policy learning 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) [1] and its autoregressive variants (37 38 achieve significant success in the NLP field and produce impacts in other domains. MLM masks a portion of word tokens from the input sentence and trains the model to predict the masked tokens, which has been demonstrated to be generally effective in learning language representations for various downstream tasks. For computer vision (CV) tasks, similar to MLM, masked image modeling (MIM) learns representations for images/videos by pretraining the neural networks to reconstruct masked pixels from 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. With the recent popularity of the Transfomer-based architectures, series of works 9 4 21 47 46] dust off the idea of MIM and show impressive performance on learning representations for vision tasks. Inspired by MLM and MIM, we explore the mask-based modeling for RL to exploit the high correlation in vision data to improve agents' awareness of global-scope dynamics in learning state representations. Most importantly, space, instead of the pixel space like the aforementioned MIM works, which better coordinates the representation learning and the policy learning in RL.
propose to predict the masked information in the latent
3 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) 24).
formulated as (O, A. p. r. 7). where O, A. p, r and ry denote the observation space (e.g., pixels), the action space, the transition dynamics p Pr(oc-1lo<: az), the reward function O x A - R, and the discount factor, respectively. Following common practices [34], an observation o, is constructed by few RGB frames. The transition dynamics and the reward function can be formulated as Pt = Pr(o.41 oo a.) and rt = r (o,, az), respectively. The objective of RL is to learn a policy (a, o,) that maximizes the cumulative discounted return E, Eizo 7'T,, where y 10,
3.2 Mask-based Latent Reconstruction
Mask-based Latent Reconstruction (MLR) is an auxiliary objective to promote representation learning in vision-based RL and is generally applicable for different RL algorithms, e.g., Soft Actor-Critic Critic (SAC) and Rainbow [22] The core idea of MLR is to facilitate state representation learning by recon- structing spatially and temporally masked pixels in the latent space. This mechanism enables better use of context information when learning state representations, further enhancing the understanding of RL agents for visual signals. We illustrate the overall framework of MLR in Figure I and elaborate on it below.
Framework. In MLR, shown in Figure 1 we mask portion of pixels in the input observation sequence along its spatial and temporal dimensions. We encode the masked sequence and the original sequence from observations latent states with an encoder and a momentum encoder, respectively. We perform predictive reconstruction from the states corresponding to the masked sequence while taking the states encoded from the original sequence as the target. We add reconstruction supervisions between 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

---

Masked Observation Seq.
Original Observation Seq.
Mask
0:13
Oc+X-1
Aug()
Aug(()
Moving Average
Online Encoder
Momentum Encoder
Sees Sce2 R:43 K
8:41 Sea2 S1a-11
4:
Action Seq.
Postion Emb.
Predictive Latent Decoder
& Stes 8eez SLOK-1
Moving Average
Online Proj.Heat
Momentum Proj. Head
Online Pred, Head
Stop Gradient
Cosine Similarity
Jiez Pers
Yeok-1
Pe-1 PerK-1
Figure 1: The framework of the proposed MLR. We perform random spatial-temporal masking (i.e., cube masking) the sequence of consecutive observations in the pixel space. The masked observations are encoded to be the latent states through an online encoder. We further introduce a predictive latent decoder decode/predict the latent states conditioned the corresponding action sequence and temporal positional embeddings. Our method trains the networks to reconstruct the feature representations of the missing contents in appropriate lalent space using cosine similarity based distance metric applied hetween the predicted features of the reconstructed states and the target features inferred from original observations by momentum networks.
(i) Masking. Given an observation sequence of K timesteps Th {o,, 0.+1.-.- 0.++ -1) and each observation containing n RGB frames, all the observations in the sequence to be a cuboid with the shape K x H x W where H W is the spatial size. Note that the actual shape is K x x W x D where D 3n is the number of channels in each ob- servation. We omit the channel dimension for simplicity. As illustrated in Figure 2 we divide the cuboid into regular non-overlapping cubes with the shape of h x w. We then ran- domly mask a portion of the cubes following uniform distribution and obtain a masked obser- vation sequence 7g {ōr- 0.41, -1 Following [52]. we perform stochastic image augmentation Aug(.) (e.g., random crop and in- tensity) on each masked observation 7-- The objective of MLR to predict the state representations of the unmasked observation sequence from the masked one in the latent space.
stacked
W
h
0.+3
0r+2
0c+1
Oc
Figure 2: Illustaaionn of our cube masking. We divide the input observation sequence to non- overlapping cubes (I: x h x w). In this example, have k = K h H and the observation sequence has K = timesteps and spatial size of H x W.
1 W where

---

(ii) Encoding. We adopt two encoders to learn state representations from masked observations and original observations respectively. A CNN-based encoder f is used to encode each masked observation Oc+; into its corresponding latent state Sets E R. After the encoding, we obtain a sequence of the masked latent states FK = {S,, S1+1. K[+K-1] for masked observations. The parameters of this encoder are updated based on gradient pack-propagation in an end-to-end way. We thus call it "online" encoder. The state representations inferred from original observations are taken 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 I has the same architecture as the online encoder, and its parameters updated by exponential moving average (EMA) of the online encoder weights By with the momentum coefficient m E [0, 1), as formulated below:
0, f m0, + (1-11m)) m-)0-
the online encoder in our MLR predicts the information of the masked contents from the unmasked contents. As opposed to pixel-space reconstruction in [2 47. MLR performs the reconstruction in the latent space to better perform RL policy learning. Moreover, pixels usually have high information densities [2T] which may contain distractions and redundancies for the policy learning in RL. Considering the fact that in RL the next state is determined by the current state as well as the action, we propose to utilize both the actions and states as the inputs for prediction to reduce the possible ambiguity to enable an RL-tailored mask-based modeling design. To this end, adopt a transformer- based latent decoder to infer the reconstruction results via a global message passing where the actions and temporal information are exploited. Through this process, the predicted information is "passed" to its corresponding state representations.
(iii) Decoding. Similar to the mask-based modeling in the CV field, e.g.. 21 47
As shown in Figure [] the input tokens of the latent decoder consist of both the masked state sequence -k (i.e., state tokens) and
State Tokens
the corresponding action sequence Th {ar, a,+1, K (i.e., action tokens).
Action Seq. deree--
Emb
A0D 1-00--000 9; A.i3
Each action token is embedded as a feature vec- tor with the same dimension the stated to- ken using an embedding layer, through an em- bedding layer Emb(-). Following the common practices in transformer-basee models [45), we adopt the relative positional embeddings code relative temporal positional information into both state and action tokens with element- wise addition denoted by "4" in the following equation (see Appendix B.1 for more details). Notably, the state and action token at the same timestep + share the same positional embed- ding Pt-r E Rdd. Thus, the inputs of latent decoder can be mathematically represented as follows:
Action Tokens
POOO-O Mes2 Blaz-1
6: Beet Au
Transformer
Bcov-1
Figure 3: Illustration of predictive latent decoder.
L a, S1+1, A,+1,"" -, + Pp: Pt, Pc+1, Pet1,'' 1P++K-1,PP4K' (2)
The input token sequence is fed into Transformer encoder 45 consisting of L attention layers. Each layer is composed of a Multi-Headed Self-Attention (MSA) layer [45]. a layer normalisation (LN) B) and multilayer perceptron (MLP) blocks. The process can be described as follows:
z' MSA(LNNx*)) x,
x+1 MLP(LN(22)) + +
(4)
The output tokens of the latent decoder, i.e., A% {8,, S S are the predictive reconstruction results for the latent representations inferred from the original observations. We elaborate on the reconstruction loss between the prediction results and the corresponding targets in the following.
in self-supervised learning, we use an asymmetric architecture for calculating the distance hetween the predicted/recoonstuutedd latent states and the target states, similar to [39 (2). For the outputs of the latent decoder, we use projection
(iv) Reconstruction. Motivated by the success of BYOL 12
5

---

head and prediction head to get the final prediction result Yi++ ((g($++:)) corresponding to St-q- For the encoded results of original observations, we use a momentum-upee uupuuttteddddddddoeeeoee-aae projection head 9 whose weights are updated with an EMA of the weights of the online projection head. 'heee two projection heads have the same architectures. The outputs of the momentum projection head g- Le., Vz+i = (5++)) are the final reconstruction targets. Here, we apply stop-gradient operation as illustrated in Figure T to avoid model collapse, following [12].
The objective of MLR is to enforce the final prediction result Se+i to be as close as possible to its corresponding target Ve+i- To achieve this, we design the reconstruction loss our proposed MLR by calculating the cosine similarity between Ve+i and Ve+i, which can be formulated below:
Ye++
Lmir
(5)
18t+ell2 1[
The loss Limrr is used to update the parameters of the online networks, including encoder f. predictive latent decoder é. projection head g and prediction head q. Through our proposed self-supervised auxiliary objective MLR, the learned state representations by the encoder will be more informative and thus can further facilitate policy learning.
Objective. The proposed MLR serves as an auxiliary task, which is optimized together with the policy learning. Thus, the overall loss function Liotal for RL agent training is:
Liotel C+t + XCmir,
(6)
where Lyt and Lmir are the loss functions of the base RL agent (e.g., SAC [5] and Rainbow (2
and the proposed mask-based latent reconstruction, respectively. À is hyperparameter for balancing the two terms. Notably, the agent of vision- based RI. commonly consists of two parts, i.e., the (state) representation network (i.e., encoder) and the policy learning network. The encoder of MLR is taken as the representation network to encode observations into the state representations for RL training. The latent decoder can be discarded during testing since it is only needed for the optimization with our proposed auxiliary objective during training. More details can be found Appendix
4 Experiment
4.1 Setup
Environments and Evaluation. We evaluate the sample efficiency of our MLR on both the con- tinuous control benchmark DeepMind Control Suite (DMControl) (43 and the discrete control benchmark Atari [5]. On DMControl, following the previous works (27 49 26 52), we choose six commonly 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-1000 and DMControl-500k benchmarks, respectively. The score of each environment ranges from 0 to 1000 43 For discrete control, we test agents the Atari-100k benchmark 55 23 27 which contains 26 Atari games and allows the agents 100k interaction steps (Le., 400k environment steps with action repeat of 4) for training. Human-normalize score (HNS) i is used to measure the performance on each game. Considering the high variance of the this benchmark, we test each run over 100 episodes [39]. To achieve a more rigorous evaluation on high-variance benchmarks such as Atari- 100k with a few runs, recent work Rliable systematically studies the evaluation bias issue in deep RL and recommends robust and efficient aggregate metrics to evaluate the overall performance (across all tasks and runs), e.g., interquartile- with percentile confidence intervals (CIs, estimated by the
mean (IQM) and optimality gap (OG) R percentile bootstrap with stratified sampling). We follow aforementioned common practices and report the aggregate metrics on the Atari- 100k benchmark with 95% Cls.
HHNS calculated by SA-Sn, where SA. SR and Sh are the scores the agent, random play and the SH-SR'
expert human, respectively.
'IQM discards the top and bottom 25% of the runs and calculates the mean score of the remaining 50% runs. OG is the amount by which the agent fails to meet a minimum score of the default buman-leva performance. Higher IQM and lower OG are better.

---

Table I: Comparison results (mean + std) the DMControl- 100k and DMControl- benchmarks.
Our method augments Baseline with the proposed MLR objective (denoted as MLR).
100k Step Scores
PlaNet
Dreamer
SAC+AR
SLAC
CURL
DrQ
Play Virtual
Raseline
MI.K
Finger, spin
136 = 216
341 70
740 64
693
141
767 = 56
9X1 104
915
853 + 112
907 +
Campole, swingup
297
19
326 27
311 11
582 + 146
759 + 92
816 36
784 +: 6.1
2066 + -48
Reacher, casy
314 155
233
593 118
Cheetab, run
138 88
235
4: 137
267
319
299 48
344 2:
474 2:
399 80
482
Walker, walk
224 48
277 12
194 22
361 = 73
401 24
612 =12==10 164
460 173
424 + 281
643 + 114
cup, catch
246
174
926
4:
Mcun
135.8
289.8
396.2
471.3
$59.7
688.3
729.3
616.8
772.8
Median
137.0
295.5
351.0
436.5
560.0
685.5
800.5
620.5
836.0
500k Step Scores
Finger, spin
561 284
796 +: 183
884 128
673
926 45
103
963
944 97
973
Cartpole, swingup
475 71
762
27
735
841 45
868 +:
865
871
872
Reacher, easy
210 =: 390
793
-: 164
627 58
929
41
942 : 71
942 = 66
943 52
957 + 41
Cheetah, run
305
131
570
253
550
34
640
518 28
660 +:
719
602 4: 67
674 d:
Walker, walk
351 58
49
847
842
902 43
921 + 4575
928
818 + 263
9:9
Ball in catch
160 +: 380
879 87
791 = 58
852 = 71
959 27
963 :
967 2:
960 :: 10
961 + 14
Mcan
393.7
782.8
739.5
845.8
856.3
896.5
Median
405.5
794.5
764.5
757.5
914.0
929.5
935.0
907.0
948.0
Implementation. SAC [15] and Rainbow [22] are taken as the base continuous-control agent and discrete-control agent, respectively (see Appendix A.1 and A.2 for the details of the two algorithms). In our experiments, we denote the base agents trained only by RL loss Z,t (as in Equation (6) as Baseline, while denoting the models of applying our proposed MLR to the base agents MLR for the brevity. Note that compared to naive SAC or Rainbow, our Baseline additionally adopts data augmentation (random crop and random intensity). We adopt this following the prior works which uncovers that applying proper data augmentation can significantly improve the sample efficiency of SAC or Rainbow. As shown in Equation 6 set weight À to balance L, and Lmtr so that the gradients of these two loss items lie in a similar range and empirically find X = 1 works well in most environments. In MLR, by default, we set the length of sampled trajectory K to 16 and mask ratio to 50%. We set the size of the masked cube (k x h X w) to 8 x 10 x 10 on most DMControl tasks and 12 12 on the Atari games. More implementation details can be found in Appendix B.
[26 49
4.2 Comparison with State-of-the-Arts
DMControl. We compare our proposed MLR with the state-of-the-art (SOTA) sample-efficient
RL methods proposed for continuous control, including PlaNet (16), Dreamer [7., SAC+AE (50),
SLAC (28), CURL (27). DrQ 49 and Virtual [52). The comparison results are shown in Table I and all results are averaged over 10 runs with different random seeds. We can observe that: (i)
MLR significantly improves Baseline in both DMComtrol-100k and DMControl-500K, and achieves consistent gains relative to Baseline across all environments. It is worthy to mention that, for the OMControl-1008 our proposed method outperforms Baseline by 25.3 % and 34.7% % in mean and median scores, respectively. This demonstrates the superiority of MLR in improving the sample efficiency of RL algorithms. (ii) The RL agent equipped with MLR outperforms most state-of-thh--a art and DMControl-
methods on DMControl-100 previous method (i.e., Virttua) by 43.5 and 35.5 respectively on mean and median scores on DMControl- 100k. Our method also delivers the best median score and reaches a comparable mean score with the strongest SOTA method on DMControl-500k.
Specifically, our method surpasses the best
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(e) (DrQ using the e-greedy parameters in
[ SPR 39 and Play Virtual [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-ne
(IQM) score
of 0.432, which 28.2% higher than SPR (IQM: 0.337) and 15.5% % higher than Play Virtual (IQM:
0.374). This indicates that MLR has the highest sample efficiency overall. For the optimality gap (OG)
metric, MLR reaches an OG of 0.522 which better than SPR (OG: 0.577) and Play Virtual (OG:
0,558), showing that our method performs closer to the desired target, i.e., human-level performance.
The full scores of MLR (over random seeds) across the 26 Atari games and more comparisons and analysis can be found in Appendix C.I.
7

---

IQM
Optimality Gap
MLR
PlayVirtual
SPR
DrQ(E)
DrQ
CURL
OTR
DER
0.2
0.3
0.4
0.6
0.7
0.8
Human Normalized Score
Figure 4: Comparison results on Atok.-100k. Aggregate metrics (IQM and optimality gap (OG) with 95% confidence intervals (CIs) are used for the evaluation. Higher IQM and lower OG are better.
4.3 Ablation Study
In this section, we investigate the effectiveness of our MLR auxiliary objective, the impact of masking strategies and the model/method design. Unless otherwise specified, we conduct the ablation studies on the DMControl- benchmark and run each model with 5 random seeds.
Effectiveness evaluation. Besides the comparison with SOTA methods, we demonstrate the effective- ness of our proposed MI. by studying its improvements compared to our Baseline. The numerical results are presented in [ while the curves of test performance during the training process are given in Appendix C.2. Both numerical results and test performance curves can demonstrate that our method obviously outperforms Baseline across different environments thanks to more informative representations 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-Ty:: we divide the input observation sequence into multiple segments along the temporal dimension and mask out a portion of segments randomly. (Here, the segment length masking (also referred
set to be equal to the temporal length of cube, i.e., k.) (iii) to as "cube masking"): as aforementioned and illustrated in Figure 2 we rasterize the observation sequence into non-overlapping cubes and randomly mask portion of them. Except for the differences described above, other configurations for masking remain the same as our proposed spatial-tempont (i.e., cube) masking. From the results in Table 2 we have the following observations: (i) All three masking strategies (i.e., MLR-S, MLR-T and MLR) achieve mean score improvements compared to Baseline by 18.5%, 12.2% and 25.0% % respectively, and achieve median by 23.4%, 25.0% and 35.9%. respectively. This demonstrates the effectiveness of the core idea of introducing mask-based reconstruction to improve the representation learning of RL. (ii) Spatial- temporal masking is the most effective strategy over these three design choices. This strategy matches better with the nature of video data due to its spatial-temporal continuity in masking. It encourages the state representations to be more predictive and consistent along the spatial and temporal dimensions, thus facilitating policy learning in RL.
improvements
Reconstruction target. In masked language/image modeling, reconstruction/prediiioo is commonly performed in the original signal space, such as word embeddings pixels. To study the influence of the reconstruction targets for the task of RL, we compare two different reconstruction spaces: (i) Pixel space reconstruction (denoted MLR-Pixel, we predict the masked contents directly by reconstruct- ing original pixels, like the practices in CV and NLP domains; (ii) Latent space reconstruction (i.e., MI.R): reconstruct the state representations (i.e., features) of original observations from masked observations, as we proposed in MIR. 2 the comparison results. The reconstruction in the latent space is superior that in the pixel space in improving sample efficiency. As discussed in the preceding sections, vision data might contain distractions and redundancies for policy learning, making the pixel-level reconstruction unnecessary. Besides, latent space reconstruction more conducive to the coordination between the representation learning and the policy learning in RL, because the state representations are directly optimized.
8

---

Table 2: Ablation study of masking strategy and reconstruction target. We compare three masking strategies: spatial masking (MLR-S), temporal masking (MLR-T) and spatial-temporal masking (MLR), and two reconstruction targets: original pixels (denoted MLR-Pixed and momentum projections in the latent space (i.e., MLR).
Environment
Baseline
MLR-S
MLR-T
MLR-Pixel
MLR
Finger, spin Cartpole, swingup
822 + 146
919 + 55
787 + 139
782 = 95
907 = 69
782 +: 74
665 = 118
829 + 33
803 = 91
791 = 50
Reacher, easy Cheetah, run
557 + 137
848 + 82
745 -
787 = 136
875 92
438 +
449 + 46
443 +
346 84
495 13
Walker, walk
414 310
556 + 189
393 + 202
490 216
597 102
Ball in cup, catch
669 + 310
927 + 6
934 + 29
675 + 292
939 + 9
Mean
613.7
727.3
688.5
647.2
767.3
Median
613.0
756.5
766.0
728.5
833.0
Mask ratio. In recent works of masked image
modeling 471 the mask ratio is found to be
844.0
840
crucial for the final performance. We study the
5
Medien
influences of different masking ratios
820
and find that the ratio of 50% is appropriate choice for our proposed MLR. An over-small value of this ratio could eliminate redun- dancy, making the objective easy to be reached by extrapolation from neighboring contents that are free of capturing and understanding seman- ties from visual signals. An over-large value leaves few contexts for achieving the reconstruc- tion goal. As discussed 21 477, the choice of this ratio varies for different modalities and depends on the information density of the input signals.
ple efficiency Figure 3 (with random seeds)
800
781
786 5
786.5
780
771.2
763
780
749.3
744
745,
740.7
747 0
740
123.
714
707.8
10
20
30
40
60
70
50
90
Figure 5: Ablation study of mask ratio.
Action token. We study the contributions of action tokens used for the latent decoder in Figure by discarding it from our proposed framework. The results are given in Table 3 Intuitively, prediction only from visual signals is of more less ambiguity. Exploiting action tokens benefits by reducing such ambiguity so that the gradients of less uncertainty can be obtained for updating the encoder.
Table 3: Ablation studies on action token, masking features and momentum decoder. MLR W.O. AcrTok denotes removing the action tokens in the input tokens of the predictive latent decoder. MLR-F R-F indicates performing masking on convolutuional feature maps. MLR-MoDec indicates adding momentum predictive latent decoder in the target networks.
Environment
Baseline
MLR w.o. ActTok
MIR-F
MLR-MoDec
MLR
Finger, spin Cartpole, swingup
822 + 146
832 +: 46
828 + 143
843 + 135
907 +: 69
782 + 74
816 + 27
789 55
766 + 88
791 50
Reacher, easy Cheetah, run
557
835
753 159
800 + 49
875
438 + 33
433 + 72
477 38
470 1 12
495 = 13
Walker, walk
414 + 310
412 - 210
673 - 33
571 + 152
597 + 102
Ball in cup, catch
669 + 310
837 + 114
843 + 119
788 + 155
939 - 9
Mean
613.7
694.2
727.2
706,4
767.3
Median
613.0
824.2
771.0
777.0
833.0
Masking features. We compare "masking pixels" and "masking features" in Table 3 Masking features (denoted by MIR-F) does not perform equally well compared with masking pixels as proposed in MLR, but it still achieves significant improvements relative to Baseline.

---

Why not use latent decoder for targets? We have tried to add coomentum-updateer latent decoder for obtaining the target states S1+- illustrated in Figure|] [T The results show that adding a momentum decoder leads to performance drops. Since there is no information prediction for the state representation learning from the original observation sequence without masking, we do not need to refine the implicitly predicted information like that in the outputs 8 of the online encoder.
Decoder depth. We analyze the influence of us- ing Transformer-based latent decoders with dif- ferent depths. We show the experimental results in Table 4 Generally, deeper latent decoders lead worse sample efficiency with lower mean and median scores. Similar to the designs in 21 471. it is appropriate to use lightweight decoder in MIR, because we expect the predict- ing masked information to be mainly completed by the encoder instead of the decoder. Note that the state representations inferred by the encoder are adopted for the policy learning in RL.
Table 4: Ablation study of predictive latent decoder depth. We report the number of parameters, mean and median scores on DMControl-100k:
Layers
Param.
Mean
Median
20.4K
726.8
719.0
2
40.8K
767.3
833.0
4
81.6K
766.2
789.5
8
163.2K
728.3
763.5
Similarity loss. We compare models using two kinds of similarity metries to measure the distance in the 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.2.
39 10] have shown that in self-supervised learning, supervising the feature representations in the projected space via the projection/predictinn heads is often better than in the original feature space. We investigate the effect of the two heads and find that both improve agent performance (see Appendix C.2).
Projection and prediction heads. Previous works 12
Sequence length and cube size. These two factors can be viewed as hyperparameters. Their corresponding experimental analysis and results are in Appendix C.2.
4.4 More Analysis
We make more detailed investigation and analysis of our MLR from the following aspects with details shown in Appendix C.3. (i) We provide more detailed analysis of the effect of each masking strategy. (ii) The quality of the learned representations are further evaluated through pretraining evaluation and regression accuracy test. (ii) We also investigate the performance of MLR on more challenging control tasks and find that MLR is still effective. (iv) We additionally discuss the relationship between Virtual and MLR. They both achieve leading performance in improving sample efficiency but from different perspectives. (v) We discuss the applications and limitations of MLR. We observe that MLR is more effective tasks with backgrounds and viewpoints that do not change drastically than on tasks with drastically changing backgrounds/viewpoints or vigorously moving objects.
5 Conclusion
In this work, we make the first effort to introduce the mask-based modeling to RL for facilitating policy learning by improving the learned state representations. We propose MI.,, a simple yet effective self-supervised auxiliary objective reconstruct the masked information in the latent space. In this way, the learned state representations are encouraged to include richer and more informative features. Extensive experiments demonstrate the effectiveness of MLR and show that MLR achieves the state-of-the-art performance on both DeepMind Control and Atari benchmarks. We conduct a detailed ablation study for our proposed designs in MLR and analyze their differences from that in NLP and CV domains. We hope our proposed method can further inspire research for vision- based RL from the perspective of improving representation learning. Moreover, the concept of masked latent reconstruction is also worthy of being explored and extended in CV and NLP fields. We are looking forward to seeing more mutual promotion between different research fields.
Acknowledgments
We thank all the anonymous reviewers for their valuable comments our paper.
10

---

References
[1] Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A. C., and Bellemare, M. Deep rein- forcement learning the edge of the statistical precipice. Advances in Neural Information Processing Systems, 34, 2021.
[2] Anand, A., Racah, E., Ozair, S., Bengio, Y., Côté, M.-A., and Hjelm, R. D. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, 2019.
[3] Ba, J. L.,, Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv: FrXiv:/607.064500 2016.
[4] Bao, H., Dong, L., Piao, S., and Wei, F. Beit: Bert pre-training of image transformers. In International Conference on Learning Representations, 2021.
[5] Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
[6] Bellman, R. A markovian decision process. Journal of mathematics and mechanics, 6(5): 684, 1957.
[7] Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., al. Language models are few-shot learners. arXiv preprint wrxiv:2005.14165. 2020.
[8] Castro, S., Moitra, S., Gelada, C., Kumar, S., and Bellemare, M. G. Dopamine: research framework for deep reinforcement learning. arXiv preprint arXiv:/812.00 1812.66/ 10, 2018.
[9] Chen, M., Radford, A., Child, R., Wu, J., Jun, H., Luan, D., and Sutskever, I. Generative pretraining from pixels. In International Conference on Machine Learning, pp. 1691-1703. PMLR, 2020.
[10] Chen, T. Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual representations. In International Conference Machine Learning, pp- 1597-1607. PMIR, 2020.
[11] Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04800 2018
[12] Grill, J.-B., Strub, F., Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., Doersch, C., Avila Pires, B., Guo, Z., Gheshlaghi Azar, M., Piot, B., kavukcuoglu, k., Munos, R., and Valko, M. Bootstrap your own latent new approach to self-supervised learning. In Advances in Neural Information Processing Systems, 2020.
[13] Guo, Z. D., Pires, B. A., Piot, B., Grill, J.-B., Altché, F., Munos, R., and Azar, M. G. Boot- strap latent-predictive representations multitask reinforcement learning. In International Conference Machine Learning, pp. 3875-3886. PMLR, 2020.
[14] Guo, Z. D., Thakoor, S., Pislar, M., Pires, B. A., Altché, F., Tallec, C., Saade, A., Calandriello, D., Grill, J.-B., Tang. Y., et al. Byol-explore: Exploration by bootstrapped prediction. arXiv preprint 2rXiv:2206.08332. 2022.
[15) Haarnoja, T., Zhou, A., Hartikainen, K., Tucker, G., Ha, S., Tan, J., Kumar, V., Zhu, H., Gupta, A., Abbeel, P., al. Soft actor-critic algorithms and applications. arXiv preprint 18Xiv:1812.0555 2018.
[16] Hafner, D., Lillicrap. T., Fischer, I., Villegas, R., Ha, D... Lee, H., and Davidson, J. Learning latent dynamics for planning from pixcls. In International Conference on Machine Learning. PMLR, 2019.
pp.
[17] Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
11

---

[18] Hafner, D., Lillicrap. T., Norouzi. M., and Ba, J. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193, 2020.
[19] Hansen, N. and Wang, X. Generalization in reinforcement learning by soft data augmentation. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pp. 1361 1-13617. IEEE, 2021.
[20] Hansen, S., Dabney, W., Barreto, A., Van de Wiele, T., Warde-Farley, D., and Mnih, V. Fast task inference with variational intrinsic successor features. arXiv preprint arXiv:/906. 1906005000 2019.
[21] He, K., Chen, X., Xie. S., Li, Y., Dollár, P., and Girshick, R. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/VV Conference on Computer Vision and Pattern Recognition, pp. (6000-16009, 2022.
[22] Hessel, M., Modayil, J., Van Hasselt, H.. Schaul, T., Ostrovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M., and Silver, D. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
[23] Jaderberg, M., Mnih, V., Czarnecki, W. M., Schaul, T., Leibo, J. Z., Silver, D., and Kavukcuoglu, K. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv: 161 1.05397, 2016.
[24] Kaelbling, L. P., Littman, M. L.,, and Cassandra, A. R. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-133 1998.
[25] Kielak, K. Do recent advancements model-based deep reinforcement learning really improve data efficiency? arXiv preprint arXiv:2003. 10181, 2020.
[26] Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., and Srinivas, A. Reinforcement learning with augmented data. Advances in Neural Information Processing Systems, 2020.
[27] Laskin, M., Srinivas, A., and Abbeel, P. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, pp. 5650- PMLR, 2020.
[28] Lee, A. X., Nagabandi, A., Abbeel, P., and Levine, Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. In Advances in Neural Information Processing Systems, 2020.
[29] Lee, K.-H., Fischer, I., Liu, A., Guo, Y., Lee, H., Canny, J., and Guadarrama, S. Predictive information accelerates learning in RI.. arXiv preprint arXiv:2007. /2401, 2020.
[30] Liu, G., Zhang, C., Zhao, L., Qin, T., Zhu, J., Jian, L., Yu, N., and Liu, T.-Y. Return-based contrastive representation learning for reinforcement learning. In International Conference on Learning Representationss 2021.
[31] Liu, H. and Abbeel, P. Aps: Active pretraining with successor features. In International Conference on Machine Learning, pp. 6736-6747. PMLR, 2021.
[32] Liu, H. and Abbeel, P. Behavior from the void: Unsupervised active pre-training. Advances in Neural Information Processing Systems, 14:18499-18777 2021.
[33] Mazoure, B., Tachet des Combes, R., DOAN, T. L., Bachman, P., and Hjelm, R. D. Deep reinforcement and infomax learning. Advances in Neural Information Processing Systems, 2020.
[34] Mnih, V., Kavukcuoglu, K., Silver, D.,, Graves, A., Antonoglou, I., Wierstra, D., and Riedmiller, M. Playing atari with deep reinforcement learning, arXiv preprint 1rXiv:1312.5602, 2013.
[35] Oord, A. v. d., Li, Y., and Vinyals, O. Representation learning with contrastive predictive coding. arXiv preprint arXiv:
2018.
[36] Pathak, D., Krahenbuhl, P., Donahue, J., Darrell, T., and Efros, A. A. Context encoders: Feature learning by inpainting. In Proceedings of the IEFE conference on computer vision and pattern recognition, pp. 2344,,544, 2016.
12

---

[37] Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I. Improving language understanding by generative pre-training, 2018.
[38] Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[39] Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A., and Bachman, P. Data-efficient reinforcement learning with self- ooff-predietive representations. In International Conference Learning Representations, 2021.
[40] Schwarzer, M. Rajkumar, N., Noukhovitch, M., Anand, A., Charlin, L., Hjelm, R. D., Bachman, P., and Courville, A. Pretraining representations for data-efficient reinforcement learning. Advances in Neural Information Processing Systems, 34: 14:11266-16699 2021.
[41] Shelhamer, E., Mahmoudieh, P., Argus, M., and Darrell, T. Loss is its own reward: Self- supervision for reinforcement learning. ArXiv, abs/1612.07307, 2017.
[42] Stooke, A., Lee, K., Abbeel, P., and Laskin, M. Decoupling representation learning from reinforcement learning. arXiv preprint arXiv:2009.083/9 2020.
[43] Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A.. Merel, J., Lefrancq, A., et al. Deepmind control suite. arXiv preprint 1rXiv:1801.000 006900 2018.
[44] van Hasselt, H. P., Hessel, M., and Aslanides, J. When to use parametric models in reinforcement learning? In Advances Neural Information Processing Systems, 2019.
[45] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.., Gomez, A. N., Kaiser, E.., and Polosukhin, I. Attention is all you need. In Advances in neural information processing systems, pp. 5998-99908 2017.
[46] Wei, C., Fan, H., Xie, S., Wu, C.-Y., Yuille, A., and Feichtenhofer, C. Masked feature prediction for self-supervised visual pre-training. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
2022.
[47] Xie, 7.,, Zhang, Z., Cao, Y., Lin, Y., Bao, J., Yao, Z.,, Dai, Q., and Hu, Simmim: simple framework for masked image modeling. In P Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 9653 133-99633,2 2022.
[48] Yarats, D., Fergus, R., Lazaric, A., and Pinto, L. Reinforcement learning with prototypical representations. In International Conference on Machine Learning, pp. 11920-11931. PMLR, 2021.
[49] Yarats, D., Kostrikov, I., and Fergus, R. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. International Conference on Learning Representations, 2021.
[50] Yarats, D., Zhang, A., Kostrikov, I., Amos, B., Pincau, J., and Fergus, R. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10674-10681, 2021.
[51] Ye, W., Liu, S., Kurutach, T., Abbeel, P., and Gao, Y. Mastering atari games with limited data. Advances Neural Information Processing Systems, 34, 2021.
[52] Yu, T., Lan, C., Zeng, W., Feng, M., Zhang, Z., and Chen, Z. Playvirtual: Augmenting cycle- consistent virtual trajectories for reinforcement learning. In Advances in Neural Information Processing Systems, 2021.
[53] Zhang, A., McAllister, R. T., Calandra, R., Gal. Y., and Levine, S. Learning invariant repre- sentations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
[54] Zhu, J., Xia, Y., Wu, I.,, Deng. J., Zhou, W., Qin, T., and Li, H. Masked contrastive representation learning for reinforcement learning, arXiv preprint arXi::2010.07470, 2020.
13

---

[55] Lukasz Kaiser, Babaeizadeh, M., Mitos, P., Osiński, B., Campbell, R. H., Czechowski, K., Erhan, D., Finn, C., Kozakowski, P., Levine, S., Mohiuddin. A., Sepassi, R., Tucker, G., and Michalewski, H. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020.
14

---

Checklist
1. 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.3.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See Appendix D.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. 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]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material a URL)? [Yes] https://github.comm/
(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 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., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix B.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you eite 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 the supplemental material or as URL? [Yes]
Lated-LLatent- Witeet-Reet Reconstruction
(d) Did you discuss whether and how consent obtained from people whose data you're using/curating? [Yes] See Appendix B.
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [Yes] See Appendix
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
15


Reviews:
[{"review_id": "kiAL7mL44NV", "reviewer": "Reviewer_uZPQ", "summary": "This paper built upon recent advancement in NLP and CV to introduce an auxiliary task in RL. More specifically, the authors proposed to reconstruct latents from masked observations, with the goal of improving representations used in RL. A few experiments across both discrete and continuous control to demonstrate the promise of the proposed methods.", "questions": "1. The authors stated in L122 that MLR ``...is the theoretically applicable to different RL algorithms''. Is there any theory behind the proposed algorithm for the application in RL?\n\n2. Do the features used for RL come from the output of online encoder, i.e., \\tilde{s}?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n1. Using the idea of auxiliary tasks to boost the RL performance has been explored and this paper belongs to this line of research. The idea from the proposed method looks reasonable and interesting.\n\n2. The authors designed different types of experiments to justify the promise of the proposed method, and the results generally look promising.\n\n3. This paper is mostly well written and organized.\n\n\nWeakness:\n1. One issue is that the current version lacks the explanation and investigation of why the proposed method should outperform the previous algorithms combining representation learning with RL, e.g., CURL. We can see some promise based on experimental results, but in-depth exploration on why MLR works is integral to make the claims convincing.\n\n2. The numbers and plots in the experiments section definitely show that MLR achieves competitive performance. From Table 1, however, it seems that MLR has very similar performance when compared against PlayVirtual and CURL with 500k steps, which may suggest that the representations from MLR may not be sufficiently distinguishable from previous work. \n\n3. My another concern regarding the experiments corresponds to Atari games. In Table 1 of the appendix, the authors showed that MLR can outperform Baseline, which denotes Rainbow. The issue is that the scores here are much lower when compared with those in the Rainbow paper. Just pick a simple one: For the game Pong, the reported score here is -3.1 while it's 19 in the Rainbow paper. Consequently, it's not convincing to me that the auxiliary task can help the RL task on Atari games.\n\n4. A minor point regarding the background in Section 3.1: POMDP distinguishes state and observation in its formulation, but the text there lacks the discussion on states.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "P1k4k_iZ0dE", "reviewer": "Reviewer_djAT", "summary": "The paper introduces mask based reconstruction (MLR) for RL with image inputs. MLR is an auxiliary objective optimized together with the policy learning objective. The idea is to reconstruct masked pixels in the latent state instead of the input space where state representations are inferred from original unmasked frames. The paper proposes to mask portions on both spatial and temporal dimensions in the observation sequence. There is an “online encoder” for masked observations and a momentum-encoder for original images. Furthermore, there is a Transformer-based latent decoder that also takes action embeddings as input. MLR uses a projection + prediction head to get the final prediction result, and cosine similarity to compute the distance between masked and original sequence. The paper presents several ablations about the masking strategies, depth of decoder, similarity losses, etc. Results on DMC-100K and Atari-100K show that MLR outperforms leading algorithms like SPR, PlayVirtual.  \n", "questions": "1. The method was not compared with EfficientZero which is an important baseline. How does MLR compare with EfficientZero?\n2. The Appendix talks about the training time and details. How they compare with leading methods like SPR? \n3. There are no discussions on scenarios where the mask-based reconstruction is helping. For example, are improvements on games like MsPacman where the background is not moving or games with moving viewpoints like Crazy Climber? This could also help understand the limitations of the proposed method. ", "limitations": "The paper does not highlight the limitations of MLR. The paper should talk about potential bottlenecks when scaling MLR to more complex tasks. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- Paper is well written and easy to follow.\n- Results and ablation studies are thorough and demonstrates the efficacy of each component.\n- Interesting direction to apply mask-based reconstructions in latent space as it is challenging to learn good latent representations while solving RL tasks.\n\n### Weaknesses\n- Novelty is limited as Masked based reconstruction are already known to work well in Vision tasks.\n- Limitations of the proposed method are not discussed.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "BrRw_0SuJH", "reviewer": "Reviewer_du6d", "summary": "The paper tackles the field of state representation learning in visual RL. They propose Mask-based Latent Reconstruction (MLR) for learning state representations based on masking sequences of input images. This practice is inspired by the computer vision and natural language processing communities. The approach is different than the ones used traditionally in these fields by reconstructing the missing content in the latent space instead of the original one. The authors provide an extensive evaluation to demonstrate the improvement in performance their method achieves to SOTA methods. Their evaluation includes discrete and continuous control environments as well as ablations of the different design decisions introduced in their approach.", "questions": "- If I understood correctly, \"baseline\" in your tables refers to either SAC or rainbow. In table 1, that should be SAC. Can you explain how SAC is outperforming other methods which add on top of it such as SAC-AE and CURL? these methods were built with the motivation that SAC on pixels isn't good enough. Can you explain how these results relate to the original ones from those papers?\n- this brings me to the next question: how do you ensure a fair comparison in your evaluation? do all the methods use the same SAC implementation, same encoder, decoder architectures, and hyperparameters, etc...? Please elaborate on the standardization of experiments.\n- Related to what I mentioned in the weakness section, can you elaborate on the intuition behind your choice of targets?  to me, this seems like a random choice. In other words, what are you trying to enforce on the state representation? what kind of property? and how does your method allow for this property to emerge?  For instance, when using autoencoders, the property is compressing, with contrastive methods, it's somehow clustering... Please give a clear explanation here and not a high-level explanation.\n- Related to the previous comment, if MRL is not integrated with policy learning and used instead for pretaining only, would it actually help? I know that one of the motivations of the approach is that the two objectives are simultaneously and jointly trained, but such a test would be important to validate the importance of the learned state representation. Can you include such an experiment in the final version? for the rebuttal it doesn't have to be over all the environments...(whatever can be managed).", "limitations": "The paper lacks an evaluation of the obtained state representations. An evaluation could be something like the regression accuracy (or a correlation of the two variables) of the original states (provided by simulation) based on the representation. The pertaining experiment I mentioned in my last question would also qualify as a (qualitative) evaluation of the obtained representations. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n========\n\n- The idea of using masked observations for state representation learning is to my knowledge novel in RL\n- The paper is well situated with respect to previous related work\n- The paper is mostly (with few exceptions) well-written and easy to follow (clear)\n- The proposed method is simple to implement\n- The evaluation is extensive: it includes multiple baselines, different kinds of environments (discrete and continuous action spaces), and ablations of the different components of the method\n- The results show that the proposed method mostly outperforms the baselines and also leads to a low reward variance. \n\nWeaknesses:\n==========\n\n- In section 3.2, paragraph iii) (seconding) should be rewritten and made clearer. Multiple details in this paragraph are ambiguous and hard to follow. The first sentence in this paragraph is itself very confusing and could be interpreted in many ways (what is meant by \"visible ones\" and \"implicit way\"). The motivation for using the actions is missing or at least not convincing. Understanding relative positional embedding requires reading [43], which makes the paper not self-contained. Additionally, it's unclear what the plus in equation 2 means, is it an actual addition or a concatenation? it can't really be an addition since the two vectors have different dimensions, so is it a concatenation or something else? please specify.\n- The proposed approach claims to reconstruct the masked content in the latent space. This is done using the cosine loss based on a target that is obtained using the original images and a moving average of the online encoder and projection head. It is unclear to me what is the meaning of such a loss, why would we want to push the representation to be similar to another one produced by a moving average over its previous weights but using the original images? I would have been more convinced if the method simply relied on reconstructing the original pixels as done in the ablation. Let me rephrase, I'm not saying that the method doesn't make sense, I'm just saying that the intuition behind the state representation learning method is unclear (at least to me). I'm willing to change my rating if this aspect is made clearer and of course convincing.\n\nMinor Problems:\n============\n\n- the paper repeatedly mentions that the presented approach reconstructs the masked content in the latent space, which is a false statement. The method attempts to reconstruct the information available in the masked content in the latent space.\n- some typos and grammar problems here and there. I would recommend a non-author to go over the paper + some grammar check software. e.g. in line 262 neural --> natural. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "2tTCUVoTAXU", "reviewer": "Reviewer_ZSQ9", "summary": "This paper proposes a mask-based latent reconstruction/prediction auxiliary task for visual RL. The proposed auxiliary task takes as input a sequence of observations, applies a uniformly sampled spatiotemporal mask to the sequence, and then task the network with predicting the latent representation of the unmasked sequence encoded by an EMA of the online encoder. The proposed method, MLR, is evaluated on DMControl and Atari tasks, and a number of ablation studies uncover the key design choices of the proposed framework.", "questions": "I believe I have already made my biggest concerns clear in the \"weaknesses\" section above. However, I do have a few questions for the authors that should be addressable during the rebuttal:\n- It is not clear to me after reading the paper + appendix exactly which data augmentation is used and where. Can you elaborate on this?\n- I'd also like the authors to clarify on how the \"Baseline\" SAC implementation differs from that of other published works. Appendix describes the architecture to be based on CURL; would \"Baseline\" be equivalent to CURL w/o the auxiliary loss? Is data augmentation applied?\n- How do the listed wall-times for the proposed method compare to (1) the \"Baseline\" implementation without the auxiliary loss, and (2) prior work such as CURL, DrQ, SPR, or PlayVirtual?", "limitations": "I believe the authors provide a satisfactory discussion of potential negative societal impact of their work, but discussion of limitations is minimal. I'd like to see more explicit discussion of when the proposed auxiliary loss can be expected to be useful, and when it may fail (e.g., which key assumptions are made).", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Conceptually simple, clearly presented, and ablations are insightful.\n- Experimental results are encouraging, the provided rliable metrics are appreciated, and code is provided.\n\nWeaknesses:\n- DMControl evaluation. I understand that the particular choice of DMControl tasks used in the evaluation is convenient due to the number of available baseline results. However, performance on these tasks have largely been saturated by prior work and are thus not a good indicator of algorithmic improvements (e.g., does a median performance of 914 vs 929.5 vs 935 vs 907 vs 948 @ 500k steps really matter? I understand that 100k steps may be more informative than 500k but point still stands, and the improvement of curves for these tasks in Figure 2 of appendix don't appear significant to me, presumably for the same reason). I strongly suggest the authors evaluate on more challenging tasks / benchmarks in future revisions, although I realize that this may not be feasible within the strict time frame of a rebuttal.\n- Related work. I would like to see more discussion of RL methods that use BYOL-style auxiliary objectives and data augmentation/masking in the related work section, given the strong similarity to such methods. I will be happy to provide references, but am under the impression that the authors are familiar with the literature given some of these works (e.g. SPR) are cited and/or briefly mentioned in the experiment section.\n\n**Post-rebuttal:** In light of the author response I am willing to raise my score by 1. I expect the authors to include reviewer suggestions in a future revision.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept