Mask-based Predictive Representations for Reinforcement Learning

Abstract: 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.

• The paper presents MPR, a novel self-supervised method that predicts masked latent features using a transformer-based decoder to improve sample efficiency in RL.
• Empirical evaluations on DMControl and Atari benchmarks show MPR outperforms several SOTA methods, achieving up to 7.5% higher mean scores and surpassing human-level performance in some Atari games.
• MPR strategically employs spatial and temporal masking with EMA-updated encoders and optimized transformer depth to efficiently transfer contextual information for robust policy learning.

Mask-based Predictive Representations for Reinforcement Learning

Introduction

The paper "Mask-based Predictive Representations for Reinforcement Learning" (2607.04153) presents a self-supervised auxiliary task for vision-based deep reinforcement learning (RL) that leverages masked sequence modeling. The method, termed Mask-based Predictive Representations (MPR), jointly optimizes policy and representation learning by predicting masked features of observation sequences in the latent space. MPR draws on successes from mask-based pretraining in NLP and vision, driving sample efficiency and enhancing the quality of learned policies in both continuous and discrete control settings.

Mask-based Predictive Representation Architecture

MPR extends the paradigm of joint-embedding architectures by augmenting observation sequences with random spatial and temporal masks, then encoding the masked and target sequences via online and target encoders respectively. Representations are predicted and reconstructed in the latent (feature) space using a predictive decoder, typically a transformer. This setup efficiently transfers contextual information and avoids direct pixel-space reconstruction, thereby reducing computational cost and minimizing the influence of redundant pixel-level features.

Figure 1: The architecture of MPR implements spatial and temporal masking on observation sequences, encoding them and reconstructing masked features in the latent space using a predictive decoder.

Distinctively, the masking is designed independently of patch size, decoupling the masking strategy from vision backbone details, and the target encoder weights are updated using exponential moving average (EMA) from the context encoder to stabilize learning. The final representations from the target encoder are used as state inputs for downstream RL algorithms.

Figure 2: Masking process on observation sequences showing the original images, their masks, augmented contexts, and targets.

On Atari and DMControl benchmarks, masking configurations are tailored with different replication parameters ($R$) and block sizes for maximizing the relevance of spatial and temporal context. All self-supervised losses are based on cosine similarity in the latent space, jointly optimized with the RL loss.

Experimental Evaluation and Ablation

Sample Efficiency and SOTA Comparisons

The MPR method is systematically evaluated against several state-of-the-art model-free RL approaches, including CURL, DrQ, and MLR, on DeepMind Control Suite (DMControl) and the Atari-100k benchmarks.

On DMControl, MPR outperforms MLR in 5 out of 6 continuous control tasks at both 100k and 500k step regimes, achieving 7.5% higher mean score at 100k steps and requiring fewer parameters. In discrete control (Atari-100k), MPR achieves the highest average score in 11 out of 26 environments and consistently matches or exceeds other baselines across the board. Notably, MPR surpasses human-level performance in five Atari games.

Figure 3: Schematic depiction of MPR's application to discrete control (Atari) environments.

Masking Strategies and Architectural Components

A comprehensive ablation suite examines the impact of mask ratio, masking policy (spatial/temporal), sequence length, and predictive decoder depth.

• Optimal mask ratio is identified at 40%, with deviations in either direction reducing policy performance.
• Masking policy: Temporal masks (block size 8) tend to favor environments with high temporal redundancy (e.g., Walker, Ball-in-cup), while spatial or hybrid (spatial-temporal) masks are preferred in tasks like Cheetah-run.
• Increasing sequence length to 8 enhances the use of contextual information, whereas further elongation degrades compositionality and policy quality.
• Predictive decoder depth is optimal at two transformer layers; deeper decoders reduce sample efficiency.

  Figure 4: Ablations demonstrating the effects of action token inclusion, mask replication (R), loss metric selection, and feature masking.

In alternative designs, adding action tokens or excessive sequence replication ($R=2$) was detrimental in continuous control, indicating that task-specific context is best exploited with moderate sequence lengths and focused on visual token reconstruction rather than precise latent embedding reconstruction.

Theoretical and Practical Implications

MPR brings forth multiple implications for vision-augmented RL:

• Theoretical: Joint self-supervision via masked context promotes representations that encode both short- and long-range temporal dependencies, improving the predictive efficiency of state abstractions and potentially enhancing generalization across task variants.
• Practical: The masking approach is decoupled from the particulars of the vision encoder; this flexibility accommodates future integrations of segmentation-aware or multimodal masks, such as those generated by SAM and BLIP. MPR is compatible with diverse RL algorithms (e.g., SAC, Rainbow) and scales without bespoke augmentations for either the image pipeline or action sequence modeling.

The demonstrated empirical improvements, especially under strict sample constraints, highlight MPR's effectiveness for real-world or resource-limited reinforcement learning applications.

Future Directions

MPR opens avenues for integrating more sophisticated and semantically informed mask generation. For example, masks derived from image segmenters (e.g., object-centric or language-guided) could further enrich context and potentially yield more disentangled and robust state abstractions. Additionally, scaling MPR with larger or multimodal datasets, or in the context of world models, could push the boundary of sample-efficient policy learning in partially observed or highly stochastic domains.

Conclusion

Mask-based Predictive Representations (MPR) offer an effective, computationally efficient, and highly adaptable self-supervised auxiliary task for RL in both continuous and discrete settings. By relying on masked sequence reconstruction in the latent space and leveraging transformer-based predictive decoders, MPR advances the state of sample-efficient policy learning, consistently outperforming prior work in standard RL benchmarks. The method's decoupling from traditional patch-based masks and action sequence conditioning marks a distinctive theoretical and practical contribution to the field, setting a robust foundation for future work on representation-driven reinforcement learning.