Masked World Models for Visual Control
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.