# Mask-based Latent Reconstruction for Reinforcement Learning

arXiv: https://arxiv.org/abs/2201.12096 · HTML: https://mlsanity.com/paper/2201.12096/

## Key facts

- Cited by 39 papers in the map (data through 14 Oct 2026).
- Builds on 54 of its 70 references that are in the map.
- First posted on arXiv 28 Jan 2022.

## Builds on (most cited first, 30 of 54)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Context Encoders: Feature Learning by Inpainting](https://mlsanity.com/paper/1604.07379/) (2016)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [SimMIM: A Simple Framework for Masked Image Modeling](https://mlsanity.com/paper/2111.09886/) (2021)
- [Prioritized Experience Replay](https://mlsanity.com/paper/1511.05952/) (2015)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Dueling Network Architectures for Deep Reinforcement Learning](https://mlsanity.com/paper/1511.06581/) (2015)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Masked Feature Prediction for Self-Supervised Visual Pre-Training](https://mlsanity.com/paper/2112.09133/) (2021)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)

## Built on by (most cited first, 30 of 39)

- [Masked World Models for Visual Control](https://mlsanity.com/paper/2206.14244/) (2022)
- [Reinforcement Learning with Action-Free Pre-Training from Videos](https://mlsanity.com/paper/2203.13880/) (2022)
- [PLASTIC: Improving Input and Label Plasticity for Sample Efficient Reinforcement Learning](https://mlsanity.com/paper/2306.10711/) (2023)
- [Towards General-Purpose Model-Free Reinforcement Learning](https://mlsanity.com/paper/2501.16142/) (2025)
- [Revisiting Plasticity in Visual Reinforcement Learning: Data, Modules and Training Stages](https://mlsanity.com/paper/2310.07418/) (2023)
- [A Comprehensive Survey of Data Augmentation in Visual Reinforcement Learning](https://mlsanity.com/paper/2210.04561/) (2022)
- [Learning Latent Dynamic Robust Representations for World Models](https://mlsanity.com/paper/2405.06263/) (2024)
- [A Survey of State Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2506.17518/) (2025)
- [Light-weight probing of unsupervised representations for Reinforcement Learning](https://mlsanity.com/paper/2208.12345/) (2022)
- [MaDi: Learning to Mask Distractions for Generalization in Visual Deep Reinforcement Learning](https://mlsanity.com/paper/2312.15339/) (2023)
- [Learning to Identify Critical States for Reinforcement Learning from Videos](https://mlsanity.com/paper/2308.07795/) (2023)
- [Learning Better with Less: Effective Augmentation for Sample-Efficient Visual Reinforcement Learning](https://mlsanity.com/paper/2305.16379/) (2023)
- [Learning to Discover Skills through Guidance](https://mlsanity.com/paper/2310.20178/) (2023)
- [MA2CL:Masked Attentive Contrastive Learning for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2306.02006/) (2023)
- [DEAR: Disentangled Environment and Agent Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2407.00633/) (2024)
- [Investigating Pre-Training Objectives for Generalization in Vision-Based Reinforcement Learning](https://mlsanity.com/paper/2406.06037/) (2024)
- [World Model Robustness via Surprise Recognition](https://mlsanity.com/paper/2512.01119/) (2025)
- [Aerial Image Object Detection With Vision Transformer Detector (ViTDet)](https://mlsanity.com/paper/2301.12058/) (2023)
- [Sequential Action-Induced Invariant Representation for Reinforcement Learning](https://mlsanity.com/paper/2309.12628/) (2023)
- [Learning Future Representation with Synthetic Observations for Sample-efficient Reinforcement Learning](https://mlsanity.com/paper/2405.11740/) (2024)
- [MOOSS: Mask-Enhanced Temporal Contrastive Learning for Smooth State Evolution in Visual Reinforcement Learning](https://mlsanity.com/paper/2409.02714/) (2024)
- [Revisiting Bisimulation Metric for Robust Representations in Reinforcement Learning](https://mlsanity.com/paper/2507.18519/) (2025)
- [Spectral Representation-based Reinforcement Learning](https://mlsanity.com/paper/2512.15036/) (2025)
- [Representation Learning for Sequential Volumetric Design Tasks](https://mlsanity.com/paper/2309.02583/) (2023)
- [Efficient Reinforcement Learning Through Adaptively Pretrained Visual Encoder](https://mlsanity.com/paper/2502.05555/) (2025)
- [PvP: Data-Efficient Humanoid Robot Learning with Proprioceptive-Privileged Contrastive Representations](https://mlsanity.com/paper/2512.13093/) (2025)
- [Mask-Based Modeling for Neural Radiance Fields](https://mlsanity.com/paper/2304.04962/) (2023)
- [Learning Fused State Representations for Control from Multi-View Observations](https://mlsanity.com/paper/2502.01316/) (2025)
- [Diffusion Transformer Meets Random Masks: An Advanced PET Reconstruction Framework](https://mlsanity.com/paper/2503.08339/) (2025)
- [COVR:Collaborative Optimization of VLMs and RL Agent for Visual-Based Control](https://mlsanity.com/paper/2601.06122/) (2026)

## Abstract (arXiv)

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.