A Closer Look at Invalid Action Masking in Policy Gradient Algorithms
Shengyi Huang, Santiago Ontañón
cs.LG, cs.AI, stat.ML
2020-06-25
Masking invalid logits with a large negative still yields a valid policy gradient; on μRTS it hits return 40, while penalties miss the first reward as maps grow.
Games such as StarCraft and Dota have legal actions that change with the state. Engineering practice unions every state's actions into one fixed discrete set, then samples from it. Dota 2's full set has 1,837,080 entries; a random draw is often an unaffordable item, a walk into a wall, or an enemy unit. The usual fix is to mask invalid actions and sample only from the legal set. The paper treats this as more than an implementation trick: it needs a policy-gradient justification, and a measurement of how the gap grows as invalid actions explode.
The folk alternative is a penalty on invalid actions so the policy learns not to pick them. In a huge action space that signal drowns exploration.
The policy emits logits, then a softmax. Masking replaces invalid logits with a large negative M (for example -1e8). After softmax those coordinates have near-zero probability and near-zero gradient. The key proposition: for each state the mask is either identity or a constant, both differentiable, so the masked policy π' still meets the policy gradient theorem. Updates use log π', not the unmasked log-probabilities.
Four PPO agents play a μRTS harvest task: +1 for mining, +1 for returning to base, about 40 at ceiling, 200-step cap. Actions are 8-way MultiDiscrete; source-unit and attack-target ranges grow with the square of map side. Maps are 4×4, 10×10, 16×16, 24×24. The four treatments:
The mask is incomplete. It covers source unit and attack target only; action-type parameters can still be illegal.
Proper masking hits episode return 40 on every map. The first positive reward arrives at 0.05%–0.08% of training, and solving takes about 9%–18% of the budget. Penalties still work on 4×4 (return 40 at rinvalid=-0.01) and collapse to 0–1 on 10×10 and larger, sometimes spending several percent of training on the first +1. rinvalid=-1 suppresses exploration on the small map and is consistently worst.
| Strategy | Map | Episode return | First reward (% of steps) |
| Masking | 4×4 to 24×24 | 40.00 | 0.05–0.08% |
| Penalty -0.01 | 4×4 | 40.00 | 0.51% |
| Penalty -0.01 | 10×10 | 0.50 | 1.57% |
| Penalty -0.01 | 24×24 | 0.50 | 1.92% |
| Naive masking | 24×24 | 38.50 | 0.07% (but 49% to solve) |
Naive masking can score higher, 59.61 on 4×4, because the agent learns to harvest the far side of the map. PPO's KL between successive policies blows up, and 24×24 takes 49% of training to solve, a bad sign for harder tasks. After training with a mask and stripping it at test time, return is 33.53 on 4×4 and 17.37 on 24×24, still well above penalties. Some of the "do not pick illegal units" behavior is in the weights; large maps do not keep it.
This turns "mask illegal actions" from folklore into a policy-gradient fact, with a scaling picture as the illegal set grows. For large discrete actions, masking should be the default. Do not expect a penalty to teach legality. Sampling and the update must share the masked distribution; masking only at sample time looks like learning and then detonates KL.
The task is harvesting, not a full match. Read the result as evidence about exploration efficiency, not RTS strength.
The mask covers two action heads, not the full rulebook. The environment skips 9 frames, so timing differs from a real RTS. Four seeds per cell, so the curves are noisy. Performance drops hard on large maps once the mask is removed; if a rules engine cannot supply a mask at deployment, do not assume the policy still works. μRTS legality is cheap to compute. The paper does not say where the mask comes from when the action space is messier.