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

Feb 23, 2026

Sparse Masked Attention Policies for Reliable Generalization

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.

Caroline Horsch, Laurens Engwegen, Max Weltevrede +2

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

Reinforcement Learning Generalization in RL Representation Learning for RL

Ask PDF

Oct 20, 2026

Confessions of a Masked MDP

In reinforcement learning, one presumes the action set to be known ahead of time, and, thus, one must know if any state-action pairs are not permissible ahead of computations. In some applications, determining this action mask is the main computational bottleneck. Isolating this problem to highlight its core properties, we study how to combine the search for an optimal policy with the learning of the mask itself, under a linear MDP model in which each attempted action reveals the exact value of a linear score determining its validity. In this framework, we first construct an algorithm based on deterministic optimistic optimization and show that it is minimax optimal, which we show is also computationally efficient. However, it is sub-optimal on non-degenerate instances, as we show. We propose an alternative, which is also computationally efficient, which plans its exploration strategy. This algorithm admits a formal interpretation as an analogue of information-directed-sampling for deterministic feedback, which may be of independent interest. Finally, we study an application of this framework to transfer learning, where the invalid action set changes over time.

Sebastian Reboul, Lorenzo Croissant, Hélène Halconruy

SAMOVAR, Télécom SudParis, Institut Polytechnique de Paris · Modal’X, Université Paris-Nanterre · Wiremind Cargo, AI & ML Group +1

Reinforcement Learning Constrained RL Linear Markov Decision Processes

Ask PDF

Mar 19, 2026

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

Infinite-shoe casino blackjack provides a rigorous, exactly verifiable benchmark for discrete stochastic control under dynamically masked actions. Under a fixed Vegas-style ruleset (S17, 3:2 payout, dealer peek, double on any two, double after split, resplit to four), an exact dynamic programming (DP) oracle was derived over 4,600 canonical decision cells. This oracle yielded ground-truth action values, optimal policy labels, and a theoretical expected value (EV) of -0.00161 per hand. To evaluate sample-efficient policy recovery, three model-free optimizers were trained via simulated interaction: masked REINFORCE with a per-cell exponential moving average baseline, simultaneous perturbation stochastic approximation (SPSA), and the cross-entropy method (CEM). REINFORCE was the most sample-efficient, achieving a 46.37% action-match rate and an EV of -0.04688 after 10^6 hands, outperforming CEM (39.46%, 7.5x10^6 evaluations) and SPSA (38.63%, 4.8x10^6 evaluations). However, all methods exhibited substantial cell-conditional regret, indicating persistent policy-level errors despite smooth reward convergence. This gap shows that tabular environments with severe state-visitation sparsity and dynamic action masking remain challenging, while aggregate reward curves can obscure critical local failures. As a negative control, it was proven and empirically confirmed that under i.i.d. draws without counting, optimal bet sizing collapses to the table minimum. In addition, larger wagers strictly increased volatility and ruin without improving expectation. These results highlight the need for exact oracles and negative controls to avoid mistaking stochastic variability for genuine algorithmic performance.

Kevin Song

Department of Biomedical Engineering The University of Alabama at Birmingham

Policy Optimization Reinforcement Learning Model-Free RL Sample-Efficient RL

Ask PDF

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