mlsanityData to 14 Oct 2026

Mask-based Predictive Representations for Reinforcement Learning

arXiv 2607.04153 ↗First posted 5 Jul 2026cs.LG cs.AIcited by 0

On the map: Deep reinforcement learning methods → Deep RL algorithms and representations

Abstract (arXiv)

Vision-based deep reinforcement learning involves dealing with high-dimensional inputs of image information. It is crucial to abstract effective states from high-dimensional image inputs and limited samples for sample-efficient reinforcement learning. To address this challenge, inspired by fields such as natural language processing and computer vision, we propose a self-supervised task based on mask prediction as an auxiliary task for reinforcement learning. This non-reconstruction method uses the sequence information collected by the agent from the environment and the context information in the sequence to predict the masked information, thereby strengthening the agent's understanding of the task and learning effective representations. Combined with transformers, we find that the model reconstructs the masked input sequence in the latent space. By feeding the compressed representations learned by this method into reinforcement learning models, we observe an improvement in the sample efficiency of reinforcement learning. Moreover, the model outperforms state-of-the-art sample-efficient reinforcement learning methods on multiple continuous and discrete control benchmarks.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods. Builds on 39 of its 47 references that are in the map. First posted on arXiv 5 Jul 2026.

Builds on

39 of 47 references.

References · 39

  1. 2017Attention Is All You Needcited by 60,361
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
  3. 2020Language Models are Few-Shot Learnerscited by 32,419
  4. 2021Masked Autoencoders Are Scalable Vision Learnerscited by 7,403
  5. 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,885
  6. 2013Playing Atari with Deep Reinforcement Learningcited by 3,274
  7. 2021BEiT: BERT Pre-Training of Image Transformerscited by 2,375
  8. 2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,653
  9. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,353
  10. 2018Learning Latent Dynamics for Planning from Pixelscited by 1,264
  11. 2023Mastering Diverse Domains through World Modelscited by 1,253
  12. 2021SimMIM: A Simple Framework for Masked Image Modelingcited by 1,153
27 more
  1. 2019Dota 2 with Large Scale Deep Reinforcement Learningcited by 1,103
  2. 2020Mastering Atari with Discrete World Modelscited by 1,007
  3. 2018DeepMind Control Suitecited by 956
  4. 2017Rainbow: Combining Improvements in Deep Reinforcement Learningcited by 786
  5. 2020CURL: Contrastive Unsupervised Representations for Reinforcement Learningcited by 760
  6. 2021Offline Reinforcement Learning as One Big Sequence Modeling Problemcited by 644
  7. 2020Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixelscited by 585
  8. 2019Model-Based Reinforcement Learning for Ataricited by 502
  9. 2020Learning Invariant Representations for Reinforcement Learning without Reconstructioncited by 436
  10. 2021Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learningcited by 372
  11. 2020Data-Efficient Reinforcement Learning with Self-Predictive Representationscited by 318
  12. 2020Agent57: Outperforming the Atari Human Benchmarkcited by 303
  13. 2019Improving Sample Efficiency in Model-Free Reinforcement Learning from Imagescited by 274
  14. 2022Masked Visual Pre-training for Motor Controlcited by 248
  15. 2021Mastering Atari Games with Limited Datacited by 237
  16. 2019Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Modelcited by 231
  17. 2019Unsupervised State Representation Learning in Ataricited by 144
  18. 2020Bootstrap Latent-Predictive Representations for Multitask Reinforcement Learningcited by 102
  19. 2021DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representationscited by 74
  20. 2016Loss is its own Reward: Self-Supervision for Reinforcement Learningcited by 72
  21. 2022Masked Autoencoding for Scalable and Generalizable Decision Makingcited by 45
  22. 2021Learning Markov State Abstractions for Deep Reinforcement Learningcited by 44
  23. 2022Understanding Self-Predictive Learning for Reinforcement Learningcited by 44
  24. 2021PlayVirtual: Augmenting Cycle-Consistent Virtual Trajectories for Reinforcement Learningcited by 41
  25. 2022Mask-based Latent Reconstruction for Reinforcement Learningcited by 39
  26. 2021Return-Based Contrastive Representation Learning for Reinforcement Learningcited by 36
  27. 2022Stabilizing Off-Policy Deep Reinforcement Learning from Pixelscited by 32

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025Self-Predictive Dynamics for Generalization of Vision-based Reinforcement Learning
  2. 2025A Method of Selective Attention for Reservoir Based Agents
  3. 2025MaskFocus: Focusing Policy Optimization on Critical Steps for Masked Image Generation
  4. 2025MAGIC-MASK: Multi-Agent Guided Inter-Agent Collaboration with Mask-Based Explainability for Reinforcement Learning
  5. 2026Observation-Grounded Self-Predictive Reinforcement Learning for Visual Continuous Control
  6. 2022MaskViT: Masked Visual Pre-Training for Video Prediction
  7. 2026Neuro-symbolic Action Masking for Deep Reinforcement Learning
  8. 2024Learning Future Representation with Synthetic Observations for Sample-efficient Reinforcement Learning