SAINT: Attention-Based Policies for Discrete Combinatorial Action Spaces

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 18 distinct combinatorial environments across three task domains, including environments with $1.35 \times 10^{18}$ possible actions, SAINT consistently outperforms strong baselines.

💡 Research Summary

The paper tackles the long‑standing challenge of reinforcement learning (RL) in discrete combinatorial action spaces, where the joint action is a Cartesian product of many sub‑actions and the number of possible actions grows exponentially. Traditional RL methods either treat the action space as a flat categorical distribution—becoming intractable for large combinatorial problems—or impose restrictive structural assumptions. Existing approaches fall into two main families: (1) factorized policies that assume conditional independence across sub‑actions (π(a|s)=∏π_i(a_i|s)), which cannot capture interactions, and (2) autoregressive policies that impose a fixed ordering on sub‑actions (π(a|s)=∏π_i(a_i|s,a_{<i})), which break permutation invariance and can suffer when the imposed order does not reflect the true dependency structure. Both families are inadequate for domains such as drug combination therapy, traffic signal control, or resource allocation where sub‑action indices are arbitrary and the essential structure lies in the interactions among sub‑actions.

The authors propose SAINT (Sub‑Action Interaction Network using Transformers), a novel policy architecture that treats a joint action as an unordered set of sub‑actions and models their dependencies with self‑attention conditioned on the global state. SAINT consists of three stages. First, each sub‑action i receives a learnable embedding e_i∈ℝ^d. A global state vector s is processed by a FiLM (Feature‑wise Linear Modulation) network g(s)→(γ,β), and the same (γ,β) are applied to every e_i, yielding state‑aware embeddings \tilde e_i=γ⊙e_i+β. This injects state information efficiently without increasing dimensionality. Second, the matrix of embeddings \tilde E∈ℝ^{A×d} passes through L Transformer blocks. Positional encodings are deliberately omitted so that the self‑attention operation is permutation‑equivariant: the output does not depend on the order of sub‑action tokens. Multi‑head attention allows each sub‑action to attend to all others, learning high‑order interaction patterns while preserving the set‑like nature of the input. Third, each resulting context‑aware token x_i is fed into a sub‑action‑specific decision MLP f_i, producing logits over the K_i discrete choices of that sub‑action. A softmax yields a categorical distribution π_i(a_i|s). Because each x_i already incorporates information from the rest of the set, the joint policy can be expressed as the product of these per‑sub‑action distributions without loss of expressive power, keeping inference tractable (O(A·d) rather than O(∏K_i)).

SAINT is compatible with any actor‑centric RL objective that maximizes the weighted log‑likelihood of sampled joint actions, such as PPO, A2C, IQL, or AWAC. The weight w_Φ(s,a) can be an advantage estimate, a Q‑value, or any non‑negative scalar. The architecture does not require a factorized critic; the critic only supplies the scalar weight. However, methods that need to evaluate expectations or maxima over the entire combinatorial action space (e.g., E_{a′∼π}