Overcoming Valid Action Suppression in Unmasked Policy Gradient Algorithms

TL;DR

This paper identifies a failure mode in unmasked policy gradient algorithms where valid actions are suppressed due to shared parameters, and proposes a solution using feasibility classification to improve performance in environments with action validity constraints.

Contribution

The paper reveals a systematic suppression issue in unmasked policies caused by parameter sharing and proposes feasibility classification as a practical remedy.

Findings

Unmasked training suppresses valid actions at unvisited states.

Feasibility classification enables deployment without oracle masks.

Empirical results confirm exponential suppression and effectiveness of the proposed method.

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 $\pi(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.