Masked World Models for Visual Control
Authors: Younggyo Seo
Danijar Hafner
Hao Liu
Fangchen Liu
Stephen James
Kimin Lee
Pieter Abbeel
Venue: CoRL 2022 Poster (2022)
Keywords: Visual model-based RL
World models

Abstract: Visual model-based reinforcement learning (RL) has the potential to enable sample-efficient robot learning from visual observations. Yet the current approaches typically train a single model end-to-end for learning both visual representations and dynamics, making it difficult to accurately model the interaction between robots and small objects. In this work, we introduce a visual model-based RL framework that decouples visual representation learning and dynamics learning. Specifically, we train an autoencoder with convolutional layers and vision transformers (ViT) to reconstruct pixels given masked convolutional features, and learn a latent dynamics model that operates on the representations from the autoencoder. Moreover, to encode task-relevant information, we introduce an auxiliary reward prediction objective for the autoencoder. We continually update both autoencoder and dynamics model using online samples collected from environment interaction. We demonstrate that our decoupling approach achieves state-of-the-art performance on a variety of visual robotic tasks from Meta-world and RLBench, e.g., we achieve 81.7% success rate on 50 visual robotic manipulation tasks from Meta-world, while the baseline achieves 67.9%. Code is available on the project website: https://sites.google.com/view/mwm-rl.

Masked World Models for Visual Control
Younggyo Seo Danijar Hafner2.3.4
Hao Liu
Fangchen Liu²
Stephen James
Kimin Lee
Pieter Abbeel
KAIST 2 UC Berkeley
Gooole Research
University of Toronto
Abstract: Visual model-based reinforcement learning (RL) has the potential enable sumple-efficien robot learning from visual observations. Yet the current approaches typically train a single model end-to-end for learning both visual rep- resentations and dynamics, making it difficult to accurately model the interaction between robots and small objects. In this work, we introduce a visual model-based RL framework that decouples visual representation learning and dynamics learn- ing. Specifically, we train an autoencoder with convolutional layers and vision transformers (ViT) to reconstruct pixels given masked convolutional features, and learn a latent dynamics model that operates on representations from the autoen- coder. Moreover, to encode task-relevant information, we introduce an auxiliary reward prediction objective for the autoencoder. We continually update both au- toencoder and dynamics model using online samples collected from environment interaction. We demonstrate that our decoupling approach achieves state-of-the-arr performance variety of visual robotic tasks from Meta- world and RLBench, egg., we achieve 81.7% success rate on 50 visual robotic manipulation tasks from Meta-worde world, while the baseline achieves 67.9 9... Code is available the project website: https: /ttt///stttseeeeeeeeeeeeeeeeeeee wwwiiee..goggeeee..ooeoooo com/ wwvew//wee wwwee/w/wmmmm... ht///mwmll
Introduction
Model-based reinforcement learning (RL) holds the promise of sample-efficient robot learning by learning a world model and leveraging it for planning [1, 2, 3] or generating imaginary states for behavior learning [4, 5J. These approaches have also previously been applied to environments with visual observations, by learning an action-conditional video prediction model [6, 7] or latent dynamics model that predicts compact representations in an abstract latent space [8, 9]. However, learning world models environments with complex visual observations, e.g., accurately modeling interactions with small objects, is an open challenge.
We argue that this difficulty comes from the design of current approaches that typically optimize the world model end-to-end for learning both visual representations and dynamics [9, 10]. This imposes trade-off between learning representations and dynamics that can prevent world models from accurately capturing visual details, making it difficult to predict forward into the future. Another approach is to learn representations and dynamics separately, such as earlier work by Ha and Schmidhuber [11] who train variational autoencoder (VAE) [12] and dynamics model on top of the VAE features. However, separately-trainet VAE representations may not be amenable to dynamics learning [8, 10] or may not capture task-relevant details [11].
On the other hand, masked autoencoders (MAE) [13] have recently been proposed as an effecttve and scalable approach visual representation learning. by training self-supervised vision trans- former (ViT) [14] to reconstruct masked patches. While it motivates us to learn world models on top of MAE representations, we find that MAE often struggles capture fine-grained details within patches. Because capturing visual details, e.g., object positions, is crucial for solving visual control tasks, is desirable to develop a representation learning method that captures such details but also achieves the benefits of MAE such as stability, compute-efficiency. and scalability.
Work done while visiting UC Berkeley. Correspondence to younggyo.. seokkaist. ac, kr.
I Now at Dyson Robot Learning Lab.
6th Conference on Robot Learning (CoRI. 2022), Auckland, New Zealand.

---

Visual Representation Leaming
Dynamics Learning
2
Ōc
St
Sr+1
Masked Conv Features
C,0
G.0
2++1
Figure 1: Illustaation of our approach. We continually update visual representations and dynamics using online samples collected from environment interaction, by repeating iterative processes of training (Left) an autoencoder with convolutional feature masking and reward prediction and (Right) a latent dynamics model the latent space of the autoencoder. We note that autoencoder parameters are not updated during dynamics learning.
In this paper, we present Masked World Models (MWM), visual model-based RL algorithm that decouples visual representation learning and dynamies learning. The key idea of MWM is train an autoencoder that reconstructs visual observations with convolutional feature masking, and latent dynamics model on top of the autoencoder. By introducing early convolutional layers and masking out convolutional features instead of pixel patches, our approach enables the world model to capture fine-gan fine-grained visual details from complex visual observations. Moreover, in order to learn task- relevant information that might not be captured solely by the reconstruction objective, an auxiliary reward prediction task for the autoencoder. Specifically, we separately update visual representations and dynamics by repeating the iterative processes of (i) training the autoencoder with convolutional feature masking and reward prediction, and (ii) learning the latent dynamics model that predicts visual representations from the autoencoder (see Figure 1).
introduce
Contributions
We highlight the contributions of our paper below:
We demonstrate the effectiveness of decoupling visual representation learning and dynamics learning for visual model- RL. MWM significantly outperforms state-of-the-att model- based baseline [15] on various visual control tasks from Meta-world [16] and RLBench [17].
We show that seff-supervised VIT trained reconstruct visual observations with convolutional feature masking can be effective for visual model-based RI. Interestingly, we find that masking convolutional features can be more effective than pixel patch masking [13], by allowing for capturing fine-grained details within patches, This is in contrast to the observation in Touvron al. [18], where both perform similarly on the ImageNet classification task [19].
We show that an auxiliary reward prediction task can significantly improve performance by encoding relevatlevant information into visual representations.
Related Work
World models from visual observations There have been several approaches to learn visual
representations for model-based approaches via image reconstruction [6, 7, 8, 9, 10, 11, 15, 20, 21, 22], e.g., learning a video prediction model [6, 23] or latent dynamics model [8, 9, 10]. This has been followed by series of works that demonstrated the effectiveness of model-base approaches for solving video games [15, 24, 22] and visual robot control tasks [7, 21, 25, 26]. There also have been several works that considered different objectives, including bisimulation [27] and contrastive learning [28, 29, 30). While most prior works optimize a single model to learn both visual representations and dynamics, we instead develop framework that decouples visual representation learning and dynamics learning.
Self-supervised vision transformers Self-supervised learning with vision transformers (ViT) [14] has been actively studied. For instance, Chen et al. [31] introduced MoCo-v3 which trains ViT with

---

(a) Pick Place
(b) Shelf Place
(c) Reach Target
(d) Push Button
(e) Reach Duplo
Figure 2: Examples of visual observations used in our experiments. We consider a variety of visual robot control tasks from Meta-world [16], RLBench [17], and DeepMind Control Suite [41].
contrastive learning. Caron et al. [32] introduced DINO which utilizes self-distillation loss [33], and demonstrated that self-supervised ViTs contain information about the semantic layout of images. Training sulf--up VITs with masked image modeling [13, 34, 35, 36, 37, 38, 39] has also been successful. In particular, He et al. [13] proposed masked autoencoder (MAE) that reconstructs masked pixel patches with asymmetric encoder-decoder architecture. Unlike MAE, propose to randomly mask features from early convolutional layers [40] instead pixel patches and demonstrate that supervidd ViT' can also be effective for visual model-based RL.
We provide more discussion on related works in more detail in Appendix C.
3 Preliminaries
Problem formulation We formulate a visual control task as partially observable Markov decision process (POMDP) [42], which is defined is the action space, (orl a<te) is the transition dynamics, is the reward function that maps previous observations and actions to a reward r, (oga. ace), and y E [0, is the discount factor.
tuple (O, r, 7). O is the observation space, A
Dreamer Dreamer [15, 21] is a visual model-based RL method that learns world models from pixels and trains actor-critic model via latent imagination. Specifically, Dreamer learns Recurrent State Space Model (RSSM) [9]. which consists of following four components:
Representation model: 8: go( 1 4t-1,4-1,00) Image decoder: 6, ~ 00 (x))
(1)
Transition model: 8, (s. 8,-1,
Reward predictor: f. PP s,)
The representation model extracts model state s, from previous model state 8,-1, previous action 4- and current observation or. The transition model predicts future state š, without the access to current observation o,- The image decoder reconstructs raw pixels provide learning signal, and the reward predictor enables us to compute rewards from future model states without decoding future frames. All model parameters 0 are trained jointly learn visual representations and environment dynamics by minimizing the negative variational lower bound [12]:
[())= = (o1
ΣT (- (nlooo(lsss)) - Inpo(re(s:) + KL (q0(s(/1t-1,,
where is hyperparameter that controls the tradcoff between the quality of visual representation learning and the accuracy of dynamics learning [43]. Then, the critic is learned to regress the values computed from imaginary rollouts, and the actor is trained to maximize the values by propagating analytic gradients back through the transition model (see Appendix A for the details).
Masked autoencoder Masked autoencoder (MAB) [13] is self-supervised visual representation technique that trains an autoencoder to reconstruct raw pixels with randomly masked patches con- sisting of pixels. Following a scheme introduced in vision transformer (ViT) [14], the observation Or E XX is processed with patchify stem that reshapes or into a sequence of 2D patches
h E (PPCC) where P is the patch size and N HW/P p2 is the number of patches. Then subset of patches is randomly masked with ratio of to construct hT" MM
Patchify stem: h4 = (oa)c
Masking:
hT" ~ pask [4, m)
(3)

---

A ViT encoder embeds only the remaining patches h²" into D-dimensional vectors, concatenates the embedded tokens with learnable CLS token, and processes them through a series of Transformer layers [44]. Finally, ViT decoder reconstructs the observation by processing tokens from the encoder
and learnable mask tokens through Transformer layers followed by a linear output head:
(4) All the components paramaterized by ø are jointly optimized to minimize the mean squared error (MSE) between the reconstructed and original pixel patches. MAE computes & without masking, and utilizes its first component (i.e., CLS representation) for downstream tasks (e.g., image classification).
ViT encoder: 2t ~ (2 ht(") ViT decoder: Or ~ Po 2²"")
4 Masked World Models
In this section, we present Masked World Models (MWM), a visual model-based RL framework for learning accurate world models by separately learning visual representations and environment dynamics. Our method repeats (i) updating an autoencoder with convolutional feature masking and an auxiliary reward prediction task (see Section 4.1), (ii) learning dynamics model in the latent space of the autoencoder (see Section 4.2), and (iii) collecting samples from environment interaction. We provide the overview and pseudocode of MWM in Figure 1 and Appendix D, respectively.
4.1 Visual Representation Learning
It has been observed that masked image modeling with a ViT architecture [13, 34, 36] enables compute-efficiet and stable supe-supe visual representation learning. This motivates us to adopt this approach for visual model-based RL, but we find that masked image modeling with commonly used pixel patch masking [13] often makes it difficult to learn fine-grained details within patches, e.g., small objects (see Appendix B for a motivating example). While one consider small-size patches, this would increase computational costs due to the quadratic complexity self-attention layers.
To handle this issue, we instead propose to train an autoencoder that reconstructs raw pixels given randomly masked convolutional features. Unlike previous approaches that utilize patchify stem and randomly mask pixel patches (see Section 3), we adopt a convolution stem [14, 40] that processes of through series of convolutional layers followed by flatten layer, to obtain hg E TR. No where N₂ is the number of convolutional features. Then by is randomly masked with a ratio of m to obtain h, c.c.m E R M. D. and ViT encoder and decoder process h; to reconstruct raw pixels.
Convolution stem: hf feony (or)
Masking:
h ~ ~p p (h, (h.rm) h,, m)
(5)
VIT encoder:
2t c. po(2 1) ViT decoder: 6 6 24 m))
Because early convolutional layers mix low-level details, we find that our autoencoder can effectively reconstruct all the details within patches by learning to extract information from nearby non-masked features (see Figure for examples). This enables us to learn visual representations capturing such details while also achieving the benefits of MAE, e.g., stability and compute-efficiency
Reward prediction In order to encode task-relevant information that might not be captured solely by the reconstruction objective, we introduce an auxiliary objective for the autoencoder to predict rewards jointly with pixels. Specifically, we make the autoencoder predict the reward r, from zf in conjunction with raw pixels.
(6) In practice, we concatenate one additional learnable mask token inputs of the ViT decoder, and utilize the corresponding output representation for predicting the reward with a linear output head.
ViT decoder with reward prediction: o,, fg ~ Po Fr "()"")
High masking ratio Introducing carly convolutional layers might impede the masked reconstruc- tion tasks because they propagate information across patches [18], and the model can exploit this to find a shortcut solve reconstruction tasks. However, we find that high masking ratio (i.e., 75%) can prevent the model from finding such shortcuts and induce useful representations (see Fig- ure 6(b) for supporting experimental results). This also aligns with the observation from Touvron et al. [18), where masked image modeling [34] with a convolution stem [45] can achieve competitive performance with the patchify stem on the ImageNet classification task [19].
We also note that Xiao et al. [40] showed that introducing early convolutional layers can stabilize the training of ViT models, which can be helpful for RL that requires stability from the beginning of the training.

---

Place
102
Dreamer/2
DreamerV2
100
DreamerV2 MWW
Noaly Mwl.
n Muuptt WWW
Success (55)
Success Rate (%5)
Succets Rato (%)
92
40-10" meito
100-10t
6x10f
10-10'
122+10 180+10"
260-10t
Ox10"
40-10"
120 100-10" 200-10t
Environment Steps
Environment Sleps
Environment Steps
Push
Shelf Place
Sweep
LMM
DreamerV2
MWM
MWM
MWMM
Sucozees Rade (8)
Suroverss Rale (8)
Succiss Rate (N)
MY MwM
0-32" 40-10"
20-12"
120-10" 192-10" 200-10"
0-10"
60"12"
123-10"
190-10" 240 300-10"
0x10"
20-10"
40x10"
00-10"
30-10" 100-10"
Enviromment Stopa
Environment Stops
Environment Staps
Figure 3: Learning curves on six visual robotic manipulation tasks from Meta-world as measured on the success rate. We select the tasks that require modeling interactions between small objects and robot arms. Learning curves on tasks are available in Appendix G. The solid line and shaded regions represent the mean and bootstrap confidence intervals, respectively, across five runs.
4.2 Latent Dynamics Learning
Once we learn visual representations, we leverage them for efficiently learning dynamics model in the latent space of the autoencoder. Specifically, we obtain the frozen representations 2..0 from the autoencoder, and then train variant of RSSM whose inputs and reconstruction targets are 2 by replacing the representation model and the image decoder in Equation with following components:
Representation model:
~ [o( 1, 0,,,, c,0
Visual representation decoder: a (2 |s))
Because visual representations capture both high- and low-level information in an abstract form, the model can focus more on dynamics learning by reconstructing them instead of raw pixels (see Section 5.5 for relevant discussion). Here, we also note that we utilize all the elements of 2. unlike MAE that only utilizes CLS representation for downstream tasks. We empirically find this enables the model to receive rich learning signals from reconstructing all the representations containing spatial information (see Appendix for supporting experiments).
Optimization
Given a random batch {(0,, ry, (;) a;)] MWM objective is defined as follows: B
2IN (0,
0)
Inpp(ool=,""") - 11n-1-111 (r
(8)
B
3=1
visual representation leamning
-Inpo 1s;;) In + KL qo (s(5))s)-1.. |
dynamics learning
Experiments
We evaluate MWM various robotics benchmarks, including Meta-world [16] (see Section 5.1), RLBench [17] (see Section 5.2), and DeepMind Control Suite [46] (see Section 5.3). We remark that these benchmarks consist of diverse and challenging visual robotic tasks. We also analyze algorithmic design choices in-depth (see Section 5.4) and provide a qualitative analysis of how our decoupling approach works by visualizing the predictions from the latent dynamics model (see Section 5.5).
Implementation We use visual observations of 64 64 3. For the convolution stem, we stack 3 convolution layers with the kernel size of 4 and stride 2, followed by linear projection layer.
5

---

1/0
Dreamar/22
1.0
Dreame V2
1.0
Dieamar12
MMM
Success Rate (50
(S3)
0.0
mormon W Wy
Suceess Rato (8)
0.6
0.8
04
0.4
0.4
0.
Environeent
Environment
(a) Meta- world aggregated
RLBench: Reach Target
(c) RLBench: Push Button
Figure 4: (a) Aggregate performance on all 50 Meta-world tasks. We normalize environment steps by maximum steps in each task. The solid line and shaded regions represent the mean and stratified bootstrap confidence intervals, respectively. across 250 runs. We report the learning curves on (b) Reach Target and (c) Push Button from RLBench. Performances are not directly comparable to previous results [47, 48] due to the difference in setups (see Section 5.2). The solid line and shaded regions represent the mean and bootstrap confidence intervals, respectively, across eight runs.
We use 4-layer ViT encoder and 3-layer VIT decoder. We find that initializing the autoencoder with a warm-up schedule at the beginning of training is helpful. Unlike MAE, we compute the loss on entire pixels because we do not apply masking implementation on top DreamerV2 [15]. To take a sequence of autoencoder representations as inputs, replace CNN encoder and decoder with 2-layer Transformer encoder and decoder. We use same hyperparameters within the same benchmark. More details are available in Appendix E.
pixels. For world models, we build our
5.1 Meta-world Experiments
Environment details In order to use single camera viewpoint consistently over all 50 tasks, use the modified corner2 camera viewpoint for tasks. our experiments, we classify 50 tasks into easy, medium, hard, and very hard tasks where experiments environments steps with action repeat of 2, respectively. More details are available in Appendix
over 500K, 1M, 2M, 3M
of selected six challenging tasks that require agents to control robot arms to interact with small objects. We find that MWM significantly outperforms DreamerV2 in terms of both sample-efficiency and final performance. In particular, MWM achieves > 80% success rate on Pick Place while Dreamer V2 struggles to solve the task. These results show that our approach of separating visual representation learning and dynamics learning can learn accurate world models on challenging domains. We also note that DreamerV2 learns visual representations with the reward prediction loss, which implies that our superior performance does not come from utilizing additional information. Figure 4(a) shows the aggregate performance over all the 50 tasks from the benchmark, demonstrating that our method consistently outperforms DreamerV2 overall. We also provide learning curves on all individual tasks in Appendix G, where MWM consistently achieves similar or better performance on most tasks.
Results In Figure 3, report the performance
5.2 RLBench Experiments
Environment details In order to evaluate our method on more challenging visual robotic manipula- tion tasks, we consider RLBench [17], which has previously acted as effective proxy for real-robot performance [48]. Since RLBench consists of reward and challenging tasks, solving them typically requires expert demonstrations, specialized network architectures, additional inputs (c.g.. point cloud and proprioceptive states), and an action mode that requires path planning [47, 48, 49, 50]. While we could utilize some of these components, we instead leave this as future work in order to maintain a consistent evaluation setup across multiple domains. In our experiments, we instead consider two relatively easy tasks with dense rewards, and utilize an action mode that specifies the delta of joint positions. We provide more details in Appendix F.
Results As shown Figure 4(b) and Figure 4(c), we observe that our approach can also be effective RLBench tasks, significantly outperforming DreamerV2. In particular, DreamerV2

---

Duplo
Quadrupad Run
Quadripad Walk
292
Dreamer/2
1003
DreamerV2
1920
202
803
820
Epleced Returns
Eploade Retums
Epilocao Riduma
190
800
820
000
600
600
200
200-
DreamerV2
MMMM
0410
10+10"
20-13"
30ri0r 60-10"
Dx10"
20-100
40<10"
60-19" 80+10"
100
0+10"
20warr
40-10"
60-10"
80-10
100m10"
Enntronment Steps
Environment Stops
Environment Stops
Figure 5: Learning curves three visual robot control tasks from DeepMind Control Suite as measured on the episode return. The solid line and shaded regions represent the mean and bootstrap confidence intervals, respectively, across eight runs.
DreamenN2
75%
DreamerV2
Conv feature
MAM Reward
Success Raio (36)
(MAR))
Success Rala (36)
Rata (%)
predction
60
60
40
20
20
20
0.0
04
0.
0.2
04 06
1.0
2.0
0.2
0.0
Nommalaae and Environment Slaps
Normlized Enrironment Stape
Normallzed Ensironment Steps
(a) Feature masking
(b) Masking ratio
(c) Reward prediction
Figure 6: Learning curves on three manipulation tasks from Meta-world that investigate the effect of (a) convolutional feature masking, (b) masking ratio, and (c) reward prediction. The solid line and shaded regions represent the mean and stratified bootstrap confidence interval across 12 runs.
achieves 20% success rate Reach Target, while our approach can solve the tasks with > 80% success rates. We find that this because DreamerV2 fails to capture target positions in visual observations, while method can capture such details (see Section 5.5 for relevant discussion and visualizations). However, we also note that these results are preliminary because they are still too cumple-inefficien to be used for real-world scenarios. We provide more discussion in Section 6.
5.3 DeepMind Control Suite Experiments
Environment details In order to demonstrate that our approach is generally applicable to diverse visual control tasks, we also evaluate our method on visual locomotion tasks from the widely used DeepMind Control Suite benchmark. Following a standard setup Hafner et [21], we use an action repeat of 2 and default camera configurations. We provide more details in Appendix F.
Results Figure shows that our method achieves competitive performance to Dreamer on visual locomotion tasks (i.e., Quadruped tasks), demonstrating the generality of our approach across diverse visual control tasks. which is one of a few manipulation tasks in the benchmark (see Figure 2(e) for an example), This implies that our method effective on environments where the model should capture fine-grained details like object positions. More results are available in Appendix H, where trends are similar.
also observe that our method outperforms DreamerV2 on Reach Duplo,
5.4 Ablation Study
Convolutional feature masking We compare convolutional feature masking with pixel masking (ie., MAE) in Figure 6(a), which shows that convolutional feature masking significantly outperforms pixel masking. Both results are achieved with the reward prediction. This demonstrates that enabling the model to capture fine inn-inri details within patches can be important for visual control. We also report the performance with varying masking ratio m E (0.25, 0.5, 0.75, 0.9) in Figure 6(b). As we discussed in Section 4.1, we find that m 0.75 achieves better performance than E (0.25, 0.5) because strong regularization can prevent the model from finding a shortcut from input pixels. However, also find that too strong regularization (i.e., 0.9) degrades the performance.

---

True
Predicted Predicted Recon (Dreamer) (MWM) (MWM)
Figure 7: Future frames reconstructed with the autoencoder (i.e.. Recon) and predicted by latent dynamics models (i.e., Predicted). Best viewed as video provided in Appendix B.
Reward prediction In Figure 6(c). we find that performance significantly degrades without reward prediction, which shows that the reconstruction objective might not be sufficient for learning task- relevant information. It would be an interesting future direction to develop representation learning scheme that learns task-relevant information without rewards because they might not be available in practice. We provide more ablation studies and lcarning curves individual tasks in Appendix I. Moreover, we also show that reward prediction allows for learning useful representations that can be generally applicable diverse robotic manipulation tasks in Appendix L.
5.5 Qualitative Analysis
We visually investigate how our world model works compared the world model DreamerV2. Specifically, we visualize the future frames predicted by latent dynamics models on Reach Target from RLBench in Figure 7. this task, robot arm should reach target position specified by a red block in visual observations (see Figure 2(c)), which changes every trial. Thus it is crucial for the model to accurately predict the position red blocks for solving the tasks. We find that our world model effectively captures the position of red blocks, while Dreamer V fails. Interestingly, observe that our latent dynamics model ignores the components that are not task-relevant such as blue and orange blocks, though the reconstructions from the autoencoder This shows how our decoupling approach works: it encourages the autoencoder to focus on learning representations capturing the details and the dynamics model to focus modeling task-relevant components of environments. We provide more examples in Appendix B.
also
capturing all the details.
Discussion
We have presented Masked World Models (MWM), which is visual model-based RI framework that decouples visual representation learning and dynamics learning. By learning latent dynamics model operating in the latent space of self-supervised ViT, we find that our approach allows for solving variety of visual control tasks from Meta-worid, RLBench, and DeepMind Control Suite.
Limitation Despite the results, there number of areas for improvement. As have shown in Figure 6(c), the performance of our approach heavily depends on the auxiliary reward prediction task. This might be because our autoencoder is not learning temporal information, which crucial for learning task-relevant information. It would be interesting to investigate the performance of video representation learning with VITs [36, 51]. would also interesting to study introducing auxiliary prediction for other modalities, such as audio. Another weakness is that our model operates only on RGB pixels from a single camera viewpoint; we look forward a future work that incorporates different input modalities such as proprioceptive states and point clouds, building on top of the recent multi-modal learning approaches [52, 53). Finally, our approach trains behaviors from scratch, which makes it still too sample-inefficient to be used in real-world scenarios. Leveraging a small number of demonstrations, incorporating the action mode with path planning [47], or pre-training a world model on video datasets [54] are directions hope investigate in future works.

---

Acknowledgments
We would like to thank Jongjin Park and Sihyun Yu for helpful discussions. We also thank Cirrascale Cloud Services" for providing compute resources. This work was partially supported by Office of Naval Research under grant N00014-21-1-2769 and grant 900014-22-1-2121, the Darpa RACER program, the Hong Kong Centre for Logistics Robotics, BMW, and Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2019-0-000)) Artificial Intelligence Graduate School Program (KAIST)).
References
[1] K. Chua, R. Calandra, R. McAllister, and S. Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Advances neural information processing systems, 2018.
[2] M. Deisenroth and C. E. Rasmussen. Pilco:
A model-based and data-efficient approach to policy search. In International Conference on Machine Learning, 2011.
[3] I. Lenz, R. A. Knepper, and A. Saxena. Deepmpe: Learning deep latent features for model predictive control. In Robotics: Science and Systems, 2015.
[4] T. Kurutach, I. Clavera, Y. Duan, A. Tamar, and P. Abbeel. Model-ensemble trust-region policy optimization. In International Conference on Learning Representations, 2018.
[5] M. Janner, J. Fu, M. Zhang, and S. Levine. When to trust your model: Model-based policy optimization. Advances in Neural Information Processing Systems, 2019.
[6] C. Finn and S. Levine. Deep visual foresight for planning robot motion. In 2017 IEEE International Conference on Robotics and Automation (ICRA), 2017.
[7] F. Ebert, C. Finn, S. Dasari, A. Xie, A. Lee, and S. Levine. Visual foresight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint 18/2.08//200050 2018.
[8] M. Watter, J. Springenberg, J. Boedecker, and M. Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in neural information processing systems, 2015.
[9] D. Hafner, T. Lillicrap, Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, 2019.
[10] M. Zhang, S. Vikram, L. Smith, P. Abbeel, M. Johnson, and S. Levine. Solar: Deep struc- tured representations for model-based reinforcement learning. In International Conference on Machine Learning, 2019.
[11] D. Ha and J. Schmidhuber. World models. In Advances in Neural Information Processing Systems, 2018
[12] D. P. Kingma and M. Welling. Auto-cncoding variational bayes. International Conference on Learning Representations, 2014.
[13] K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick. Masked autocncoders are scalable vision learners. arXiv preprint arXiv:2111.06377 2021.
[14] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Hcigold, S. Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representaitons, 2021.
[15] D. Hafner, T. Lillicrap, M. Norouzi, and J. Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.

---

[16] T. Yu, D. Quillen, Z. He. R. Julian, K. Hausman, C. Finn, and S. Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning, 2020.
[17] S. James, Z. Ma. D. R. Arrojo, and A. J. Davison. RLBench:
The robot learning benchmark & learning environment. IEEE Robotics and Automation Letters, 2020.
[18] H. Touvron, M. Cord, A. El-Nouby, J. Verbeek, and H. Jégou. Three things everyone should know about vision transformers. arXiv preprint arXiv:2203.09799, 2022.
[19] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. Imagenet: A large-scale hierarchical image database. 2009 IEEE conference on computer vision and pattern recognition, 2009.
[20] C. Finn, X. Y. Tan, Y. Duan, T. Darrell, S. Levine, and P. Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), 2016.
[21] D. Hafner, T. Lillicrap, Ba, and M. Norouzi. Dream to control: imagination. In International Conference on Learning Representations, 2020.
Learning behaviors by latent
[22] L. Kaiser, M. Babaeizadeh, P. Milos, B. Osinski, H. Campbell, K. Czechowski, D. Erhan, C. Finn, P. Kozakowski, S. Levine, et al. Model-based reinforcement learning for atari. In International Conference on Learning Representations, 2019.
[23] A. Gupta, S. Tian, Y. Zhang, J. Wu, R. Marttn-Martin, and L. Fei-Fei. Maskvit: Masked visual pre-training for video prediction. arXiv preprint arXiv:2206. /8894 2022.
[24] W. Ye, S. Liu, T. Kurutach, Abbeel, and Y. Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 2021.
[25] Seyde, Schwarting. S. Karaman, and D. Rus. Learning to plan optimistically: Uncertainty- guided deep exploration via latent model ensembles. In Conference on Robot Learning, 2021.
[26] O. Rybkin, C. Zhu, Nagabandi, K. Daniilidis, I. Mordatch, and S. Levine. Model-based reinforcement learning via latent-space collocation. In International Conference on Machine Learning, 2021.
[27] C. Gelada, S. Kumar, J. Buckman, O. Nachum, and M. G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. International Conference on Machine Learning, 2019.
[28] T. D. Nguyen, R. Shu, T. Pham, H. Bui, and S. Ermon. Temporal predictive coding for model-based planning in latent space. In International Conference on Machine Learning, 2021.
[29] M. Okada and T. Taniguchi. Dreaming: Model-based reinforcement learning by latent imagina- tion without reconstruction. In 2021 IEEE International Conference on Robotics and Automation (ICRA), 2021.
[30] Deng, Jang, and S. Ahn. Dreamerpro: Reconstruction-free model-based reinforcement learning with prototypical representations. arXiv preprint 10Xiv:2110.14565, 2021.
[31] X. Chen, S. Xie, and K. He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021.
[32] M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEF/CVF International Conference on Computer Vision, 2021.
[33] G. Hinton, O. Vinyals, Dean, et al. Distilling the knowledge in a neural network. arXiv preprint arXiv: 1503.02531, 2015.
[34] H. Bao, I.. Dong, and F. Wei. Beit: Bert pre-training of image transformers. arXiv preprint 22X0:22006.0824,, 2021.
10

---

[35] Z. Li, Z. Chen, F. Yang, W. Li, Y. Zhu, C. Zhao, R. Deng, L. Wu, R. Zhao. M. Tang, et al. Mst: Masked self-supervised transformer for visual representation. In Advances in Neural Information Processing Systems, 2021.
[36] C. Feichtenhofer, H. Fan, Y. Li, and K. He. Masked autoencoders spatiotemporal learners. arXiv preprint wwX.::2005.0911.. 2022.
[37] Z. Xie, Z. Zhang, Y. Cao, Y. Lin, J. Bao, Z. Yao, Dai, and H. Hu. Simmim: A simple framework for masked image modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
[38] C. Wei, H. Fan, S. Xie, C.-Y. Wu, A. Yuille, and C. Feichtenhofer. Masked feature prediction for self-supervised visual pre-training. Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, 2022.
[39] J. Zhou, C. Wei, H. Wang. W. Shen, C. Xie, A. Yuille, and 'T. Kong. ibot: Image bert pre-training with online tokenizer. arXiv preprint arXiv:2111.0720 11.0322, 2021.
[40] T. Xiao, M. Singh, H. Mintun, T. Darrell, P. Dollár, and R. Girshick. Early convolutions help transformers see better. In Advances Neural Information Processing Systems, 2021.
[41] Y. Tassa, S. Tunyasuvunakool, A. Muldal, Y. Doron, Liu, S. Bohez, J. Merel, Erez, T. Lillicrap, and N. Heess. dm control: Software and tasks for continuous control. arXiv preprint arXiv:2006. 12983, 2020.
[42] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT Press, 2018.
[43] A. Alemi, B. Poole, I. Fischer, J. Dillon, R. A. Saurous, and K. Murphy, Fixing broken elbo.
In International Conference on Machine Learning, 2018
[44] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, E. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
[45] B. Graham, A. El-Nouby, H. Touvron, P. Stock, A. Joulin, H. Jégou, and M. Douze. Levit: a vision transformer in convnet's clothing for faster inference. Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021.
[46] Y. Tassa, T. Erez, and E. Todorov. Synthesis and stabilization of complex behaviors through online trajectory optimization. In International Conference on Intelligent Robots and Systems, 2012.
[47] S. James and A. J. Davison. Q-attention: Enabling efficient learning for vision-based robotic manipulation. IEEE Robotics and Automation Letters, 2022.
[48] James, K. Wada, T. Laidlow, and A. J. Davison. Coarse-to-Fine Q-attention: Efficient learning for visual robotic manipulation via discretisation. Proceedings of the IEEF/CVE Conference on Computer Vision and Pattern Recognition, 2022.
[49] S. James and P. Abbeel. Coarse-to-Fine Q-attention with Learned Path Ranking. arXiv preprint arXiv:2204.0157/ 2022.
[50] S. James and P. Abbeel. Coarse-to-Fine Q-attention with Tree Expansion. arXiv preprint arXiv:2204. 12471, 2022.
[51] A. Arnab, M. Dehghani, G. Heigold, C. Sun, M. Lučic, and C. Schmid. Vivit: video vision transformer. In Proceedings of the IEEE/CVF International Conference Computer Vision, 2021.
[52] X. Geng, H. Liu, L. Lec, D. Schuurams, S. Levine, and P. Abbeel. Multimodal masked autoencoders learn transferable representations. arXiv preprint arXiv:2205. 14204, 2022.
[53] R. Bachmann, D. Mizrahi, A. Atanov, and A. Zamir. Multimae: Multi-modal multi-task masked autoencoders. arXiv preprint arX1v:220400167,, 2022.

---

[54] Y. Seo, K. Lee, S. James, and P. Abbeel. Reinforcement learning with action- free pre-training from videos. In International Conference on Machine Learning. 2022.
[55] Schulman, Moritz, S. Levine, M. Jordan, and Abbeel. High-dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv: 5t0:::/006002438, 2015.
[56] K. Hc, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, 2016.
[57] H. Wu, B. Xiao, N. Codella, M. Liu, X. Dai, L. Yuan, and L. Zhang. Cvt: Introducing convolutions to vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021.
[58] Dai, K. He, and J. Sun. Convolutional feature masking for joint object and stuff segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1992-4000, 2015.
[59] Tompson, R. Goroshin, A. Jain, Y. LeCun, and C. Bregler. Efficient object localization using convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, 2015.
[60] G. Ghiasi, T.-Y. Lin, and Q. V. Le. Dropblock: regularization method for convolutional networks. Advances in neural information processing systems, 2018.
[61] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, 2014.
[62] J. Park, Y. Seo, C. Liu, L. Zhao, T. Qin, J. Shin, and T.-Y. Liu. Object-aware regularization for addressing causal confusion in imitation learning. Advances in Neural Information Processing Systems, 2021.
[63] A. Van Den Oord, O. Vinyals, et al. Neural discrete representation learning. In Advances in Neural Information Processing Systems, 2017.
[64] M. Jaderberg, V. Mnih, W. M. Czarnecki, TT. Schaul, J. 7. Leibo, D. Silver, and K. Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations, 2017.
[65] M. Schwarzer, A. Anand, R. Goel, R. D. Hjelm, A. Courville, and Bachman, Data-efficient reinforcement learning with self-predictive representations. In International Conference on Learning Representations, 2021.
[66] Yu, C. Lan, W. Zeng, M. Feng, Z. Zhang, and Z. Chen. Playvirtual: Augmenting cycle- consistent virtual trajectories for reinforcement learning. In Advances in Neural Information Processing Systems, 2021.
[67] Yu, Z. Zhang, C. Lan, Z. Chen, and Y. Lu. Mask-based latent reconstruction for reinforcement learning, arXiv preprint arXiv:2201.12000 2022.
[68] P. S. Castro. Scalable methods for computing state similarity in deterministic markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence, 2020.
[69] A. Zhang, R. McAllister, R. Calandra, Y. Gal, and Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
[70] A. v. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. In Advances in Neural Information Processing Systems, 2018.
[71] A. Anand, E. Racah, S. Ozair, Y. Bengio, M.-A. Côté, and R. D. Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, 2019.
[72] B. Mazoure, R. T. d. Combes, Doan, P. Bachman, and D. Hjelm. Deep reinforcement and infomax learning. In Advances in Neural Information Processing Systems, 2020.
12

---

[73] A. Srinivas, M. Laskin, and P. Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In Internatial Conference on Machine Learning. 2020.
[74] H. Liu and P. Abbeel. Behavior from the void: Unsupervised active pre-training, arXiv preprint arXiv:2103.04551, 2021.
[75] D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Reinforcement learning with prototypical representations. In International Conference on Machine Learning. 2021.
[76] S. Nair, A. Rajeswaran, V. Kumar, C. Finn, and A. Gupta. R3mm: universal visual representation for robot manipulation. arXiv preprint arXiv: 12601, 2022.
[77] S. Parisi, A. Rajeswaran, S. Purushwalkam, and Gupta. The unsurprising effectiveness of pre-trained vision models for control. arXiv preprint arXiv:2203.0353 033580 2022.
[78] M. Minderer, C. Sun, R. Villegas, F. Cole, K. P. Murphy, and H. Lee. Unsupervised learning of object structure and dynamics from videos. Advances in Neural Information Processing Systems, 2019.
[79] L. Manuelli, Y. Li, P. Florence, and R. Tedrake. Keypoints into the future: Self-supervised correspondence in model-based reinforcement learning. arXiv preprint arXiv:2009.0500 05000, 2020.
[80] M. Lambeta, P.-W. Chou, S. Tian, B. Yang, B. Maloon, V. R. Most, D. Stroud, Santos, A. Byagowi, G. Kammerer, et al. Digit: A novel design for a low-cost compact high-resolution tactile sensor with application in-hand manipulation. IEEE Roboties and Automation Letters, 2020.
[81] N. Das, S. Bechtle, T. Davehev, D. Jayaraman, A. Rai, and F. Meier. Model-based inverse reinforcement learning from visual demonstrations. In Conference on Robot Learning, 2021.
[82] D. Yarats, A. Zhang, Kostrikov, B. Amos, J. Pincau, and R. Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence, 2021.
[83] E. Shelhamer, P. Mahmoudieh, M. Argus, and T. Darrell. Loss is its own reward: Self- supervision for reinforcement learning. arXiv preprint 111 2016.
[84] D. Yarats, R. Fergus, Lazaric, and L. Pinto. Mastering visual continuous control: Improved data- autte-u reinforcement learning. arXiv preprint arXiv:2107.09645, 2021.
[85] M. Laskin, K. Lee, A. Stooke, L. Pinto, P. Abbeel, and A. Srinivas. Reinforcement learning with augmented data. Advances in Neural Information Processing Systems, 2020.
[86] I. Kostrikov, D. Yarats, and R. Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. In International Conference Learning Representations, 2021.
[87] T. Xiao, I. Radosavovic, T. Darrell, and J. Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173 2022.
13


Reviews:
[{"review_id": "tBN_gCT4GzB", "reviewer": "Reviewer_DGEG", "summary": "This paper presents Masked World Models (MVW), a model-based reinforcement learning framework for visual control.\n\nThere are a few unique designs of MVM compared to existing approaches:\n- MVM decouples visual representation learning from dynamics learning, this is different from the existing Dreamer framework that uses a single model to learn both visual representations and dynamics together.\n- MVM uses a self-supervised Visual Transformer for representation learning, and it performs masking on convolutional features, instead of masking on pixels.\n- MVM also uses an auxiliary reward prediction to encode task-relevant information.\n\nExperiments on synthetic datasets show that MVW significantly outperforms the state of the art.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- It seems to me that, MVM incorporates a few design changes (mentioned in the summary part of my review) to the Dreamer approach and showed that these design changes lead to significant performance gain.\n- The paper reads very well. Although I am not in this field, I was able to quickly grasp the related work and the unique contributions of this paper.\n\nWeaknesses:\n- Although the paper is well written, the contributions seem incremental -- it is adding small changes and tweaks to existing approaches.\n- Real-world robotic experiments would be more convincing.", "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": "Overall, interesting ideas and good results. I vote for weak accept.", "technical_quality": "Very Good:  Convincing.", "clarity_of_presentation": "Very Good:  Well organized and clear.", "potential_impact": "", "robotics_focus": "Highly relevant to robotics but no hardware experiments", "quality_of_limitations": "Limitations are addressed clearly", "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": {"impact": "4: The work contains interesting new ideas that will potentially have major impact in robotics or machine learning.", "issues": "I am curious does decoupling the visual learning from the dynamics learning make it \"slower\" to train the approach, compared to having a single network doing both tasks? ", "reviewer_expertise": "3: The reviewer is fairly confident that the evaluation is correct"}, "extra_text": {}}, {"review_id": "i8cKgyii7oNI", "reviewer": "Reviewer_abdd", "summary": "The paper considers model-based RL with visual input. The key idea is to decouple visual representation learning and dynamics learning. For visual representation learning, inspired by the recent masked auto-encoder (MAE), the paper trains an autoencoder with convolutional layers and vision transformers (ViT) to reconstruct pixels given masked convolutional features. To encode task-specific information, the paper introduces an auxiliary reward prediction objective for the autoencoder. For dynamics learning, the paper adopts the DreamerV2 approach to learn a latent dynamics model that operates on visual representations from the autoencoder. The paper presents performance results on a variety of visual robotic tasks from Meta-world and RLBench.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\nThe paper proposes MWM consisting of an autoencoder with convolutional layers and vision transformers (ViT) to reconstruct pixels given masked convolutional features,  an auxiliary reward prediction objective for the autoencoder, a latent dynamics model that operates on visual representations from the autoencoder. \n\nIt shows that self-supervised ViT trained to reconstruct visual observations with convolutional feature masking can be effective for visual model-based RL. \n\nWeaknesses\n\nHow MWM differs from DreamerV2 is not clearly stated. The main difference seems to be a redesign of the representation model and image decoder. The two models of MWM use ViT encoder and decoder with conv feature masking. \n\nThere are several questions regarding the comparison with DreamerV2. What codebase does the paper use? Since DreamerV2 only evaluated on Atari games, does the paper make any changes to DreamerV2? The poor performance might be due to no adaptation of model architecture and hyperparameters. How does MWM compare with DreamerV2 on Atari games?\n\nThe most relevant paper is cited in appendix,\nT. Xiao, I. Radosavovic, T. Darrell, and J. Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173, 2022.\n\nwhich uses features from a masked auto-encoder. The paper should compare MWM with this paper.\n\n\n\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": "The paper introduces masked auto-encoder features and conv feature masking to DreamerV2. This redesign of DreamerV2 inspired by vision transformer and masked auto-encoder seems to be very appealing and promising since the feature learning can potentially be much better. However, the comparison with DreamerV2 lacks details. Comparing DreamerV2 designed for Atari games on Meta-world and RLBench with no change does not seem to be convincing. One way to address this concern is to compare MWM directly with DreamerV2 on Atari games.\n\nFurthermore, the Xiao et al. paper has applied masked auto-encoder to motion control. The paper should also consider comparing with the this most related paper cited in appendix.\n\n===\nSeveral of my concerns are clarified during rebuttal. I upgrade my rating to weak accept from weak reject.", "technical_quality": "Fair: Significant additional evidence required to support the claims.", "clarity_of_presentation": "Good: Minor revision needed to improve organization or clarity.", "potential_impact": "", "robotics_focus": "Highly relevant to robotics but no hardware experiments", "quality_of_limitations": "Limitations are not well addressed", "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": {"impact": "3: The work is incremental and unlikely to have much impact, even though it may be technically correct and well executed.", "reviewer_expertise": "4: The reviewer is confident but not absolutely certain that the evaluation is correct"}, "extra_text": {"issues": "It seems that predicting reward is very important to achieve high performance. The paper should provide more information on what features  are enabled by predicting reward. How does this finding change with different environment and tasks? Since the paper evaluates only in a simulator, can reward be effectively predicted for real world data? \n\nIt seems that conv feature masking is a type of dropout. Instead of conv feature masking, if dropout is used, will it achieve the same performance results?"}}, {"review_id": "_EFIVRlSLcH", "reviewer": "Reviewer_GUgv", "summary": "This paper successfully achieved visual representation learning for world models in visual model-based RL with good accuracy using appropriate transformer models and masking processes.\nAs a result, the proposed method achieved higher success rates than conventional methods in variety of robot tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The proposed method was validated in a wide variety of ways and demonstrated sufficient experimental value.\n- The intention of the model design was clearly explained qualitatively.\n\nWeaknesses:\n- Although the proposal is a decoupling of representation learning and dynamics learning, the reward function is included in representation learning, which makes the visual representation task-relevant and lacks versatility.\n- Although the masking method was changed, its benefit was only shown qualitatively and experimentally, so I cannot find mathematical explanations.\n- The analysis of why the learning curve was better than the baseline was insufficient, since only the learning curve was included in this paper.", "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": "The performance of the proposed method was experimentally verified by a large number of simulations, including ablation studies.\nOn the other hand, the proposed method was not mathematically convincing.\nToo much emphasis was placed on exceeding the baseline in terms of the control performance, so further analysis may be needed to understand why the proposed method was better.", "technical_quality": "Good: Generally convincing, but some additional experiments or analyses needed.", "clarity_of_presentation": "Very Good:  Well organized and clear.", "potential_impact": "", "robotics_focus": "Highly relevant to robotics but no hardware experiments", "quality_of_limitations": "Additional details required", "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": {"impact": "3: The work is incremental and unlikely to have much impact, even though it may be technically correct and well executed.", "reviewer_expertise": "3: The reviewer is fairly confident that the evaluation is correct"}, "extra_text": {"issues": "I think, for the versatility, it would be better to learn visual representation and dynamics together because they are task-irrevant (and ultimately, environment-irrevant) components, and the reward prediction should be decoupled from them.\nIt seems to me that some other task-irrelevant information should be utilized to prioritize observations, as discussed in Limitation.\n\nThe addition of the reward prediction finally led to exceed the baseline, so I feel that it would be more important than masking and decoupling, which are in the title of this paper.\n\nI know the page limitation is strict, but the authors putted too much information in Appendix, causing bad readability.\nAt least, in my opinion, the optimization problem for the proposed method should be inserted in the main text, rather than the definition of Dreamer."}}, {"review_id": "YU8EZ9xRjMG", "reviewer": "Reviewer_btzk", "summary": "This paper improves the DreamerV2 model-based RL framework, by replacing the the discrete world model encoder (looks like a VQ-VAE style encoder with recurrency? but I' m not sure) to a MAE-like encoder (The main difference is that it's masking the feature instead of the origin image). Experiments show that the proposed method achieves much better performance tha the original DreamerV2.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- Strength: It achieves much better result than the baseline method in several tasks in the settings that the author propose.\n\n- Weakness: Experiment settings seem to have a huge gap between this submission and the original DreamerV2 paper. To be more specific, in the dreamer V2 paper (https://arxiv.org/pdf/2010.02193.pdf), it reports result after 4e7 steps for Humanoid Walk (2010.02193.pdf Appendix Fig A.2). Yet in this paper, there seems to be no more than 3e6 steps executed in any of its setups. As we all know, the training process of RL agents are unstable and the may suffer from zero-reward at the beginning of its training process, it would be important to wait until the training process to converge. Furthermore, in Figure 3 of this submission, Hand Insert / Shelf Place task increase concern about this issue, where we notice the DreamerV2 reward begin to increase after ~2e6/~3e6 steps, which is quite similar to 2010.02193.pdf Appendix Fig A.2.\n\nI'm also quite confused why the paper does not report any Atari Game or Humanoid task in the DM Control Suite, which are reported by the original DreamerV2 paper. \n\nSince the core contribution of the submission is additional (i.e. change the encoder), experiments details and solidity are important.", "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": "See the Weakness section. My main concern about this paper is the experiment part that, as has been mentioned in the Weakness section. Since the core contribution of the submission is additional (i.e. change the encoder), experiments details and solidity are important. Since the experiment setups may be not aligned with the DreamerV2 paper, and the paper does not report any Atari Game or Humanoid task in the DM Control Suite, which are reported by the original DreamerV2 paper, I am not very certain if the experiment part of this paper is convincing enough.\n\nThat being said, I myself do not have actual experience training model-based RL agents, so I rate the confidence (Reviewer Expertise) to be 2. If there are any severly technically wrong part of my review please the authors and the ACs point them out.", "technical_quality": "Fair: Significant additional evidence required to support the claims.", "clarity_of_presentation": "Good: Minor revision needed to improve organization or clarity.", "potential_impact": "", "robotics_focus": "Relevant but unlikely to deploy to hardware in near future", "quality_of_limitations": "Additional details required", "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": {"impact": "3: The work is incremental and unlikely to have much impact, even though it may be technically correct and well executed.", "reviewer_expertise": "2: The reviewer is willing to defend the evaluation, but it is quite likely that the reviewer did not understand central parts of the paper"}, "extra_text": {"issues": "- Concerns about experiment details mentioned in the Weakness section.\n- There are some obvious mistakes in the figures. For example, I noticed that unit of success rate is \"%\" yet the scale is from 0.0 - 1.0. Obviously the author means 0.0 - 100.0."}}, {"review_id": "Xk1VTXykbnC", "reviewer": "Reviewer_iNVH", "summary": "The paper presents a method for decoupling visual representation and latent dynamics learning for control. It does so by combining  recent approaches for visual representation learning (ViT) and latent dynamics learning (Dreamer) for control such that they are decoupled. This decoupling allows for learning a better prediction model as demonstrated through experiments in a diverse set of environments. ", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper is well written and well structured \n- has a number of interesting experiments that show the proposed methods superiority to the chosen baseline Dreamer V2 as well as an extensive ablation study\n- Learning the reward alongside latent features in ViT is an interesting way of encoding task relevant information in visual representation architecture.\n\nWeaknesses\n- Missing relevant works\n- Only one baseline that does not decouple visual representation and dynamics learning \n- Will the auxilliary reward prediction in the visual representation learning not limit the generalizability of the learned visual representor to only one RL task?\n- Minor issues such as missing acronym expansion\n- Lack of hardware experiments\n\nQuestions/Doubts:\n- In Figure 6.a, is the plot for Pixel patch (MAE) produced with reward learning?", "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": "The paper combines two recent approaches for visual representation learning and latent dynamics learning for control such that they are decoupled. While this proposal is not entirely novel, this paper demonstrates the advantage of this decoupling empirically with several experiments in different simulation environments. There is one novel idea (in my opinion) - the auxilliary reward learning for encoding task relevant information. However, this might limit the learned visual representation's generalizability to different tasks.", "technical_quality": "Good: Generally convincing, but some additional experiments or analyses needed.", "clarity_of_presentation": "Very Good:  Well organized and clear.", "potential_impact": "", "robotics_focus": "Highly relevant to robotics but no hardware experiments", "quality_of_limitations": "Limitations are addressed clearly", "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": {"impact": "3: The work is incremental and unlikely to have much impact, even though it may be technically correct and well executed.", "reviewer_expertise": "3: The reviewer is fairly confident that the evaluation is correct"}, "extra_text": {"issues": "- Missing relevant works specifically with respect to decoupled visual representation learning and dynamics learning:\n\t- [1] uses an special bottlenecked autoencoder that generates relevant 2D keypoints on the image for learning visual representation separately from learning a dynamics model on top of these features.  [2, 3] demonstrate this method's effectiveness on hardware including arm robots [3]. This method can be considered an additional baseline.\n\t- Another relevant work in this regard includes  [4] which learns Dense Descriptions as visual representation  and learns a dynamics model from these representations separately. This method requires RGBD images, however.\n- Minor:\n\t- On Page 3, CLS is not defined.\n\nReferences :\n\n[1] Minderer, M., Sun, C., Villegas, R., Cole, F., Murphy, K. P., & Lee, H. (2019). Unsupervised learning of object structure and dynamics from videos. Advances in Neural Information Processing Systems, 32.\n\n[2] Lambeta, M., Chou, P. W., Tian, S., Yang, B., Maloon, B., Most, V. R., ... & Calandra, R. (2020). Digit: A novel design for a low-cost compact high-resolution tactile sensor with application to in-hand manipulation. IEEE Robotics and Automation Letters, 5(3), 3838-3845.\n\n[3] Das, N., Bechtle, S., Davchev, T., Jayaraman, D., Rai, A., & Meier, F. (2020). Model-based inverse reinforcement learning from visual demonstrations. CoRL 2020.\n\n[4] Manuelli, L., Li, Y., Florence, P., & Tedrake, R. (2020). Keypoints into the future: Self-supervised correspondence in model-based reinforcement learning. arXiv preprint arXiv:2009.05085."}}]

Author rebuttal:

Decision: Accept (poster)