A Closer Look at Invalid Action Masking in Policy Gradient Algorithms

Shengyi Huang; Santiago Ontañón

Source-linked AI summary

A Closer Look at Invalid Action Masking in Policy Gradient Algorithms

Shengyi Huang, Santiago Ontañón

arXiv:2006.14171v3Jun 25, 2020cs.LGcs.AIstat.ML

TL;DR

Invalid action masking is widely used in policy-gradient reinforcement learning, but its theoretical basis and empirical effects remain under-investigated. This paper analyzes its policy gradient and masking regimes, finding that masking scales as invalid actions increase while invalid-action penalties struggle.

Problem
Theoretical foundations and empirical effects of invalid action masking in policy-gradient algorithms remain under-investigated despite its use in games with state-dependent valid actions.
Method
The paper theoretically analyzes masked policy gradients and empirically compares invalid action masking, invalid-action penalties, naive masking, and removing masks after training.
Results
Invalid action masking scales well as invalid actions increase, with tsolve roughly 12% across map sizes and tfirst about 0.05–0.08% of training steps.
Takeaways & Limitations
Invalid action masking produces a valid policy gradient, scales better than invalid-action penalties, and retains some useful behavior when masking is removed.

Abstract

from arXiv · show

In recent years, Deep Reinforcement Learning (DRL) algorithms have achieved state-of-the-art performance in many challenging strategy games. Because these games have complicated rules, an action sampled from the full discrete action distribution predicted by the learned policy is likely to be invalid according to the game rules (e.g., walking into a wall). The usual approach to deal with this problem in policy gradient algorithms is to "mask out" invalid actions and just sample from the set of valid actions. The implications of this process, however, remain under-investigated. In this paper, we 1) show theoretical justification for such a practice, 2) empirically demonstrate its importance as the space of invalid actions grows, and 3) provide further insights by evaluating different action masking regimes, such as removing masking after an agent has been trained using masking. The source code can be found at https://github.com/vwxyzjn/invalid-action-masking

Introduction

DRL agents face state-dependent valid action spaces and extremely large full action spaces where sampled actions are often invalid. The paper studies invalid action masking as a way to sample only valid actions and compares it with penalizing invalid actions, finding masking scales better as invalid actions grow.

Motivation: Valid discrete action spaces vary in size across states because challenging games have complicated rules.One state may have 5 valid actions while another has 7.
Motivation: Dota 2 illustrates the scale problem with a full discrete action space of 1,837,080 actions.In such large spaces, sampled actions will typically be invalid.
Invalid action masking: Invalid action masking masks out invalid actions and samples only from actions valid in the current state.The approach avoids repeatedly sampling invalid actions in full discrete action spaces.
Empirical comparison: When the space of invalid actions grows, invalid action masking scales well and solves the desired task, whereas invalid action penalty struggles to explore the first reward.Invalid action penalty gives negative rewards for invalid actions so the agent learns to avoid them.

Background

The background formulates reinforcement learning as an MDP with a stochastic policy and explains that policy gradient methods maximize expected discounted reward through gradient ascent.

MDP formulation: Policy gradient methods train agents in a Markov Decision Process with state space S, discrete action space A, transitions P, rewards r, discount factor γ, and horizon T.The stochastic policy πθ maps states and actions to probabilities, while trajectories begin from ρ0 and follow the transition and policy dynamics.
Policy gradients: The core objective is the expected discounted return J, whose policy gradient ∇θJ measures how policy parameters θ should change to improve reward.The paper describes gradient ascent, θ = θ + ∇θJ, as maximizing expected discounted reward.

Invalid Action Masking

Invalid action masking is theoretically valid in policy-gradient algorithms: it applies a state-dependent differentiable transformation to policy logits, yielding valid policy-gradient updates while eliminating gradient contributions from invalid actions.

Implementation: Invalid action masking replaces invalid-action logits with a large negative number and renormalizes the remaining logits into a valid action distribution.The example uses M = −1 × 10^8, making the invalid action’s probability virtually zero when M is sufficiently negative.
Implementation: Masking makes the gradient corresponding to invalid-action logits zero, rather than merely renormalizing the probability distribution.This distinguishes invalid action masking from simple probability redistribution among valid actions.
Theoretical justification: Proposition 1 establishes that the invalid-action policy gradient is the policy gradient of the masked policy π′θ.The masking process is modeled as a state-dependent differentiable function applied to the logits produced by πθ.
Theoretical justification: Because masking is identity for valid logits and constant for invalid logits, π′θ is differentiable and satisfies the policy-gradient theorem’s assumptions.The mask depends on the state, so different states can apply different transformations to the same logit vector.

Experimental Setup

Experiments use µRTS2, a minimalistic real-time strategy testbed whose combinatorial action space produces many invalid actions. Agents are trained with PPO and evaluated on resource harvesting under several invalid-action handling strategies.

Testbed: µRTS2 preserves simultaneous durative actions, large branching factors, and real-time decision-making while making invalid actions grow with map size.The Source Unit and Attack Target Unit ranges grow linearly with map size, increasing the invalid-action space.
Environment: Observations are tensors of shape (h, w, nf) with 27 binary feature planes, while actions are 8-dimensional discrete vectors.Action components encode the selected source unit, action type, and parameters for that action type.
Task: Agents harvest resources for Player 1, receiving +1 when a worker harvests and another +1 when it returns the resource to base.Episodes last at most 200 time steps but may end earlier when all resources are harvested.
Invalid-action handling: The study trains agents with Proximal Policy Optimization and compares four strategies for handling invalid actions.The supplied setup introduces invalid-action penalties and masking as comparison regimes.
Invalid-action handling: Penalty experiments test rinvalid ∈ {0, −0.01, −0.1, −1}, applying a non-positive reward whenever the agent issues an invalid action.These values assess how different negative-reward scales affect learning.
Invalid-action handling: Masking restricts Source Unit and Attack Target Unit selection to valid units, but incorrect parameters can still produce invalid actions.The experiments omit a feature-complete mask because these two masked components already substantially reduce the action space.

2. Invalid action masking.

Invalid action masking scales across increasing map sizes, whereas invalid action penalties fail to scale and can discourage exploration. Naive masking achieves strong returns but causes exploding KL divergence, while removing masking after training retains some performance but degrades on larger maps.

Results: Invalid action masking scales well as map size increases, with tsolve roughly 12% across sizes and tfirst consistently about 0.05–0.08% of training steps.The agent finds its first reward quickly regardless of map size.
Results: Invalid action penalties perform well on 4 × 4 maps but fail to scale as the invalid-action space grows, sometimes struggling to discover the first reward.With rinvalid = −0.01 on 10 × 10 maps, agents spent 3.43% of training discovering the first reward.
Results: Setting rinvalid = −1 adversely affects exploration and produces consistently the worst performance across maps.The passage attributes this effect to discouraging exploration.
Results: Naive invalid action masking attains the best repisode across almost all maps, but its average Kullback–Leibler divergence explodes.It samples from the re-normalized masked distribution while updating gradients with the unmasked probability.
Results: Removing the mask at evaluation preserves some performance, but larger maps cause degradation and more invalid actions, especially invalid Source Unit selections.These agents nevertheless perform significantly better than agents trained with invalid action penalties and evaluated without masking.

Related Work

Prior work addresses invalid actions through continuous-space nearest-neighbor selection, action elimination, and broader action-space shaping. Existing ablations indicate masking can matter, but had not examined its empirical effects as invalid actions grow, which this paper studies.

Related Work: Alternative approaches embed discrete actions into continuous space and use nearest-neighbor methods to identify valid actions, or train an Action Elimination Network to reduce the action set.These approaches are attributed to Dulac-Arnold et al. (2015) and Zahavy et al. (2018), respectively.
Related Work: “Action space shaping” simplifies the full discrete action space through action removals or discretization of continuous actions.Kanervisto et al. (2020) uses this term for work that includes removing non-useful actions, such as “sneak” in Minecraft RL.
Related Work: Prior ablations suggest invalid action masking can affect agent performance, but do not measure its empirical effect as the invalid-action space grows.This growing-invalid-action-space analysis is identified as the focus addressed by the paper.

Conclusions

The paper concludes that invalid action masking yields a valid policy gradient, operates through a state-dependent differentiable function, and scales empirically as invalid-action spaces grow. Its demonstrated effectiveness motivates wider adoption to improve learning efficiency in DRL games with large, complex discrete action spaces.

Invalid action masking produces a valid policy gradient.The paper identifies this as a central theoretical result.
Invalid action masking applies a state-dependent differentiable function when calculating the action probability distribution.
Invalid action masking empirically scales well as the space of invalid actions grows.
The paper advocates wider adoption of invalid action masking because it enables agents to learn more efficiently.

Appendices

The appendices detail the factorized MultiDiscrete PPO implementation, implementation-level training optimizations, action-timing considerations, and reproducibility materials. Together, these details specify how the experiments reduce action-space computation and standardize training and evaluation.

MultiDiscrete PPO implementation: The factorized MultiDiscrete policy generates 2hw + 36 logits instead of 9216(hw)^2 logits for the full µRTS action space.Each action component is sampled from logits generated for its own range, avoiding direct distribution generation over all possible discrete actions.
Action timing: µRTS harvest, return, and move actions each take 10 game ticks, during which the acting unit is marked busy and cannot execute another action.The passage introduces this timing behavior as a consideration for preventing repeated action issuance to busy units.