# Sparse Masked Attention Policies for Reliable Generalization

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

## 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 16 of its 20 references that are in the map.
- First posted on arXiv 23 Feb 2026.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 16 of 16)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [Decoupling Value and Policy for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2102.10330/) (2021)
- [Relational Deep Reinforcement Learning](https://mlsanity.com/paper/1806.01830/) (2018)
- [Automatic Data Augmentation for Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/2006.12862/) (2020)
- [The Sensory Neuron as a Transformer: Permutation-Invariant Neural Networks for Reinforcement Learning](https://mlsanity.com/paper/2109.02869/) (2021)
- [Neuroevolution of Self-Interpretable Agents](https://mlsanity.com/paper/2003.08165/) (2020)
- [Dynamic Sparse Training for Deep Reinforcement Learning](https://mlsanity.com/paper/2106.04217/) (2021)
- [Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning](https://mlsanity.com/paper/2209.09203/) (2022)
- [DRIBO: Robust Deep Reinforcement Learning via Multi-View Information Bottleneck](https://mlsanity.com/paper/2102.13268/) (2021)
- [Cross-Trajectory Representation Learning for Zero-Shot Generalization in RL](https://mlsanity.com/paper/2106.02193/) (2021)
- [Don't flatten, tokenize! Unlocking the key to SoftMoE's efficacy in deep RL](https://mlsanity.com/paper/2410.01930/) (2024)
- [SiT: Symmetry-Invariant Transformers for Generalisation in Reinforcement Learning](https://mlsanity.com/paper/2406.15025/) (2024)
- [Sparse-Reg: Improving Sample Complexity in Offline Reinforcement Learning using Sparsity](https://mlsanity.com/paper/2506.17155/) (2025)

## Abstract (arXiv)

In reinforcement learning, abstraction methods that remove unnecessary information from the observation are commonly used to learn policies which generalize better to unseen tasks. However, these methods often overlook a crucial weakness: the function which extracts the reduced-information representation has unknown generalization ability in unseen observations. In this paper, we address this problem by presenting an information removal method which more reliably generalizes to new states. We accomplish this by using a learned masking function which operates on, and is integrated with, the attention weights within an attention-based policy network. We demonstrate that our method significantly improves policy generalization to unseen tasks in the Procgen benchmark compared to standard PPO and masking approaches.