SAINT: Attention-Based Policies for Discrete Combinatorial Action Spaces

TL;DR

SAINT addresses the challenge of large discrete combinatorial action spaces by treating joint actions as unordered sets and applying state-conditioned self-attention to model sub-action interactions. The method yields a permutation-invariant, tractable policy with parallel per-sub-action decoding, and is compatible with PPO, A2C, and offline RL objectives. Empirical results across CityFlow, CoNE, and discretized MuJoCo demonstrate that SAINT outperforms factorized, autoregressive, Wol-DDPG, and flat baselines, scaling to environments with up to approximately $1.7 \times 10^{7}$ joint actions. The work shows improved sample efficiency and robustness to dimensionality and dependency strength, and points to future directions such as lighter attention variants and dynamic action sets for even broader applicability.

Abstract

The combinatorial structure of many real-world action spaces leads to exponential growth in the number of possible actions, limiting the effectiveness of conventional reinforcement learning algorithms. Recent approaches for combinatorial action spaces impose factorized or sequential structures over sub-actions, failing to capture complex joint behavior. We introduce the Sub-Action Interaction Network using Transformers (SAINT), a novel policy architecture that represents multi-component actions as unordered sets and models their dependencies via self-attention conditioned on the global state. SAINT is permutation-invariant, sample-efficient, and compatible with standard policy optimization algorithms. In 20 distinct combinatorial environments across three task domains, including environments with nearly 17 million joint actions, SAINT consistently outperforms strong baselines.

Figures (10)

• Figure 1: Overview of SAINT. Initial sub-action embeddings are conditioned on the global state $\mathbf{s}$ to produce state-aware representations. Stacked Transformer blocks then model dependencies among sub-actions. The resulting context-aware representations are passed to independent Decision MLPs, which output per-sub-action policy distributions used for factorized sampling.
• Figure 2: Learning curves show that SAINT outperforms all baselines in both learning speed and final reward, across both the CityFlow-Linear and CityFlow-Irregular environments.
• Figure 3: Performance in discretized MuJoCo environments. While results are similar across methods in HalfCheetah, SAINT outperforms factorized and autoregressive baselines in Hopper and Walker2D, demonstrating its ability to handle complex action spaces even when sub-action dependencies are relatively weak.
• Figure 4: Visualizations of the two CityFlow traffic control configurations used in our experiments. CityFlow-Linear (Figure \ref{['fig:simple-CityFlow']}) has three intersections arranged in a row, yielding 729 possible joint actions. CityFlow-Irregular (Figure \ref{['fig:complex-CityFlow']}) has 375 joint actions but exhibits greater coordination demands and more diverse traffic interactions.
• Figure 5: Full learning curves for all baselines in CityFlow, including Wol-DDPG. Wol-DDPG performs poorly, consistent with its known limitations in environments with unordered sub-actions.