Overcoming Valid Action Suppression in Unmasked Policy Gradient Algorithms

Renos Zabounidis; Roy Siegelmann; Mohamad Qadri; Woojun Kim; Simon Stepputtis; Katia P. Sycara

cs.LGMar 10, 2026

Overcoming Valid Action Suppression in Unmasked Policy Gradient Algorithms

Authors: Renos Zabounidis, Roy Siegelmann, Mohamad Qadri, Woojun Kim, Simon Stepputtis, Katia P. Sycara

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

Abstract

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.

Figures & tables

Figure 1: Suppression mechanism illustrated on a staircase corridor. (a) The agent must traverse a corridor to reach a staircase at $s^{*}$ , where descend is the only valid goal-reaching action. At $T\!=\!0$ , all actions are equally likely ( $1/n$ ). After $N$ training steps, gradient updates at visited states reinforce right and suppress all other actions, including descend . Because parameters are shared, this suppression propagates to $s^{*}$ before the agent arrives. The descend action is suppressed at the one state where it is needed. (b) Upper bound on $\pi(\textsf{descend}\mid s^{*})$ from Theorem 1 . Without entropy regularization ( $\beta\!=\!0$ ), the probability decays exponentially. With entropy regularization ( $\beta\!>\!0$ ), a floor emerges but cannot eliminate suppression (Eq. 7 ).

Figure 2: Architecture overview. A shared encoder $\phi(s)$ feeds three heads: a classification head (blue) predicting action validity $\hat{\nu}(s,a)$ via sigmoid, a policy head (green) producing oracle-masked actions $\pi^{\text{oracle}}_{\theta}(a\mid s)$ , and a value head (orange) estimating $V(s)$ . The predicted validity constructs a predicted policy $\pi^{\text{pred}}_{\theta}(a\mid s)$ . The KL divergence between $\pi^{\text{oracle}}$ and $\pi^{\text{pred}}$ yields per-action weights $w_{a}$ for the classification loss (Eq. 12 ). The total loss combines PPO and weighted classification objectives (Eq. 9 ). Dashed arrows indicate loss aggregation. At deployment, there are two modes: (i) with oracle masks available, classification heads can be discarded; (ii) without oracle masks, the learned predictor provides validity estimates for deployment.

Figure 3: Action suppression in Craftax and MiniHack Corridor-5 (PPO-Hybrid). Left: probability of rare critical actions at valid states (log scale). Right: fraction of timesteps with valid actions selected. Unmasked training (C2, red) exhibits exponential suppression, while oracle masking (C1, blue) prevents collapse. Dashed line: uniform initialization.

Figure 4: Representational correlation between valid and invalid states (PPO-Hybrid). Oracle masking preserves highly entangled representations, while KL-balanced classification induces validity-aware features without reintroducing policy-level suppression.

PPO-MLP
	Craftax		Classic		Corridor5
	GT	Pred	GT	Pred	GT	Pred	Frames $\to$ 1.0 (%)	Ep. Len
C2: Unmasked	26.9 $\pm$ 0.5	26.9 $\pm$ 0.5	19.2 $\pm$ 0.0	19.2 $\pm$ 0.0	0.9 $\pm$ 0.0	0.9 $\pm$ 0.0	—	161.3 $\pm$ 20.2
C1: Masked	30.4 $\pm$ 0.2	-0.9 $\pm$ 0.0	19.3 $\pm$ 0.0	-0.8 $\pm$ 0.0	0.9 $\pm$ 0.0	0.1 $\pm$ 0.0	—	118.3 $\pm$ 8.6
C3: Masked + Focal	24.6 $\pm$ 2.9	13.7 $\pm$ 1.5	19.4 $\pm$ 0.0	15.3 $\pm$ 0.6	0.9 $\pm$ 0.0	0.9 $\pm$ 0.0	—	113.6 $\pm$ 8.1
C4: Masked + KL	28.7 $\pm$ 3.2	28.6 $\pm$ 3.2	19.5 $\pm$ 0.0	19.3 $\pm$ 0.1	0.9 $\pm$ 0.0	0.9 $\pm$ 0.0	—	123.6 $\pm$ 6.6

Table 1: Evaluation with ground-truth (GT) versus predicted (Pred) masks. Masked-only agents fail when oracle masks are removed. KL-balanced classification enables near-oracle performance under predicted masks across architectures. Results report mean

\pm

std over 4 seeds (10,240 episodes per seed); bold indicates best non-baseline per column.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Environments used in our experiments.

Component	Dims	Normalization
Map	8217	Flattened 9 $\times$ 11 $\times$ 83, darkness masked
Inventory	16	$\sqrt{\text{count}}/10$ for materials; level/4 for tools
Potions	6	$\sqrt{\text{count}}/10$
Intrinsics	9	value/10
Direction	4	One-hot
Armour	8	level/2 for pieces; raw int for enchantments

Appendix

Table 2: Craftax-Full observation components.

Ore	Overworld	Gnomish	Troll	Fire	Ice
Coal	0.03	0.04	0.04	0.05	0.0
Iron	0.02	0.02	0.03	0.0	0.0
Diamond	0.001	0.005	0.01	0.0	0.005
Sapphire	0.0	0.0025	0.01	0.0	0.02
Ruby	0.0	0.0025	0.01	0.025	0.0

Appendix

Table 3: Ore spawn probabilities by floor.

Environment	$\|\mathcal{A}\|$	Valid/state	Validity ratio	Observation	Steps
Craftax-Full	43	5–8	12–19%	Symbolic (8268-d)	$10^{9}$
Craftax-Classic	17	5–8	29–47%	Symbolic (1345-d)	$10^{9}$
Corridor-2/3/5	11	3–5	27–45%	GlyphNet (256-d)	$2\times 10^{8}$

Appendix

Table 4: Environment summary.

Hyperparameter	PPO-Hybrid	PPO-RNN	PPO-MLP
Parallel environments	1024	1024	1024
Steps per rollout	128 ∗	64 ∗	64 ∗
Batch size	131,072	65,536	65,536
Number of minibatches	8	8	8
Minibatch size	16,384	8,192	8,192
Update epochs	4	4	4

Appendix

Table 5: PPO hyperparameters. ∗ indicates architecture-specific values.

Figure 6: Mask predictor ablation on Craftax (PPO-Hybrid,

\lambda=20

). Left: Episode return during training. Center-left: Classification accuracy. Center-right: Return with ground-truth masks. Right: Return with predicted masks.

Figure 7: Classification coefficient sweep with oracle masking and KL-balanced loss on Craftax (PPO-Hybrid). Left: Episode return. Center-left: Classification accuracy. Center-right: Classification loss. Right: Final return evaluated with ground-truth masks (dark) and predicted masks (light).

Figure 8: Suppression profile of the descend action during early training (PPO-Hybrid,

\lambda=20

). Left: probability at valid states (log scale). Center: probability at invalid states (log scale). Right: Pearson correlation between feature representations of valid and invalid states.

Policy Gradient Methods

Policy gradient methods optimize stochastic policies by estimating how parameter changes affect expected reward.

See all 164 papers

Sparse Masked Attention Policies for Reliable Generalization
Feb 23, 2026Caroline Horsch, Laurens Engwegen, Max Weltevrede +2Reinforcement Learning Generalization in RL
Department of Intelligent Systems, Delft University of Technology, Delft, Netherlands.
Confessions of a Masked MDP
Oct 20, 2026Sebastian Reboul, Lorenzo Croissant, Hélène HalconruyReinforcement Learning Constrained RL
SAMOVAR, Télécom SudParis, Institut Polytechnique de Paris · Modal’X, Université Paris-Nanterre · Wiremind Cargo, AI & ML Group +1
Evaluating Model-Free Policy Optimization in Masked-Action Environments via an Exact Blackjack Oracle
Mar 19, 2026Kevin SongPolicy Optimization Reinforcement Learning
Department of Biomedical Engineering The University of Alabama at Birmingham

Feb 23, 2026J/K move · Enter open · S save

Sparse Masked Attention Policies for Reliable Generalization

Caroline Horsch, Laurens Engwegen, Max Weltevrede, Matthijs T. J. Spaan+1

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

Reinforcement Learning Generalization in RL Representation Learning for RL

Open paper AskPDF 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.

Abstract

Policy Gradient Methods

Explore similar work

Sparse Masked Attention Policies for Reliable Generalization

Confessions of a Masked MDP

Evaluating Model-Free Policy Optimization in Masked-Action Environments via an Exact Blackjack Oracle