mlsanityData to 7 Oct 2026

Sparse Masked Attention Policies for Reliable Generalization

arXiv 2602.19956 ↗First posted 23 Feb 2026cs.LGcited by 0

Not on the map. Why?

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.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 16 of its 20 references that are in the map. First posted on arXiv 23 Feb 2026.

Builds on

16 of 20 references.

References · 16

  1. 2017Attention Is All You Needcited by 59,932
  2. 2016The Concrete Distribution: A Continuous Relaxation of Discrete Random Variablescited by 1,135
  3. 2021A Survey of Zero-shot Generalisation in Deep Reinforcement Learningcited by 164
  4. 2020Generalization in Reinforcement Learning by Soft Data Augmentationcited by 132
  5. 2021Decoupling Value and Policy for Generalization in Reinforcement Learningcited by 81
  6. 2018Relational Deep Reinforcement Learningcited by 72
  7. 2020Automatic Data Augmentation for Generalization in Deep Reinforcement Learningcited by 71
  8. 2021The Sensory Neuron as a Transformer: Permutation-Invariant Neural Networks for Reinforcement Learningcited by 50
  9. 2021Dynamic Sparse Training for Deep Reinforcement Learningcited by 49
  10. 2020Neuroevolution of Self-Interpretable Agentscited by 48
  11. 2022Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learningcited by 32
  12. 2021DRIBO: Robust Deep Reinforcement Learning via Multi-View Information Bottleneckcited by 31
4 more
  1. 2021Cross-Trajectory Representation Learning for Zero-Shot Generalization in RLcited by 20
  2. 2024Don't flatten, tokenize! Unlocking the key to SoftMoE's efficacy in deep RLcited by 14
  3. 2024SiT: Symmetry-Invariant Transformers for Generalisation in Reinforcement Learningcited by 3
  4. 2025Sparse-Reg: Improving Sample Complexity in Offline Reinforcement Learning using Sparsitycited by 2

Built on by

Not cited yet.