arXiv · 2602.19956
Sparse Masked Attention Policies for Reliable Generalization
Abstract
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.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Caroline Horsch, Laurens Engwegen, Max Weltevrede, Matthijs T. J. Spaan, Wendelin Böhmer. 2026-02-23. Sparse Masked Attention Policies for Reliable Generalization. https://arxiv.org/abs/2602.19956
Cite the original work for its findings. Save a collection to share your selection of sources.