Generative Modelling of Stochastic Actions with Arbitrary Constraints in Reinforcement Learning
Authors: Changyu Chen
Ramesha Karunasena
Thanh Hong Nguyen
Arunesh Sinha
Pradeep Varakantham
Venue: NeurIPS 2023 poster (2023)
Keywords: Action constrained reinforcement learning
Normalizing flow
Generative modelling
TL;DR: The first work incorporating discrete flow in RL and learning in large categorical action space with state conditioned constraints.

Abstract: 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.

Generative Modelling of Stochastic Actions with Arbitrary Constraints in Reinforcement Learning
Chen Changyu' Ramesha Karunasena', Thanh Hong Nguyen', Arunesh Sinha', Pradeep Varakantham' Singapore Management University .University of Oregon', Rutgers University
cychen 2020@phdcs. smu edu.sg, sg, rameshak@smu. edu. edu, arunesh. sinhatrutgers. edu, setpeppdsmuuedu.....
thanhhng@cs.
uoregon. Boanhhon
Abstract
Many problems in Reinforcement Learning (RL) seek optimal policy with large discrete multidimensional yet unordered action spaces; these include problems in randomized allocation of resources such 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 exists. In this work, we address these challenges by (1) applying a (state) conditional normalizing flow compactly represent the stochastic policy the compactness arises due the network only producing one sampled action and the corresponding log probability 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 enabled by 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 cateeconnittionll constraints on the support of the distribution of actions in any state")
1 Introduction
Adaptive resource allocation problems with multiple resource types (e.g., fire trucks, ambulances, police vehicles) are ubiquitous in the real world [3
23 32] One example is allocating security
resources and emergency response vehicles to different areas depending on incidents (23]. There are many other similar problems in aggregation systems for mobility/tran punpiitt/ransspotattionn logistics etc [13] In this paper, we are interested in addressing the multiple difficult challenges present in such adaptive resource allocation problems: (a) Combinatorial action space, as the number of resource allocations is combinatorial; (b) Categorical action space, there is no ordering of resource allocations with respect to the overall objective (as increasing or decreasing different resource types at different locations can have different impact on overall objective): (c) Constraints on feasible allocations and switching between allocations; and (d) Finally, uncertainty in demand for resources.
Existing research in RI. for constrained action spaces has considered resource allocation problems with single type of resource, thereby introducing order in action space [. In such ordered action spaces, actions can be converted into continuous space and this allows for the usage of continuous action RL methods (c.g., DDPG). In problems of interest in this paper, we interested in problems
OOr implementation is available at attpps//jiiallccocccccccooooo/eo
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

with multiple resource types (e.g., fire truck, ambulance, police). These problems have a large action space that discrete and unordered (categorical) and there are constraints on feasible allocations (e.g., no two police vehicles can be more than 3 km away, cannot move vehicles too far away from time step to time step). Furthermore, such constraints are easy to validate by a validity oracle given any allocation action, but are hard to represent as mathematical constraints on the support of the distribution of actions (in each state) as they often require exponentially many inequalities [4] 35]. An important consideration in allocation problems is randomized (stochastic) allocation arising from issues of fair division of indivisible resources so that an allottee is not starved of resources forever Thus, aim to output stochastic optimal policies, if one exists.
Towards addressing such resource allocation problems at scale, we propose to employ generative policies in RL. Specifically, we propose a new approach that incorporates discrete normalizing flow policy in an actor-critic framework to explore and learn the aforementioned constrained, categorical and adaptive resource allocation problems. Prior RL approaches for large discrete multidimensional action space include ones that assume a factored action space with independent dimensions, which we call as the factored marginal approach, since independence implies that any joint distribution over actions can be represented as the product of marginal distributions over each action dimension. Other approaches convert the selection of actions in multiple dimensions into sequential selection approach. Both these approaches are fundamentally limited in expressivity, which we reveal in detail in our experiments. Next, we provide a formal description of the problem that we tackle.
represented by the tuple (S, A, P,r,7, bo). where an agent can be in any state 8t E S at a given time t. The agent takes an action at E A, causing the environment to transition to a new state 8141 with a probability P S x A x S + 1]. Subsequently, the agent receives a reward r : S x A H R. In the infinite-horizon setting, the discounted factor is 0 < 7 < 1. The distribution of the initial state is bo-
Problem Statement: A Markov Decision Process (MDP)
In our work, we focus on a categorical action space, A. Categorical action spaces consist of discrete, unordered actions with no inherent numerical relationship between them. We assume that for any state s, there is a set of valid actions, denoted by C(s) S A. There is an oracle to answer whether an action a E C(s). but the complex constraint over categorical space cannot be expressed succinctly using closed form mathematical formula. Note that the constraint is not the same for every state. Our objective is to learn stochastic policy, (-1ss), which generates a prohability distribution over actions for state s with support only over the valid actions C(s) in state s. We call the set of such stochastic policies as valid policies ITe given the per state constraint C. We aim to maximize the long-term expected rewards over valid policies, that is, max J ((T),, where J is as follows:
J(x) = Ea-bo (V(s; where V(s; = E | 7'r (8t, 0.) 1s0 8;
(1)
In addition, we also consider settings with partial observability where the agent observes or E O at time where the observation arises from the state st with probability O O x S + [o, 1]. In this case. the optimal policy is a stochastic policy, T-I1),, where h is the history of observations till current time. For partial observability, we consider an unconstrained setting, as the lack of knowledge of the true state results in uncertainty about which constraint to enforce, which diverges from the focus in this work but is an interesting future research direction to explore. Thus, with partial observability, we search over stochastic policies II that maximize the long term return, that is, max ( J(x) the same as stated in Equation but where the expectation is also over the observation probability distribution in addition to the standard transition and stochastic policy probability distributions.
We propose two key innovations address the problem above. First, we present a conditional Normalizing Flow-based (6) Policy Network, which leverages Argmax Flow [16] to create a minimal representation of the policy for policy gradient algorithms. To the best of our knowledge, this is the first use of discrete normalizing flow in RL. Second, we demonstrate how to train the flow policies within the A2C framework. In particular, we need an estimate of the log probability of the action sampled from the stochastic policy but Argmax Flow provides only a biased lower bound via the evidence lower bound (ELBO). Thus, we design an effective sandwich estimator for the log probability that is sandwiched between the ELBO lower bound and an upper bound based on x2 divergence. Third, we propose a policy gradient approach which able to reject invalid actions (that do not satisfy constraints) referred to as Invalid Action Rejection Advantage (IAR-A2C). IAR-A2G queries the constraint oracle and ensures validity of actions in every state (in fully observable setting) by rejecting all invalid actions. We derive new policy gradient estimator
Contribution:

---

for IAR-A2C. Figure provides an overview of our architecture. Finally, our extensive experimental results reveal that our approach outperforms prior baselines in different environments and settings.
2 Background
Normalizing flows are a family of generative models that can provide both efficient sampling and density estimation. Their main idea to construct a series of invertible and differentiable mappings that allow transforming a simple probability distribution into a more complex one. Given = Z I2d with densities pv and pz respectively. normalizing flows [26] aim to learn a bijective and differentiable transformation f Z + V. This deterministic transformation allows us to evaluate the density at any point e E V based on the density of = E 2, as follows:
Normalizing Flows:
dz v f()
pv(u) - pz(2) det
do
In this context, pz can be any density, though it is typically chosen as a standard Gaussian and is represented by neural network. Consequently, normalizing flows offer a powerful tractable framework for learning complex density functions. However, the density estimation presented in Equation is limited to continuous probability distributions. To enable the learning of probability mass functions (P) on categorical discrete data, such as natural language, Argmax Flow [16] proposed to apply the argmax operation the output of continuous flows. Let's consider e E R Dx and x E (I....,MM M)P. The argmax operation is interpreted as a surjective flow layer v + æ, which is deterministic in one direction (Ina = arg max, Ud. written compactly x arg max v) and stochastic in the other (t ~ q(-lz)). With this interpretation, the argmax operation can be considered probabilistic right-inverse in the latent variable model expressed by:
P(xx) = f ((zl)) P(xlv) 8(x argmax(v))
where argmax is applied in the last dimension of v. In this scenario, the density model p(v) is modeled using a normalizing flow. The learning process involves introducing a variational distribution q(vlz)), which models the probabilistic right-inverse for the argmax surjection, and optimizing the evidence lower bound (ELBO), which is the RHS of the following inequality:
log P(x) 2 Ee>g((1a [log P(æ](o) - log p(v) log [[v])=)] = Eeegg-la) (eoggtt))loog(()) (()gg(p)))))
log ((v[)=)] = C
The last equality holds under the constraint that the support of q(v|z) is enforced to be only over the region $ = (v E DRD x = arg max da which ensures that P(x = 1. From standard variational inference results, log P(x) L = (L(((())))) (p((aa) [:)) which also approaches 0 as the approximate posterior g(v|ar)) [x) comes closer to the true posterior p( () (z)) over the training time.
Variational inference involves proposing a family of approximating distributions and finding the family member that is closest to the posterior. Typically, the Kullback-Leibler (KL) divergence KL(q)) [p) is employed to measure closeness, where (v) x>) represents variational family. This approach yields ELBO of the evidence log ((x) as described above.
x² Upper Bound:
Instead of using KL divergence, the authors in g suggest an alternative of x² -dvieergeecce to measure the closeness. As result, they derived an upper bound of the evidence, known as CUBO: log P(x) s
(
log E (1-)
Similar the ELBO in Argmax Flow, CUBO can be further simplified
under the constraint that the support of q(-l (æ)) is restricted to the region S:
p(o) ((())a))) )²
log P(xr) S log Eu (se)
fx²²
(4)
Also, L,= log P(z) = log(1 + D2- (p(ux))) (v)x)))) hence the gap between the ELBO and CUBO approaches 0 as the approximate posterior q(- becomes closer to the true posterior p(-|z))
3 Flow-based Policy Gradient Algorithm with Invalid Action Rejection
We first present the Flow-based Policy Network, which leverages Argmax Flow to create a minimal representation of the policy for policy gradient algorithms, and then construct flow policies within
3

---

Flow Policy
Somple actions from Flow Policy
Reject invalid betions
Actions
Unitormly sample
icy
update
Inyalid Actióna
Action
Actions
Learning algorithm
Action
Invirerment
Observation
Our IAR-A2C framework. At each time step, initial batch of action samples, along with their log probabilities, are generated using the Flow Policy. Invalid actions from this batch are rejected using an oracle. A single action is then uniformly sampled from the remaining valid ones, and executed. This selected action and the valid action set are stored along with the resulting state and reward. This collective experience is subsequently utilized to update the Flow Policy.
Figure 1:
the A2C framework. Our exposition will present policies conditioned on state, but the framework works for partial observability also by using the sequence of past observations state. After this, for the fully observable setting only, we introduce our novel policy gradient algorithm called IAR-A2C that enforces state dependent constraints.
3.1 Flow-based Policy Network
In policy gradient algorithms for categorical actions, the standard approach to model the entire policy, denoted as TT r--1s), which allows us sample actions and obtain their corresponding probabili- ties. The policy is parameterized using neural network where the network generates logits for each action, and then converts these logits into probabilities. The size of the output layer matches that of the action space, which can be prohibitively large for resource allocation problems of interest in this paper. We posit that we can have better policy representation, it is sufficient to require samples from the support set of the policy, represented by ( ((|ss) and probability value π ( [s) > 0.
Based this observation, our first contribu- tion is a compact policy representation using Argmax Flow [16], which we refer as the Flow Policy. Argmax Flow is a state-of-the- art discrete normalizing flow model and has shown great capability of learning categorical data such as in sentence generation. context, Flow Policy will output the action (a sample of the flow policy) and its probabil- ity, instead of explicitly outputting the entire distribution and sampling from it as in prior work. Once trained, we can sample from the Argmax Flow and, more importantly, estimate the probability of the sample. A normaliz- ing flow model transforms a base distribution, given by random variable 2g, the desired distribution. In our approach, the desired distribution is the distribution given by policy æ. We adapt Argmax Flow approach for learning stochastic policy.
ATU
ang
P-3
3 argusa
Uo 0,
Mo(su)e)
(aa(Fale)
(a)a)
our
Figure 2: Composition of conditional flow (zx and argmax transformation resulting in the policy To (a|s). The flow maps from a base distribution P,(2o]e (s) by using a bijection F... The diagram is adapted from Hoogeboom al. [66]
Before diving into the specifics of the Flow Policy, we first discuss our rationale for choosing the normalizing flow model over other contemporary deep generative models, such as Generative Adversarial Networks (GANs) [2] and Diffusion Models [15] 29). GANs, categorized as implicit generative models, do allow for an estimation of the data's s log probability. While prior research has successfully devised an approach for exact log probability estimation with Diffusion Models, these

---

Algorithm 1: ELBO Optimization
Input: Invertible flow Fu Suck o ... fm,1, State encoder Epp. Posterior (p, rollout
Sample Stales = {{s(())] ~M
for E States do
4
for j 5 Leeebb__seep do
2(6) F ) where = Ep(s)
al)) argmax 25)
for
j m_elbo sitesssspsspp do
z'Gn) thresholdy (ubx)), where ((b6))) (o((abb)
((fm) F (z)(m))
2g
Take a gradient ascending step of ELBO w.r.t. p. w and vb 10
L N Ei (logpp(>) (1m) [s) Ex= log det
8lu.k (2(5
log q0 (in)
11
models encounter the problem of slow sampling and expensive probability estimation, necessitating the resolution of an ordinary differential equation [29]. In contrast, normalizing flow has low cost for sampling and probability evaluation in our case, we find the simple flow model is sufficient to learn good policy. So normalizing flow is natural choice to construct our policy.
We demonstrate how to construct the flow policy within the A2C framework. An illustration is first encode the state that is then used to
provided in Figure 2 To define the policy, [aass) define the base distribution 2g. In accordance with standard practice, we select 2p as a Gaussian distribution with parameters μa, defined by the state encoder state-dependent neural network): 20 (s)s) + (s) where ~ (V, 1). We write this as 2n ~ [s),) where denotes the weights of the state encoder. We then apply a series of invertible transformations given as functions fg that define the flow, followed by the argmax operation (as defined in background on Argmax flow). Consequently, the final sampled action is given by argmax (fK o .... fi (zo)). Each fi an invertible neural network (26) and we use F. = fw.k
fw. to denote the composed function, where w, is the weight of the network representing fk- Thus, sampled action a arginxx (F (20)) for E (s). We use shorthand 0 (p, w) when needed. Also, we use denote the output of fwe and (2(2)))) to denote the probability density of 2g ~ (Nup: ap).
Sw.((=x-1) to
We use a variational distribution de for the reverse mapping from a discrete action (conditional on state s) to continuous output of F. (zo). The corresponding estimate of the log probability of the sampled action a, denoted by i.. the evidence lower bound ELBO, computed as follows:
log logpp(zols) Σ K
0fk k=1 logg det dzx-1 )) log qu (2x (5)
(als) L Ezg
To ensure that our approximation closely approximates the evidence, optimize ELBO progressively, following the training scheme of Argmax flow, as shown in the subroutine Algorithm I Jused within the outer A2C framework Algorithm(2 The target distribution in this subroutine is given by actions a(5) sampled from the current policy (line 6); thus, this subroutine aims to update the flow network components (p, w) to make the distribution of overall policy be closer to this target distribution and improve the posterior qu estimate (ELBO update in line 11), which gets used in estimating the log probability of action required by the outer A2C framework (shown later in Algorithm 2). The various quantities needed for ELBO update are obtained in lines 8, 9 by passing a(7) back through the invertible flow layers (using qy for the discrete continuous inverse map).
3.2 Policy Gradient with Invalid Action Rejection
We aim to use the flow-based policy within an A2C framework, illustrated in Figure ! here we describe the same along with how to enforce arbitrary state based constraint on actions. We propose a sandwich estimator which combines ELBO and CUBO to obtain a low-bias estimation of log (als). thus improving the convergence of policy gradient learning. Moreover, to tackle the challenge of aforementioned complex constraints, our algorithm proceeds by sampling number of actions from
5

---

the flow-based policy and then utilizing the constraint oracle to filter the invalid actions. We then provide theoretical results showing the adaptation of the policy gradient computation accordingly.
Sandwich estimator: The ELBO (, is a biased estimator of log (als). but it known in literature S that i is consistent estimator (i.e., converges to log π [a]ss) in the limit) and there are generative methods based on this consistency property, such as [3]. However, we aim to use l in the policy gradient and stochastic gradient descent typically uses unbiased gradient estimate (but not always, see [6]) Therefore, we propose to use a new technique which combines ELBO and CUBO to reduce the bias, improving the convergence of our policy gradient based learning process. In particular, we estimate an upper bound of log π (als) using the following CUBO:
((
Pp(zols)
8 (x )
(als)) = Cx² 2
log LE xx (2xlaa,)
(6)
det Bfa.k
k==1
We then use weighted average of the upper and lower bounds low-bias estimate of log (als), denoted by log Pe.v alx + (1 ) where a is a hyperparameter. We call this the sandwich estimator of log probability. We observed in our experiments that an adaptive all,, , (+) as function of the two bounds provides better results than a static a = 2 (see Appendix [.. for more details).
Constraints: Since the agent only has oracle access to the constraints, existing safe exploration approaches [20] (24] are not directly applicable this setting. If the agent queries the validity of all actions, then it gains complete knowledge of the valid action set C(s). However, querying action validity for the entire action space at each timestep can be time-consuming, particularly when the action space is large, e.g., 1000 categories. We demonstrate this issue in our experiments.
To address this challenge of complex state- detteddeenndnt constraints in the full observation setting, propose a new policy gradient algorithm called Invalid Action Rejection Advantage Actor-Critic (IAR-A2C). IAR-A2C operates by sampling set of actions from the current flow-based policy and then leveraging the constraint oracle to reject all invalid actions, enabling the agent to explore safely with only valid actions. We then derive the policy gradient estimator for this algorithm. Recall that C(s) is the set of valid actions in state s and let C² (s) be the set of invalid actions. We sample an action from T (a(a))) (recall 0 (p, w) when using Flow Policy), rejecting any invalid action till a valid action is obtained. Clearly, the effective policy n' induced by this rejection produces (a(a))) obtain (a(as) Ea,ects) malails] for E C(s). For the purpose of policy gradient, need obtain the gradient of the long term reward with x': [(0)= = Ey-bo (s;+²)] We show that:
only valid actions. In fact, by renormalizing
Theorem With nt defined as above,
(s) Enec(a) (als(als) Enec(or (a(ss)
VoJ(0) = Eyr (s.
In practice, analogous to standard approach in literature [22 to reduce variance, we use the TD error (an approximation of advantage) form instead of the Q function in Theorem [ Then, given the network To. the empirical estimate of the first gradient term readily obtainable from trajectory samples of a' To estimate the second term, for every state s in the trajectories, we sample S actions from Th (als) and reject any Cs( to get / < S valid actions a1, ..., ag. We then apply Lemma |
Lemma 1 + Ezen Vo log ne (a and & are unbiased estimates of Exec(s) Vala((als) and Emec(s) To (aas)) respectively.
The full approach is shown in Algorithm 2. with the changes from standard A2C highlighted in red. Recall that, 0 = (p, w) is the parameter of the flow-based policy network, @ is the parameter of the critic, and v is the parameter of the variational distribution network qu for ELBO.
4 Related Work
Large discrete action space. Dulac-Arnold al. 110 attempt to solve this problem by embedding discrete actions into continuous space; however, this approach does not work for our unordered

---

Algorithm 2: IAR-A2C
Initialize step count 6- 1, Initialize episode counter F.
repeat
Reset gradients: d(e, < O, de < O, Initialize experience set T, Extart = t, get state 8t
rеpeat
5
Execute a; according to policy =g((a)(s)) Defined in Section 3.2 Get number of valid and sampled actions l, and St
7
Receive reward ry and new state $1+1
Add new experience <- U (st, at, re) and update
until terminal o Istart = Imaz
R R={
for terminal
10
(se)
for non-terminal st Bootstrap from last state
11
for i {t I,.... {stare] do
12
R+ r; +
log (a,1s,) = al (a.ls)) + (1 (ails) // Sandwich estimator
13
14
Accumulate gradients w.r.t (0, v): d(0, v) t
d(0, (b) (R Vs(s:)) (Vo, Po, (0) 34 Ejem Vo.elog Po. p
Accumulate gradients w.r.t. o: do de + 0(R (Vo)((ss))//0
16
Perform update of using d8, of using de and of ø using dø Execute ELBO updating by rumning Algorithm
17
| with r as input
G+E+1
19
until F.
Emaz
discrete space we reveal by thorough comparison in experiments. Another type of approach relies on being able to represent joint probability distribution over multi-dimensional discrete actions using marginal probability distribution over each dimension (also called factored representation). Tang and Agrawal [30] apply this approach to discretized continuous control problems to decrease the learning complexity. Similarly, Delalleau et al. 77 assumes independence among dimensions model each dimension independently. However, it is well-known in optimization and randomized allocation literature 35 that dimension independence or marginal representation not valid the presence of constraints on the support of the probability distribution. Another type of approach [34] [3)] converts the choice of multi-dimensional discrete action into a sequential choice of action across dimensions at each time step (using LSTM based policy network), where the choice of action in any dimension is conditioned on actions chosen for prior dimensions.
Safe actions. Our approach falls into the field of constrained action space in RL. Situated in the broader safe RL literature [I] 14), these works aim to constrain actions each step for various purposes, including safety. Notably, we do not have an auxiliary cost that we wish to bound, thus, frameworks based on Constrained MDP [1 cannot solve our problem. Also, our state dependent constraints if written inequalities can be extremely large and hence Lagrangian methods [2] that pull constraints into the objective are infeasible. Most of the methods that we know for constrained actions need the constraints to be written mathematical inequalities and even then cannot handle state dependent constraints. Although action masking [17 does offer a solution constraints, it necessitates supplementary information in the form of masks. Thus, they do not readily apply to our problem [20). Some of these methods 9 24 aim at diffcrentiating through an optimization; these methods
state-dependent
slow due to the presence of optimization in the forward pass and some approaches to make them faster[27] work only with linear constraints; all these approaches scale poorly in the number of inequality constraints.
Normalizing flow policy, The application of normalizing flows in RL has been an emerging field of study with a focus on complex policy modeling, efficient exploration, and control stability [18] 21 33 Normalizing flow was integrated into the Soft Actor-Critic framework to enhance the modelling expressivity beyond conventional conditional Gaussian policies and achieve more efficient exploration and higher rewards 21 33 In robotic manipulation, Khader al. [18] proposed improve the
control stability by novel normalizing fllo control structure. However, these works primarily focus

---

on continuous actions. To the best of our knowledge, ours is the first work that has successfully incorporated discrete flow, extending the scope of flow policies into discrete action spaces.
Experiments
Through our experiments, we aim to address two main research questions related our primary contributions: (1) Is the flow-based policy effective in representing categorical actions? (2) Does IAR-A2C offer advantages in constrained action space with oracle constraints? We address these two questions in Sections 5.1 and 5.2 accordingly. Following these, we present an ablation study that examines the effectiveness of various modules in approach.
We first describe our set-up. We evaluate IAR-A2C against prior works across a diverse set of environments, including low- dimensional discrete control tasks such as CartPole and Acrobot 33, the visually challenging Pistonball task (31 with high-dimensional image inputs and an extremely large action space (upto 59, 049 cat- egories), and emergency resource allocation simulator in a city, referred to as Emergency Resource Allocation (ERA). CartPole and Acrobot well-established environments. Pistonball is also stan- dard environment where a series of pistons on a line needs to move (up, down, or stay) to move ball from right to left (Figure 3). While the Pistonball environment was originally designed for a multi-agent setting with each piston controlled by distinct agent, we reconfig-
Figure 3: Pistonball, goal is move the ball to the left border by operating the pistons.
ure this as a single-agent control of all pistons. This modification presents a challenging task for the central controller as the action space is exponentially large. We show results for Acrobot and three versions of Pistonball: v1 with 3 v2 with 38, and v3 with 330 actions. Results for CartPole are in the appendix.
Finally, our custom environment, named ERA, simulates a city that is divided into different districts, represented by graph where nodes denote districts and edges signify connectivity between them. An action is to allocate a limited number of resources to the nodes of the graph. A tuple including graph, current allocation, and the emergency events is the state of the underlying MDP. The allocations change every time step but an allocation action is subject to constraints, namely that a resource can only move to a neighboring node and resources must be located in proximity (e.g. within 2 hops on the graph) as they collaborate to perform tasks. Emergency events arise at random on the nodes, and the decision maker aims to minimize the costs associated with such events by attending them as quickly as possible. Moreover, we explore a partially observable scenario (with no constraints) in which the optimal allocation randomized, thus, the next node for resource sampled from the probability distribution over neighboring nodes that the stochastic policy represents (see appendix for set-up details). We show results for five versions of ERA: ERA-Partial with 9 actions and partial observability in unconstrained scenario, while v2 with 73 actions, v3 with 8 actions, v4 with 93 actions, and v5 with 10 actions constrained scenario. The results for ERA-v1 are Appendix C
Benchmarks: In the unconstrained setting, we compare our approach to Wol-DDPG [0]. A2C [22]
factored representation (Factored, discussed in Section and autoregressive approach (AR, discussed in Section (4). Wol-DDPG is chosen as it is designed to handle large discrete action spaces without relying on the dimension independence assumption. In the oracle constraints setting, we compare method to action masking (MASK) [17]. which determines the action mask by querying all actions within the action space. As far as we know action masking currently the only existing approach for constraining action with oracle constraints, we also include a comparison with IAR-augmented AR (AR+IAR) which is able to handle the constraints by utilizing our invalid action rejection technique, as well as a comparison with Wol-DDPG to demonstrate the performance of a method that cannot enforce the constraints.
5.1 Learning in Categorical Action Space without Constraints
In the Acrobot environments, we use A2C as the benchmark because this task is relatively simple and A2C can easily learn the optimal policy. For the Pistonball environment, consider both Wol-DDPG and A2C as benchmarks. The results are displayed in Figure 4 We also conduct experiments on CartPole and ERA-v5 with all constraints removed, which can be found in Appendix
8

---

Acrobot-v1
Pistonball-v1
Pistonball-v2
Pistonball-33
1.00
1.00
1.00
Average Return
Average Return
Average Retum
Average Return
0.75
0.75
0.75
0.50
0.50
-3
0.25
0.25
0.00
0.00
-4
0.0
0.5 Timesteps
1.0
0.0
Timesteps
1e6
1e6
Timesteps
1e6
Timestops
1e6
Wol-DDPG
Factored
Figure 4: Learning curves time in (a) Acrobot and (b) Pistonballlv(1-3] (without constraints). All approaches are trained with 5 seeds, except for Wol-DDPG, trained with 3 seeds due to its expensive training cost on Pistonball.
lel
ERA.v2
1e1
E3A-v3
1e1
v4
le1
ERA-v5
Average Return
Average Return
Average Return
Average Return
-4
-6
-6
-6
-6
-8
-8
-8
6
4
Timesteps
Timesteps
Timesteps
1e5
Timesteps
les
Ours
Wol-DDPG
MASK
AR+IAR
Figure 5: Learning curves in ERA-v(2-5] (with constraints). All settings are trained with 5 seeds.
The results of Acrobot demonstrate that the flow-based policy exhibits comparable performance the optimal policy (A2C), which also highlights the effectiveness of our sandwich estima- tor for log π (als). more challenging environments with high- dimensional image observations and extremely large action spaces, our model has comparable performance to Factored and AR, while significantly outperforming A2C and Wol-DDPG. Interestingly, we observe that Wol-DDPG struggles to learn even in the simplest Pistonball environment, despite being de- signed for large discrete action spaces. We hypothesize that Wol-DDPG might function properly if the actions discrete yet ordered, as demonstrated in the quantized continuous con- trol task presented in the Wol-DDPG paper [10].
1e2
ERA-Partial
4
Average Return
Ours Factored
0.0
0.5
1.0
1.5
Timesteps
1e6
Figure 6: Learning curves in ERA- Partial. Our approach obtains higher return than Factored ap- proach.
ERA-Partial, we demonstrate the known fact that the opti- mal policy in environments with partial observability may be stochastic one [28). We compare with the factored approach (Figure 6p. In this environment, the optimal stochastic policy
is a joint distribution that not factorizable into a product independent distributions over each dimension. Thus, the factored approach cannot effectively represent the joint distribution due to its independence assumption among dimensions. The experiment results show our approach significantly outperforms the factored approach. We further present smaller example in the appendix which shows that the inability to represent an arbitrary joint distribution makes the factored approach extremely suboptimal in partial observable settings.
5.2 Learning in Categorical Action Space with State-Dependent Constraints
We now address the second question about constraint enforcement by IAR-A2C. The results are in Figure We observe that our approach demonstrates better or comparable performance compared to the benchmarks. AR known as a strong baseline for environments with large discrete action space, but surprisingly, it performs poorly. We hypothesize this is due to the case that the autoregressive model does not have a sense of the constraints of the remaining dimensions when outputs the action for the first dimension, thereby producing first dimension actions that may be optimal without

---

constraints but are suboptimal with constraints. Detailed analysis and experimental evidence to support our hypothesis are provided in Appendix
Also, action masking achieves its performance by querying all actions of the entire action space, whereas our approach only requires querying batch of actions, which is substantially smaller(e.g.. ... 64 versus 1, 000 for ERA-v5). Thus, while Fig- S shows IAR-A2C taking more iteration for convergence, the same figure when drawn with the x-axis wall clock time in Figure / (shown only for v4 here, others are in appendix C shows an order of magnitude faster convergence in wall clock
1e1
ERA-v4
Average Return
Moaa M M
-4
-6
MASK
Ours
time. Another critical property of our approach is the guar- anteed absence of constraint violations, similar to the action masking method. However, while action masking demands the full knowledge of the validity all actions, our method merely requires the validity of the sampled actions within batch. Note that Wol-DDPG can and does violate constraints during the learning process. Further details regarding the action violation of Wol-DDPG are provided in the appendix
20 Wall time (h)
40
Figure Learning curves for wall clock time. Our approach converges much faster than action masking.
5.3
Ablation Study
We conduct ablation studies of IAR- A2C on ERA-v4 to investigate the ef- fect of various choices modules and estimators in our approach.
ERA.v4
1el
ERA-v4
Valid action rate: 15
1.0
Full model
0.8
Average Return
flow q
-4
0.6
We compare the performance of approaches using the policy gradient provided by Theo- rem |
Policy Gradient:
0.4
model
0.2
(gradient correction) and the original policy gradient of A2C (stan- dard policy gradient), while still ap- plying invalid action rejection. We observe in Figure Sa that the number of valid actions in the batch decreases rapidly, and the program may crash if no valid actions are available.
0.0
wio grad cor
0
2
4
6
Timesteps
le5
Timesteps
1e5
(a)
(b)
Figure 8: (a) Ablation of gradient correction and sandwich
estimator (b) Ablation of posterior type.
Sandwich Estimator: We examine the effects if we use only the ELBO estimator for log probability of action instead of our sandwich estimator. We find that the ELBO estimator is also prone reduction in valid actions in Figure Sa and unstable learning as a consequence, similar to the observation when using the standard policy gradient.
The posterior g(2la, s) can be modeled by a conditional Gaussian or normalizing flow. In our experiments, we discover that modelling with flow posterior is crucial for learning, as it can approximate the true posterior more effectively than a Gaussian, as seen in Figure Sb
Posterior Type:
Conclusion and Limitations
We introduced novel discrete normalizing flow based architecture and an action rejection approach to enforce constraints on actions in order to handle categorical action spaces with state dependent oracle constraints in RL. Our approach shows superior performance compared to available baselines, and we analyzed the importance of critical modules of our approach. A limitation of our method is in scenarios where the fraction of valid actions out of all actions is very small, and hence our sampling based rejection will need a lot of samples to be effective, making training slower. This motivates future work on improved sampling; further, better estimation of the log probability actions also promising research direction for improved performance.
10

---

Acknowledgement
This research is supported by the National Research Foundation Singapore under its Al Singapore Programme (Award Number: AISG2-RP-2020-016 Dr. Nguyen was supported by grant W91 INF. 20-1-0344 from the US Army Research Office.
References
[1] Abhinav Bhatia, Pradeep Varakantham, and Akshat Kumar. Resource constrained deep rein- forcement learning. In Proceedings of the International Conference on Automated Planning and Scheduling, volume 29, pages 610-620, 2019.
[2] Steven Bohez, Abbas Abdolmaleki, Michael Neunert, Jonas Buchli, Nicolas Heess, and Raia Hadsell. Value constrained model-free continuous control. ArXiv preprint, abs/1902.04623, 2019. BtR/pttss://arxtttt http://:///aaa..oooraaaa//a/////
04623
[3] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. ArXiv preprint, abs/1606.01 540, 2016. URL https://arxiv org/ abs 01540
[4] Eric Budish, Yeon-Koo Che, Fuhito Kojima, and Paul Milgrom. Designing random allocation mechanisms: Theory and applications. American economic review, 103(():585-623 623, 2013.
[5] Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In Yoshua Bengio and Yann LeCun, editors, 4th International Conference on Learning Repre- sentations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL http:://aaar: http://arxxvv..rra///55.5..... 00519
[6] Jie Chen and Ronny Luss. Stochastic gradient descent with biased but consistent gradient estimators. ArXiv preprint, abs/1807. 1880, 2018. URL UKL.Ittss///rradv.or///bbbb/b
11880
[7] Olivier Delalleau, Maxim Peter, Eloi Alonso, and Adrien Logut. Discrete and continuous action representation for practical rl in video games. ArXiv preprint, abs/1912, 1077, 2019. URL
https://arxix: org/ abs /1912. 11077
[8] Adji Bousso Dieng, Dustin Tran, Rajesh Ranganath, John Paisley, and David Blei. Variational inference via chi upper bound minimization. Advances in Neural Information Processing Systems, 30, 2017.
[9] Priya Donti, David Rolnick, and J. Zico Kolter. DC3: A learning method for optimization with hard constraints. In 9th International Conference on Learning Representatiins, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. Jttplltttss://peerreee
net/forumh
[10] Gabriel Dulac-Amold, Richard Evans, Hado Hasselt, Peter Sunehag, Timothy Lillicrap, Jonathan Hunt, Timothy Mann, Theophane Weber, Thomas Degris, and Ben Coppin. Deep reinforcement learning in large discrete action spaces. ArXiv preprint, abs/1512.07679, 2015.
URL
abs/ 1/////1111000 07679
[11] Javier Garcia and Fernando Fernández. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research, 16(1):1437-1488) 2015.
[12] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(1 1):134-144, 2020.
[13] Sergei Pavlovich Grachev. Aleksei Aleksandrovich Zhilyaev. Vladimir Borisovich Laryukhin, Dmitrii Evgen'evich Novichkov, Vladimir Andreevich Galuzin, Elena V Simonova, IV Maiy- orov, and Peter O Skobelev. Methods and tools for developing intelligent systems for solving complex real-time adaptive resource management problems. Automation and Remote Control, 82:1857-1885, 2021. 2021.
11

---

[14] Shangding Gu, Long Yang, Yali Du, Guang Chen, Florian Walter, Jun Wang, Yaodong Yang, and Alois Knoll. review of safe reinforcement learning: Methods, theory and applications. ArXiv preprint, abs/2205. 2022. URL URI.ttpss: aww.aaiiiivoo.. org/
10330
[15] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin, editors, Advances in Neural Information Processing Systems 33: Anmual Con- ference Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12. 2020, virtual, 2020. URL https//////orrccddnngg wtwpr/rrwweennnnnnn.nnr.rpepeeee neurips wwtpcccppppppeeeeee h0sh/hash/ html
@c5bcfec8584affo 01
[16] Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. In Marc' Aurelio Ran- zato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 12454-12465, 2021. URL U htw.ooceednnnngneurrrrssse. neurips wwtpppper//02///aaeeeeee 7ra666555aeee2229966666652222222 wwwtraat.tttttte
[17] Shengyi Huang and Santiago Ontañón. A closer look invalid action masking in policy gradient algorithms. In The International FLAIRS Conference Proceedings, volume 35, 2022.
[18] Shahbaz Abdul Khader, Hang Yin, Pietro Falco, and Danica Kragic. Learning stable normalizing- flow control for robotic manipulation. In 202/ IEEE International Conference on Robotics and Automation (ICRA), pages 1644-1650. IEEE, 2021.
[19] Dexun Li, Meghna Lowalekar, and Pradeep Varakantham. Claim: Curriculum learning policy for influence maximization in unknown social networks. In Uncertainty in Artificial Intelligence, pages 1455-1465. PMLR, 2021.
[20] Jyun-Li Lin, Wei Hung, Shang- Shang-Hsuan Yang, Ping-Chun Hsieh, and Xi Liu. Escaping from zero gradient: Revisiting action-constrained reinforcement learning via frank-wolfe policy optimiza- tion. In Cassio P. de Campos, Marloes H. Maathuis, and Erik Quaeghebeur, editors, Proceedings of the Thieth-Seventh Conference on Uncertainty in Artificial Intelligence, UAI 2021, Virtual Event, 27-30 July 2021, volume 161 of Proceedings of Machine Learning Research, pages 397- 407. AUAI Press, 2021. URL https::/ppppooooo https:://proceednngg .ar.ggsmmmrr.r... mlr,
www.mmrr.........mmmmmmmmmmmmm. press/ wwt.:prrssssvvveie vwtp//i21111111bbbe/ hhmml
[21] Bogdan Mazoure, Thang Doan, Audrey Durand, Joelle Pineau, and R Devon Hjelm. Leveraging exploration in off-policy algorithms via normalizing flows. In Conference on Robot Learning, pages 430-444. PMLR, 2020.
[22] Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforce- ment learning. In Maria-Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 1928-1937. JogR. 2016. URL http::/ppoceeeiiiiiee
mlr prwss/eessseee v48 pp8ss//48// www.m1m......mm html
[23] Ayan Mukhopadhyay, Geoffrey Pettet, Sayyed Mohsen Vazirizade, Di Lu, Alejandro Jaimes, Said El Said, Hiba Baroud, Yevgeniy Vorobeychik, Mykel Kochenderfer, and Abhishek Dubey. A review of incident prediction, resource allocation, and dispatch models for emergency man- agement. Accident Analysis & Prevention, 165: 105:10650,, 2022.
[24] Tu-Hoa Pham, Giovanni De Magistris, and Ryuki Tachibana. Optlayer- -pacttcca constrained optimization for deep reinforcement learning in the real world. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pages 6236 06646223.222 IEEE, 2018
[25] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021. URL. http://jmmh jmlr. .erellrorrg//oo org/papers/voa
20 1tm4.html
12

---

[26] Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Francis R. Bach and David M. Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, ICMI. 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pages 1530-1538. JMLR.org 2015. URK.htttp://proceeding
mlr. ww.r.pre///eee. wpwpss//3///////pppppp.
v37/ rtwt///eeeeeeeeee html
[27] Sanket Shah, Arunesh Sinha, Pradeep Varakantham, Andrew Perrault, and Milind Tambe. Solving online threat screening games using constrained action space reinforcement learning. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty- Second Innovative Applications of Artificial Intelligence Conference, February 7-12, 2020. pages 2226 2235. AAAI Press, 2020. URL htttps://aat htt///aaii........ .rg. hetg/ojs/inddee. php/ AAAI/ article/view/5500
[28] Satinder P Singh, Tommi Jaakkola, and Michael I Jordan. Learning without state-estimation in partially observable markovian decision processes. In Machine Learning Proceedings 1994, pages 284-292. Elsevier, 1994.
[29] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. 9th In- lernational Conference on Learning Representations, ICLR 2021, Viriual Event, Austria, May 3-
7, 2021. OpenReview.net, 2021. URL/kttpp://pppererr// wwwp:eeeeefforooe
[30] Yunhao Tang and Shipra Agrawal. Discretizing continuous action space for on-policy op- timization. In The Thirty-Fourth AAAI Conference Artificial Intelligence, AAAI 2020,
The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020,
New York, NY, USA, February 7-12, 2020, pages 5981-5988. AAAI Press, 2020. URL https: http://aaaa.orgojj////////// wttp:///////........//rrr http://p///taartteeeeeeeee//////
[31] Justin K. Terry. Benjamin Black, Nathaniel Grammel, Mario Jayakumar, Ananth Hari, Ryan Sullivan, Luis S. Santos, Clemens Dieffendahl, Caroline Horsch, Rodrigo ViceztVeeeete,, Niall L. Williams, Yashas Lokesh, and Praveen Ravi. Pettingzoo: Gym for multi-agent reinforcement learning. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 15032-15043, 2021. URL bttts://////oeeeeediiiii wwwstattttttttttttttttt........ html
neur sourips.co cc/
[32] Ivan B Vermeulen, Sander M Bohte, Sylvia G Elkhuizen, Han Lameris, JM Bakker, and Han La Poutré. Adaptive resource allocation for efficient patient scheduling. Artificial intelligence in medicine, 46(1):67-80, 2009.
[33] Patrick Nadeem Ward, Ariella Smofsky, and Avishek Joey Bose. Improving exploration in soft-actor-critic with normalizing flows policies. ArXiv preprint, abs/1906.02771, 2019, URL
https://arxiv.. htt///as//110 02771
[34] Mason Wright, Yongzhao Wang, and Michael Wellman. Iterated deep reinforcement learning in games: History-awan training for improved stability. In EC, pages 617-636, 2019.
[35] Haifeng Xu. The mysteries of security games: Equilibrium computation becomes combinatorial algorithm design. In Proceedings of the 2016 ACM Conference on Economics and Computation, pages 497-514, 2016.
[36] Yiming Zhang, Quan Ho Vuong, Kenny Song, Xiao- Yue Gong, and Keith W Ross. Efficient entropy for policy gradient with multidimensional action space. ArXiv preprint, Abs/1806.00089) 2018.7 URL httpt:://///t //trxiv...ivveve wrciivvoorrrraee wrg/abs../a....ooo.roo...o.. http:/ab/1100..000 00589
13


Reviews:
[{"review_id": "7VmADbECEa", "reviewer": "Reviewer_NSnS", "summary": "This paper studies challenges in Reinforcement Learning (RL) problems that involve finding an optimal policy where the action space is categorical and large in which current RL approaches do not have a good performance. In addition, these problems often require validity constraints on the actions, which are difficult to express mathematically. To tackle these challenges, the authors propose using discrete normalizing flows conditioned on states to represent the stochastic policies compactly. They employ an invalid action rejection method and derive a modified policy gradient for updating the base policy. Lastly, they evaluate their method by comparing its performance with other RL benchmarks on several tasks in which they achieve state-of-the-art results.", "questions": "1. In figure 2, it is not clear what K is. It would be helpful if the caption was more descriptive.\n2. From my understanding, since the argmax is going to give us the action, it must be applied to a layer with dimension equal to the number of actions. Is my understanding correct? If yes, how is this compact compared to the categorical policies implemented using a Softmax layer?\n3. How does the computational power required by your method compare standard categorical policies for the same task? To me, it seems that you require far more computation to represent the policy as computing q and the density (requiring computing a Jacobian) is required with your method whereas it is not with categorical policies.\n4. On page 4, you show how the state encoder works. Is this the standard way to condition the normalizing flow on some given vector? I would assume that traditionally, the noise is concatenated or summed with the condition.\n5. For the gradient update on the return (shown in equation 7), have you considered using reparameterization of the policy (since it is now represented with normalizing flows) to avoid using the re-inforce trick? I believe the fact that you can reparameterize the policy is an advantage as it can lead to potentially smarter gradient updates.\n6. Can you provide any experimental results demonstrating the compactness of your method?\n7. In figure 5, why is your model slower to converge than the other benchmarks?\n8. In your experiments, do you add the same state encoder used for your approach, to the other benchmark algorithms? I believe that if the state encoder has a positive effect on your method, it has to be equally considered in other benchmarks you are comparing to.", "limitations": "Overall, this paper successfully introduces Argmax Flows as a representation method for categorical policies, achieving state-of-the-art results within the provided tasks and algorithms, which is commendable. However, at this stage, I perceive the proposed method as an alternative to standard categorical policies that may offer better performance but are computationally more demanding. To enhance the paper, further refinement is required in two key areas. First, clarify the claims regarding the superiority of this method. If its strength lies in compactness, provide a detailed explanation and demonstrate this aspect through experiments. If performance is the main advantage, showcase its generalizability across different tasks and algorithms. Additionally, it's important to note that I maintain an open mind and may reconsider my evaluation based on a compelling rebuttal.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Overall Strengths:\n1. I personally find the high-level idea of looking into ways to improve the performance of RL policies invaluable.\n2. The idea of representing stochastic policies for discrete action spaces with discrete normalizing flows is well-worth pursuing. \n3. The overall performance of the proposed method is superior to the benchmarks shown in the experimental section, which shows the effectiveness of the approach.", "weaknesses": "Overall Weaknesses:\n1. This paper emphasizes that standard categorical policies implemented using a Softmax layer require large resources as the size of the action space grows. They then motivate how their proposed Argmax Flow-based model is instead a compact representation of the policy. However, it is very vague to me if this claim is really true. I have clarified this point in the “Questions” section of the review.\n2. In some parts of the paper, there are terms used that are not described until later on in the paper which makes the paper a little bit difficult to follow. I have clarified this point in the “Questions” section of the review.\n3. The experimental setup is limited. I would personally be interested in experiments that show how the proposed method compares to others in terms of compactness (as it is one of the focuses in the write-up). Moreover, comparing against other state-of-the art algorithms and benchmark tasks would be invaluable. The current experimental results are promising, but I still wonder if they generalize across different algorithms and tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and feedback.\n\nWe have implemented an autoregressive model (unfortunately, we did not find any publicly available code) and results are in the attached pdf; also present are results for the marginal/factored approach. \n\nFigure 1 in the attached PDF shows that our approach has similar performance as the factored/marginal and the autoregressive approach in the **unconstrained case**, which is as per expectation since the dimensions are independent and the factored approach assumes underlying independence across dimensions. The autoregressive approach assumes a particular dependence structure (based on ordering of dimensions) but is equivalent to marginal A2C when there is underlying independence across dimensions.\nFigure 3 in the PDF shows that the autoregressive approach does not perform well in the **constrained case**.\n\nWe have additional experiment still running and hope to share these during the discussion phase.\n\nResponse to questions:\n\n1. We will make the caption more descriptive. K here is the last of K invertible neural network layers used in normalizing flow (described in lines 150-151)\n\n\n2. The argmax is applied to the last dimension of the continuous $v \\in \\mathbb{R}^{D \\times M}$. As a example, suppose number of dimensions are D=2 with M=3 choices in each dimension and if the continuous output by standard normalizing flow is $v = [[ 0.4, 0.5, 0.9], [2, 0.7, 1]]$ then argmax(v, dim=1) gives [2, 0], that is, choice 2 in first dimension and choice 0 in second dimension. Note that the output size of the continuous part is $D \\times M$ and not the number of $M^D$ actions possible. We do note that all this is from the previous work in Argmax flow (described in lines 82-94)\n\n\n3. Yes, compared to marginal (or independent categorical) representation, our approach will consume more computational power as we have to maintain an additional posterior network. But, our approach does not require computing Jacobians; the log determinant that shows in line 11 in Alg 1 is efficiently computed by appropriate choice of the invertible function $f_{w,k}$. There are many choices for $f_{w,k}$ used in normalizing flow literature all of which ensure efficient computation of this log determinant, more details can be found in\n[​​Papamakarios et al. 2021] (second last paragraph of page 11). \n\n   We also discuss our choice of $f_{w,k}$ in Appendix B.2.\n\n   [​​Papamakarios et al. 2021] ​​Papamakarios, George, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. \"Normalizing flows for probabilistic modeling and inference.\" The Journal of Machine Learning Research 22, no. 1 (2021): 2617-2680. \n\n4. Unlike other generative models, in normalizing flow typically the size of the noise input is the same as the size $D \\times M$ of the output (due to the invertible layers used in between). Given this, we would have been forced to generate a state encoding of size $D \\times M$, which then would impose a size restriction for baselines. Thus, we used this approach which has been used in other papers also [Ward et al. 2019].\n\n   [Ward et al. 2019] Ward, P. N., Smofsky, A., & Bose, A. J. (2019). Improving exploration in soft-actor-critic with normalizing flows policies. arXiv preprint arXiv:1906.02771.\n\n\n5. We have not considered reparameterization, though we agree that this is another interesting way of getting gradients made possible by use of the generative model as a policy network. We aim to explore this in future pertaining to what advantages would this offer over and above the reinforced style gradient we use.\n\n\n6. As stated in answer 2, the compactness is self evident in the size of outputs $D \\times M$ that we have used relative to the $M^D$ size that the standard A2C would use, given the size of the action space $|\\mathcal A| = M^D$. \n\n\n7. Our approach appears slower in the number of iterations, but the same result in actual wall clock time (shown in Figure 6 in main paper and Figure A.3 in appendix) shows the much faster convergence in time. This is because MASK takes a lot of time per iteration.\n\n8. Yes, we use the state encoding for all baselines.\n"}, {"review_id": "AyGpFhVXYT", "reviewer": "Reviewer_We8u", "summary": "The submission investigates reinforcement learning in the context of large categorical action spaces. The submission proposes parameterizing the policy using a state conditional normalizing flow in these settings. The submission derives a rejection-sampling-based method for handling valid action constraints determined by a black-box oracle. The submission shows the performance of its approach in experiments.\n\n---\n\nI am disclosing here that I have not checked the proofs of the theorems.", "questions": "> Think of the things where a response from the author can change your opinion, clarify a confusion or address a limitation\n\nI think adding baselines that use factored or autoregressive action spaces could cause me to increase my score. I will not decrease my score if these additional baselines perform well -- I think it would be valuable for the community to know in either case.\n", "limitations": "> Have the authors adequately addressed the limitations\n\nYes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The problem that the submission investigates is both important and underexplored. The submission's approach seems reasonable and principled.", "weaknesses": "The submissions writes the following:\n\n> Another type of approach relies on being able to represent a joint probability distribution over multi-dimensional discrete actions using marginal probability distribution over each dimension (also called factored representation). Tang and Agrawal [27] apply this approach on discretized continuous control problems to decrease the learning complexity. Similarly, Delalleau et al. [7] assumes independence among dimensions to model each dimension independently. However, it is well-known in optimization and randomized allocation literature [4, 32] that dimension independence or marginal representation is not valid in the presence of constraints on the support of the probability distribution. Another type of approach [33, 31] converts the choice of multi-dimensional discrete action into a sequential choice of action across the dimensions at each time step (using a LSTM based policy network), where the choice of action in any dimension is conditioned on actions chosen for prior dimensions. This heuristic also produces a deterministic policy, thus, it is not suited for problems where the optimal policy is stochastic.\n\nI have two criticisms/questions about this passage. First, I don't understand the sentence referencing \"a deterministic policy\". It is of course possible for autoregressive models to parameterize non-deterministic policies. Could the authors clarify exactly what is meant here? Second, while I think the submission's criticisms of these approaches is fair, it's point would be made more strongly if it showed that these flaws resulted in worse experimental performance. Sometimes, it is the case that a \"stupid but easy\" approach ends up performing well in practice. If the submission were to show that this is not the case for the environments it examines, it would strengthen its value.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and feedback.\n\nWe agree that that autoregressive model can also produce stochastic policies. We will modify this statement. We have implemented an autoregressive model (unfortunately, we did not find any publicly available code) and results are in the attached pdf; also present are results for the marginal/factored approach. \n\nFigure 1 in the attached PDF shows that our approach has similar performance as the factored/marginal and the autoregressive approach in the **unconstrained case**, which is as per expectation since the dimensions are independent and the factored approach assumes underlying independence across dimensions. The autoregressive approach assumes a particular dependence structure (based on ordering of dimensions) but is equivalent to marginal A2C when there is underlying independence across dimensions. Figure 3 in the PDF shows that the autoregressive approach does not perform well in the **constrained case**.\n\nWe have additional experiment still running and hope to share these during the discussion phase"}, {"review_id": "iAQU4GTcAz", "reviewer": "Reviewer_Ed9B", "summary": "The paper proposes a new approach to solve RL problems with large categorical action spaces and state-dependent constraints. The approach consists of two main components: a flow-based policy network that uses a discrete normalizing flow model to represent the stochastic policy, and an invalid action rejection method that leverages an oracle to filter out actions that violate the constraints. The paper derives a modified policy gradient estimator for the proposed method and shows its effectiveness on control tasks (Acrobot, Pistonball) and resource allocation task (ERA).", "questions": "\n1. It is not clear why 'such constraints are easy to validate by a validity oracle given any allocation action, but are hard to represent as mathematical constraints' (line 35). Can you provide a more concrete example?\n\n2. How does the set of actions sample from the flow policy (Sample actions from flow policy in Figure 1)? Is it obtained by multiple sampling from the distribution of policy outputs? If so, why not switch to a single time of sampling until getting a valid action?\n\n3. What is constraints in Acrobot and Pistonball tasks as they seems like common RL tasks?\n\n", "limitations": "Refer to 'Weakness' part. I would like to present some minor limitations that do not affect the score:\n1. \"dependent constraints..\" -> \"dependent constraints.\"\n2. Introduce the motivation for using Flow based Policy network in the introduction section.\n3. Add more descriptions about the experimental setting in the experiment section, e.g., details about the action constraints.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The idea of the combination of discrete normalizing flow is novel, and the way of invalid action rejection to address the challenges of state-dependent constraints is simple but reasonable. Besides, the paper provides a sound theoretical foundation (though I did not check it very carefully) and empirical evaluation for the proposed method, showing its advantages over existing baselines. It provides a clear definition of the problem and reasonable hypotheses. The proposed methods are well presented, and the experiments demonstrate the effectiveness of the proposed method. Overall, the conclusions are supported by empirical evidence and theoretical analysis. I believe this paper develop a reasonable method for solving large action space and resource allocation problem. ", "weaknesses": "There are some limitations/weakness I want to discuss:\n\n### Method:\n\nThe motivation is not clear. I can not get the advantage of using generative policies to address resource allocation problems. The introduction section directly proposes the discrete normalizing flow policy. Especially, MASK, without generative policies, also demonstrates good performance in the experiments. Besides, the theorem/proof is straightforward and intuitive.\n\n\n\n\n### Experiments: \n\n1. The experiment results are weak on ERA tasks. Though IAR-A2C slightly outperforms MASK in terms of the final performance. However, it's clear that MASK learns more efficient on ERA-v[2-5], as shown in Figure 5. MASK is actually a simple method that only considers the action constraints by adding an action probability masking. Though author presents the wall time of the training (Figure 6, which partly alleviates my concern), there are not many scenarios where you need to query to know about the action constraints. Thus I am wondering the necessity of using the proposed method than MASK.\n\n2. For ablation study, the paper does not compare the proposed method with other generative models for RL, such as GANs or diffusion models, which could provide alternative solutions to normalizing flow policy.  \n\n3. The paper tries to tackle the problem of resource allocation, but only evaluate the proposed method on the customized environment. It can be further improved by conducting experiment on more widely used publich benchmarks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and feedback.\n\n* MASK comparison: The MASK approach can work when the mask can be specified easily. For constraints that vary by state (need a mask per state) and where the number of actions is large specifying the mask is costly - particularly, for our main problem of emergency resource allocation specifying such a mask per state is simply not possibly (apriori specification and storing of such mask would requires enormous storage space). \n* Other generative models: The paragraph starting line 134 discusses the reason for our choice of normalizing flow - GAN does not provide probability density estimates and diffusion models are much more computationally heavy than normalizing flow.\n* Benchmark: There is no widely accepted benchmark for heterogeneous resource allocation, we design our benchmark inspired by some prior works. \n   Prior works did not follow OpenAI Gym API [1, 2, 3], which requires decent efforts on implementation to match their APIs. Also, we encountered dependency errors when we attempted to use the environments. \n\n   [1] Vehicular Network Environment: GitHub - neardws/Game-Theoretic-Deep-Reinforcement-Learning\n \n   [2] Decentralized resource allocation for vehicle-to-vehicle: GitHub - Engineer1999/Double-Deep-Q-Learning-for-Resource-Allocation\n\n   [3] Decentralized resource allocation for 5G vehicle-to-vehicle: GitHub - gundoganalperen/DIRAL\n\n\nResponse to questions:\n1. Suppose there are $M$ areas and resources $R_1, R_2, …, R_N$  to be allocated ($ M \\approx 1.5N$). A constraint can be that any subset of $N/2$ areas must have at least one area covered by at least one resource. \n\n   Now, if given an allocation, it is easy to check the number of areas that have no allocation. If this number is $\\geq N/2$ then the constraint is violated (i.e., there is a subset of size N/2 with no allocation) and if this number is $< N/2$ then the constraint is satisfied (as any subset of N/2 areas will contain an allocated one). This is easy to check.\n\n   For a mathematical representation, let $x_{n, m} \\in \\\\{0,1\\\\}$ be a binary variable indicating allocation of resource $n$ to location $m$. Simple allocation constraints make sure that a resource is only allocated once: $\\sum_m  x_{n, m} = 1$ for all n.\n   To make sure that a subset of targets $i_1, \\ldots, i_k$ has at least one area allocated, we can write $\\sum_{n,k}  x_{n, i_k} \\geq 1$\n   Then, to enforce the constraint, we need \n   $\\sum_{n,k}  x_{n, i_k} \\geq 1$ for all possible subsets $\\\\{i_1, \\ldots, i_k\\\\} \\subset \\\\{1, \\ldots, M\\\\}$ with $k = N/2$\n   The number of such constraints are $M \\choose N/2$, which is large.\n   (We note that in many cases simplification of constraints might be possible but that is done typically on a case by case basis)\n\n\n2. The samples are obtained by calling the conditional normalizing flow policy multiple times using multiple samples from the noise z. We need both valid and invalid samples to estimate the terms in the modified policy gradient. Sampling till getting the first valid sample makes such estimation complicated and likely higher variance (as sometimes the first sample itself might be valid).\n3. As stated in the paper (e.g., caption of Figure 4), Acrobot and pistonball are the unconstrained case experiments in our main paper. Further, we have experiments with a constrained version of pistonball in Appendix C that restricts the upward movement of pistons on the left side to ensure the ball continues rolling to the left. \n"}, {"review_id": "oJ5IiBqn0l", "reviewer": "Reviewer_Xx2Q", "summary": "This paper introduces an algorithm for learning with large discrete action spaces where only some actions are valid for any given state in ways that are hard to programmatically exclude. The paper proposes to use normalizing flows to learn a policy across such action spaces without having to model the full action probability distribution, which is difficult to estimate and slow to learn. Algorithmic details for how to implement such a flow policy are discussed, such as the use of a sandwich bound for better action log probability estimation, and an oracle action validator to constrain invalid actions and reweight valid ones.", "questions": "Some other suggestions (in rough decreasing order of importance):\n\n-I get why without the gradient correction the policy gets stuck in section 5.3, but it's not obvious to me why this happens with ELBO versus sandwich bounds as well since ELBO should be biased but consistent similar to sandwich (which should be a better bound, but theoretically ELBO should still converge, no?). This result should be explained a little more (and I'd also appreciate ELBO versus sandwich on e.g. pistonball to show the difference when valid action issues don't apply).\n\n-I got a little lost in the details and notation at the end of section 3.1. While I get the gist of things, this explanation of the structure of a flow policy and figure 2 illustrating it seem somewhat confusing and could be improved, since this is one of the core algorithmic contributions of the paper.\n\n-Theorem 1 seems like it could be stated more simply. This is just reweighting the gradients based on the fraction of sampled actions which are valid, no? I'd argue that doesn't need a whole theorem to justify (unless I'm missing something here)\n\n-Please move the symbol definitions from the last paragraph of section 3.2 to the algorithm 2 box, it's awkward having these definitions stated in the text rather than in the box. Also adding other givens/definitions to the box would be nice as well.\n\n-There's a couple places where several things are listed followed by \"etc.\" I'd suggest removing the \"etc\" since it looks too casual stylistically.\n\n-There's several spots of awkward grammar/phrasing throughout the paper. Not enough to impede understanding, but I'd recommend an editing pass to clean them up before camera-ready.", "limitations": "In general the limitations of the presented algorithm are clearly stated (questions about comparison to prior work above aside). There's clearly room to improve with valid action handling, but this is stated by the authors, and it's a complex enough problem that leaving it to future work is valid.\n\nRegarding negative social impact, I don't think the core algorithm has any issues. The specific application example of routing emergency services between hotspots and potentially triaging between different incidents does have non-trivial moral hazard, and using the system to coordinate police response to (for example) protests has high risk for negative societal impacts. However, as that application seems to be mostly illustrative and it's not clear that the authors intend their algorithm to be used directly in the real world I think the paper itself is fine as it is (though I might suggest a different example use case for future work to avoid such concerns if possible).", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper sets out the goal of improving RL training efficiency in the context of very large discrete action spaces without assuming any structure or ordering among the actions (though it optionally assumes an action validity checker when some actions may be invalid for some states). Very large action spaces are a known issue impeding learning, typically solved by assuming domain-specific structure or allowing  e.g. step-by-step action composition, which may be unsuitable sometimes and are inelegant in a sense, so this problem is worth addressing.\n\nThe paper itself is largely well written. I generally was able to follow the arguments and math as presented, with a few notable exceptions discussed below.\n\nThe approach presented for tackling this problem seems intuitive and draws on established algorithms like normalizing flows and ELBO/CUBO in a nontrivial way. The authors argue for the use of normalizing flows among other generative models here, and their argument seems persuasive as a nonexpert in generative modelling.\n\n", "weaknesses": "First and foremost, the current abstract needs to change. It does a very poor job summarizing the content of the paper, and seems bizarrely out of place compared to the rest of the paper, which is mostly well written and clear. I put this in \"weaknesses\" rather than \"questions\" because this seems like a major issue that needs to be resolved prior to publication.\n\nRelated tangentially to the abstract, the authors assert several times that the optimal policy needs to be stochastic for the domain they are interested in. I'm unsure why this is asserted, since it's not obvious that that is the case for the custom evaluation domain. Perhaps there's something here that I'm not seeing, but this seems like an out of place insistence that's not relevant to the core problem of very large discrete action spaces. It's useful to have a stochastic policy during training, but I'm not sure offhand what conditions would lead to the optimal policy being stochastic, especially in the discrete action domain.\n\nSegueing nicely into another concern I have, that assertion (the optimal policy must be stochastic) is used to dismiss one approach that is commonly used to address very large action spaces, to decompose a single action decision into a sequence of decisions with each decision conditioned on the previous decisions for that action. I'm sympathetic to the idea that this isn't a perfect solution (it's obviously nicer to act directly in the primitive action space), but this doesn't seem like grounds to dismiss such an approach and not benchmark against it.\n\nThe last of my major concerns is that the main comparison to prior work, versus Wol-DDPG, is not valid because Wol-DDPG assumes spatial structure among discrete actions, which the authors briefly mention as an issue in section 5.1 and then ignore. It's unclear whether Wol-DDPG was used correctly here, particularly since all the learning curves shown for that algorithm are flat- the agent does not improve over the course of training on any variant of either task. Absent that comparison, the experiments clearly show that IAR-A2C outperforms vanilla A2C and is much less computationally expensive than action masking on the two tasks respectively, but I'm concerned by the lack of comparison to previous work given that this isn't a new problem.\n\nI'll also register that the sample-and-reject approach to action validity used is somewhat unsatisfying and a \"weakness\", but this is acknowledged by the authors as a limitation of the current algorithm and I'm unsure what a better non-domain-specific solution would be, so I'm not holding that against this paper as it is.\n\nGiven all the above, while I like the core concept here, I'm hesitant to recommend acceptance as is. The algorithm seems sound, but as the problem of large action spaces is chiefly a practical (rather than theoretical) one I feel a solid performance comparison to prior methods (or a more convincing argument for why such a comparison is invalid) is needed to validate the algorithm.\n\n ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and feedback.\n\nWe will write an alternate abstract and present it to the reviewer in the discussion phase (we are just not sure if that is allowed in this first round of rebuttal).\n\nIn allocation problems, stochastic allocation, if one exists, is often preferred over deterministic ones due to reasons of fair tie breaking and not starving a recipient of resources (please see Introduction in [Budish et al. 2013]). In a sequential decision making problem, a stochastic optimal policy can exist in a MDP if there are two or more actions that maximize the Q value in a state. Further, with partial observability there can be scenarios where the optimal policy is stochastic and this optimal stochastic policy is better than any deterministic policy (see Fact 2 in [Singh et al. 1994]).\n\n[Budish et al. 2013] Eric Budish, Yeon-Koo Che, Fuhito Kojima, and Paul Milgrom. Designing random allocation mechanisms: Theory and applications. American economic review, 103(2):585–623, 2013. \n\n[Singh et al. 1994] Singh, Satinder P., Tommi Jaakkola, and Michael I. Jordan. \"Learning without state-estimation in partially observable Markovian decision processes.\" In Machine Learning Proceedings 1994, pp. 284-292. Morgan Kaufmann, 1994.\n\nWe have implemented an autoregressive model (unfortunately, we did not find any publicly available code) and results are in the attached pdf; also present are results for the marginal/factored approach. \n\nFigure 1 in the attached PDF shows that our approach has similar performance as the factored/marginal and the autoregressive approach in the **unconstrained case**, which is as per expectation since the dimensions are independent and the factored approach assumes underlying independence across dimensions. The autoregressive approach assumes a particular dependence structure (based on ordering of dimensions) but is equivalent to marginal A2C when there is underlying independence across dimensions. \nFigure 3 in the PDF shows that the autoregressive approach does not perform well in the **constrained case**. For this experiment, we used our IAR approach to enforce constraints, as the autoregressive approach by itself does not enforce constraints. Note that IAR is a general method that can be applied to a wide range of unconstrained RL algorithms.\n\nWe have additional experiment still running and hope to share these during the discussion phase\n\nResponse to questions:\n\nELBO vs sandwich: Yes, indeed ELBO is consistent but we never acquire as many samples in one step (note that there is a different ELBO every state) for ELBO to reach close to the true log probability at that time step. Moreover, during training the true log probability is also changing as we update the policy using the policy gradient. Thus, unlike a standard generative modeling set-up where the ground truth distribution is fixed, we are chasing a changing ground truth - hence, any better estimate (with limited scope of samples) should intuitively provide an advantage. We have added results for the use of only ELBO in unconstrained case to the attached pdf; Figure 2 in the attached pdf shows that only ELBO still has instability issues even in the unconstrained case.\n\nWe will improve the discussion as suggested by the reviewer.\n"}]

Author rebuttal:
We thank all the reviewers for their comments and feedback. We use this space to attach the one-page pdf with extra experiments. We discuss the results in specific responses to the reviewers who asked for the experiments.

Meta-review: The reviewers unanimously agreed on acceptance.
Decision: Accept (poster)