# Generative Modelling of Stochastic Actions with Arbitrary Constraints in Reinforcement Learning

arXiv: https://arxiv.org/abs/2311.15341 · HTML: https://mlsanity.com/paper/2311.15341/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 35 references that are in the map.
- First posted on arXiv 26 Nov 2023.

## Builds on (most cited first, 27 of 27)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions](https://mlsanity.com/paper/2102.05379/) (2021)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [PettingZoo: Gym for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2009.14471/) (2020)
- [DC3: A learning method for optimization with hard constraints](https://mlsanity.com/paper/2104.12225/) (2021)
- [A Review of Safe Reinforcement Learning: Methods, Theory and Applications](https://mlsanity.com/paper/2205.10330/) (2022)
- [A Closer Look at Invalid Action Masking in Policy Gradient Algorithms](https://mlsanity.com/paper/2006.14171/) (2020)
- [Deep Reinforcement Learning in Large Discrete Action Spaces](https://mlsanity.com/paper/1512.07679/) (2015)
- [OptLayer - Practical Constrained Optimization for Deep Reinforcement Learning in the Real World](https://mlsanity.com/paper/1709.07643/) (2017)
- [Discretizing Continuous Action Space for On-Policy Optimization](https://mlsanity.com/paper/1901.10500/) (2019)
- [Leveraging exploration in off-policy algorithms via normalizing flows](https://mlsanity.com/paper/1905.06893/) (2019)
- [Discrete and Continuous Action Representation for Practical RL in Video Games](https://mlsanity.com/paper/1912.11077/) (2019)
- [Improving Exploration in Soft-Actor-Critic with Normalizing Flows Policies](https://mlsanity.com/paper/1906.02771/) (2019)
- [Value constrained model-free continuous control](https://mlsanity.com/paper/1902.04623/) (2019)
- [Resource Constrained Deep Reinforcement Learning](https://mlsanity.com/paper/1812.00600/) (2018)
- [Stochastic Gradient Descent with Biased but Consistent Gradient Estimators](https://mlsanity.com/paper/1807.11880/) (2018)
- [Escaping from Zero Gradient: Revisiting Action-Constrained Reinforcement Learning via Frank-Wolfe Policy Optimization](https://mlsanity.com/paper/2102.11055/) (2021)
- [Solving Online Threat Screening Games using Constrained Action Space Reinforcement Learning](https://mlsanity.com/paper/1911.08799/) (2019)
- [A Review of Incident Prediction, Resource Allocation, and Dispatch Models for Emergency Management](https://mlsanity.com/paper/2006.04200/) (2020)
- [Learning Stable Normalizing-Flow Control for Robotic Manipulation](https://mlsanity.com/paper/2011.00072/) (2020)
- [CLAIM: Curriculum Learning Policy for Influence Maximization in Unknown Social Networks](https://mlsanity.com/paper/2107.03603/) (2021)
- [Efficient Entropy for Policy Gradient with Multidimensional Action Space](https://mlsanity.com/paper/1806.00589/) (2018)
- [The Mysteries of Security Games: Equilibrium Computation Becomes Combinatorial Algorithm Design](https://mlsanity.com/paper/1603.02377/) (2016)

## Built on by (most cited first, 10 of 10)

- [Efficient Action-Constrained Reinforcement Learning via Acceptance-Rejection Method and Augmented MDPs](https://mlsanity.com/paper/2503.12932/) (2025)
- [An Extended Benchmarking of Multi-Agent Reinforcement Learning Algorithms in Complex Fully Cooperative Tasks](https://mlsanity.com/paper/2502.04773/) (2025)
- [Discretizing Continuous Action Space with Unimodal Probability Distributions for On-Policy Reinforcement Learning](https://mlsanity.com/paper/2408.00309/) (2024)
- [BraVE: Offline Reinforcement Learning for Discrete Combinatorial Action Spaces](https://mlsanity.com/paper/2410.21151/) (2024)
- [Leveraging Constraint Violation Signals For Action-Constrained Reinforcement Learning](https://mlsanity.com/paper/2502.10431/) (2025)
- [Action-Constrained Imitation Learning](https://mlsanity.com/paper/2508.14379/) (2025)
- [Towards Interpretable Reinforcement Learning with Constrained Normalizing Flow Policies](https://mlsanity.com/paper/2405.01198/) (2024)
- [SAINT: Attention-Based Policies for Discrete Combinatorial Action Spaces](https://mlsanity.com/paper/2505.12109/) (2025)
- [Improving and Accelerating Offline RL in Large Discrete Action Spaces with Structured Policy Initialization](https://mlsanity.com/paper/2601.04441/) (2026)
- [APC-RL: Exceeding Data-Driven Behavior Priors with Adaptive Policy Composition](https://mlsanity.com/paper/2601.19452/) (2026)

## Abstract (arXiv)

Many problems in Reinforcement Learning (RL) seek an optimal policy with large discrete multidimensional yet unordered action spaces; these include problems in randomized allocation of resources such as placements of multiple security resources and emergency response units, etc. A challenge in this setting is that the underlying action space is categorical (discrete and unordered) and large, for which existing RL methods do not perform well. Moreover, these problems require validity of the realized action (allocation); this validity constraint is often difficult to express compactly in a closed mathematical form. The allocation nature of the problem also prefers stochastic optimal policies, if one exists. In this work, we address these challenges by (1) applying a (state) conditional normalizing flow to compactly represent the stochastic policy -- the compactness arises due to the network only producing one sampled action and the corresponding log probability of the action, which is then used by an actor-critic method; and (2) employing an invalid action rejection method (via a valid action oracle) to update the base policy. The action rejection is enabled by a modified policy gradient that we derive. Finally, we conduct extensive experiments to show the scalability of our approach compared to prior methods and the ability to enforce arbitrary state-conditional constraints on the support of the distribution of actions in any state.
