Sparse Masked Attention Policies for Reliable Generalization

Caroline Horsch; Laurens Engwegen; Max Weltevrede; Matthijs T. J. Spaan; Wendelin Böhmer

cs.LGFeb 23, 2026

Sparse Masked Attention Policies for Reliable Generalization

Authors: Caroline Horsch, Laurens Engwegen, Max Weltevrede, Matthijs T. J. Spaan, Wendelin Böhmer

Organizations: Department of Intelligent Systems, Delft University of Technology, Delft, Netherlands.

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.

Figures & tables

Panel of Figure 1 — Figure 1: A demonstration of the importance of removing unnecessary information from observations for the purpose of policy generalization. In this task, the agent needs to collect the coin at the far right. Both states look very different, but share the same immediate policy (move right). This decision can be reduced from considering the entire observation to just some areas highlighted in green (around the agent, coin location, and position to the right of the agent). In this reduction, the states are almost identical.

Figure 2: The CNN feature extractor and positional encoding setup for processing inputs. This set of features are the input tokens to the attention block.

Figure 3: The normalized returns in unseen tasks for each of the Procgen games. The mean and standard error over 10 seeds is plotted.

Figure 4: The test returns (mean and standard error over 10 seeds) on unseen tasks from the dodgeball and maze environments. As expected, sparse masking results in a large improvement in generalization for dodgeball, which has an underlying sparse dependent policy structure, but shows much smaller improvement in maze, whose underlying policy structure is relatively dense.

Figure 5: The normalized test return (mean and standard error over 10 seeds) of our method for different sparsity levels $\alpha$ in the dodgeball and maze environments. For more environments see Figure 8 in the appendix.

Panel of Figure 6 — Figure 6: The relative importance of each area of the observation according to the attention maps of policies learned from Attention, Input Masked Attention, and Sparse Masked Attention networks on the Dodgeball environment.

Panel of Figure 7 — Figure 7: The relative importance of each area of the observation according to the attention maps of policies learned from Attention, Input Masked Attention, and Sparse Masked Attention networks on the Maze environment.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Environment	CNN	Attention	Input Masked Attention	Sparse Masked Attention (ours)
bigfish	2.35 $\pm$ 0.35	11.32 $\pm$ 2.14	11.96 $\pm$ 1.16	21.61 $\pm$ 0.87
fruitbot	27.68 $\pm$ 0.28	30.07 $\pm$ 0.21	29.61 $\pm$ 0.18	30.74 $\pm$ 0.20
coinrun	8.92 $\pm$ 0.06	8.42 $\pm$ 0.15	8.21 $\pm$ 0.12	8.84 $\pm$ 0.10
dodgeball	2.19 $\pm$ 0.12	3.85 $\pm$ 0.69	4.86 $\pm$ 0.38	10.95 $\pm$ 0.30
heist	1.98 $\pm$ 0.09	2.62 $\pm$ 0.23	2.74 $\pm$ 0.16	4.12 $\pm$ 0.23
starpilot	42.56 $\pm$ 0.57	36.69 $\pm$ 1.21	38.43 $\pm$ 1.32	46.38 $\pm$ 0.89

Appendix

Table 1: Returns in unseen test environments (mean and standard error over 10 seeds) on various Procgen environments for our method (Sparse Masked Attention) and the baselines we compare against.

Environment	CNN	Attention	Input Masked Attention	Sparse Masked Attention (ours)
bigfish	9.73 $\pm$ 0.97	14.42 $\pm$ 2.64	15.75 $\pm$ 1.41	26.51 $\pm$ 1.31
fruitbot	31.15 $\pm$ 0.09	31.45 $\pm$ 0.14	31.46 $\pm$ 0.10	32.49 $\pm$ 0.13
coinrun	9.78 $\pm$ 0.04	9.68 $\pm$ 0.10	9.32 $\pm$ 0.11	9.51 $\pm$ 0.11
dodgeball	6.90 $\pm$ 0.16	5.33 $\pm$ 0.79	6.81 $\pm$ 0.51	13.24 $\pm$ 0.36
heist	7.41 $\pm$ 0.10	6.46 $\pm$ 0.42	5.51 $\pm$ 0.22	8.36 $\pm$ 0.08
starpilot	44.68 $\pm$ 0.59	37.27 $\pm$ 1.04	38.15 $\pm$ 1.03	48.61 $\pm$ 1.12

Appendix

Table 2: Returns in training environments (mean and standard error over 10 seeds) on various Procgen environments for our method (Sparse Masked Attention) and the baselines we compare against.

Figure 8: The normalized return (mean and standard error over 10 seeds) on unseen test tasks for Sparse Masked Attention policies at different sparsity levels $\alpha$ in each Procgen environment.

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,727 papers

Overcoming Valid Action Suppression in Unmasked Policy Gradient Algorithms
Mar 10, 2026Renos Zabounidis, Roy Siegelmann, Mohamad Qadri +3Policy Gradient Methods
Carnegie Mellon University · Massachusetts Institute of Technology · Virginia Tech, Department of Mechanical Engineering
TransMASK: Masked State Representation through Learned Transformation
Mar 5, 2026Sagar Parekh, Preston Culbertson, Dylan P. LoseyRepresentation Learning Robot Imitation Learning
Mechanical Engineering Department, Virginia Tech · Computer Science Department, Cornell University
Mask-Aware Policy Gradients for Diffusion Language Models
Jul 16, 2026Haran Raajesh, Kulin Shah, Adam Klivans +1Masked Diffusion Models RL for Language Model Reasoning
Department of Computer Science The University of Texas at Austin

Mar 10, 2026J/K move · Enter open · S save

Overcoming Valid Action Suppression in Unmasked Policy Gradient Algorithms

Renos Zabounidis, Roy Siegelmann, Mohamad Qadri, Woojun Kim+2

Carnegie Mellon University · Massachusetts Institute of Technology · Virginia Tech, Department of Mechanical Engineering

Policy Gradient Methods

Open paper AskPDF arXiv

In reinforcement learning environments with state-dependent action validity, action masking consistently outperforms penalty-based handling of invalid actions, yet existing theory only shows that masking preserves the policy gradient theorem. We identify a distinct failure mode of unmasked training: it systematically suppresses valid actions at states the agent has not yet visited. This occurs because gradients pushing down invalid actions at visited states propagate through shared network parameters to unvisited states where those actions are valid. We prove that for softmax policies with shared features, when an action is invalid at visited states but valid at an unvisited state

s^*

, the probability

π(a \mid s^*)

is bounded by exponential decay due to parameter sharing and the zero-sum identity of softmax logits. This bound reveals that entropy regularization trades off between protecting valid actions and sample efficiency, a tradeoff that masking eliminates. We validate empirically that deep networks exhibit the feature alignment condition required for suppression, and experiments on Craftax, Craftax-Classic, and MiniHack confirm the predicted exponential suppression and demonstrate that feasibility classification enables deployment without oracle masks.

Abstract

Reinforcement Learning

Explore similar work

Overcoming Valid Action Suppression in Unmasked Policy Gradient Algorithms

TransMASK: Masked State Representation through Learned Transformation

Mask-Aware Policy Gradients for Diffusion Language Models