Title: Generative Modelling of Stochastic Actions with Arbitrary Constraints in Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Action constrained reinforcement learning

Normalizing flow

Generative modelling
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

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 1 .

1 Introduction

Adaptive resource allocation problems with multiple resource types (e.g., fire trucks, ambulances, police vehicles) are ubiquitous in the real world [13, 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/transportation, 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, as 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 RL for constrained action spaces has considered resource allocation problems with a single type of resource, thereby introducing order in action space [1]. In such ordered action spaces, actions can be converted into continuous space and this allows for the usage of continuous action RL methods (e.g., DDPG). In problems of interest in this paper, we are interested in problems

with multiple resource types (e.g., fire truck, ambulance, police). These problems have a large action space that is 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 [4]. Thus, we 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 in 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 or 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 a 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.

Problem Statement: A Markov Decision Process (MDP) is represented by the tuple ⟨S , A , P, r, γ, b 0 ⟩ , where an agent can be in any state s t ∈ S at a given time t . The agent takes an action a t ∈ A , causing the environment to transition to a new state s t +1 with a probability P : S × A × S ↦→ [0 , 1] . Subsequently, the agent receives a reward r : S × A ↦→ R . In the infinite-horizon setting, the discounted factor is 0 < γ < 1 . The distribution of the initial state is b 0 .

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 ) ⊆ A . There is an oracle to answer whether an action a ∈ 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 a stochastic policy, π ( ·| s ) , which generates a probability 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 Π C given the per state constraint C . We aim to maximize the long-term expected rewards over valid policies, that is, max π ∈ Π C J ( π ) , where J is as follows:

In addition, we also consider settings with partial observability where the agent observes o t ∈ O at time t where the observation arises from the state s t with probability O : O×S ↦→ [0 , 1] . In this case, the optimal policy is a stochastic policy, π ( ·| h ) , 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 Π that maximize the long term return, that is, max π ∈ Π J ( π ) . J ( π ) is the same as stated in Equation 1 but where the expectation is also over the observation probability distribution in addition to the standard transition and stochastic policy probability distributions.

Contribution: We propose two key innovations to address the problem above. First , we present a conditional Normalizing Flow-based [26] 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 χ 2 divergence. Third , we propose a policy gradient approach which is able to reject invalid actions (that do not satisfy constraints) referred to as Invalid Action Rejection Advantage Actor-Critic (IAR-A2C). IAR-A2C queries the constraint oracle and ensures validity of actions in every state (in fully observable setting) by rejecting all invalid actions. We derive a new policy gradient estimator

for IAR-A2C. Figure 1 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: Normalizing flows are a family of generative models that can provide both efficient sampling and density estimation. Their main idea is to construct a series of invertible and differentiable mappings that allow transforming a simple probability distribution into a more complex one. Given V = Z = R d with densities p V and p Z 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 v ∈ V based on the density of z ∈ Z , as follows:

In this context, p Z can be any density, though it is typically chosen as a standard Gaussian and f is represented by a neural network. Consequently, normalizing flows offer a powerful tractable framework for learning complex density functions. However, the density estimation presented in Equation 2 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 on the output of continuous flows. Let's consider v ∈ R D × M and x ∈ { 1 , . . . , M } D . The argmax operation is interpreted as a surjective flow layer v ↦→ x , which is deterministic in one direction ( x d = arg max m v d , written compactly as x = arg max v ) and stochastic in the other ( v ∼ q ( ·| x ) ). With this interpretation, the argmax operation can be considered a probabilistic right-inverse in the latent variable model expressed by:

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 ( v | x ) , 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:

The last equality holds under the constraint that the support of q ( v | x ) is enforced to be only over the region S = { v ∈ R D × M : x = arg max v } which ensures that P ( x | v ) = 1 . From standard variational inference results, log P ( x ) -L = KL ( q ( v | x ) || p ( v | x )) , which also approaches 0 as the approximate posterior q ( v | x ) comes closer to the true posterior p ( v | x ) over the training time.

χ 2 Upper Bound: 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 q ( v | x ) represents a variational family. This approach yields ELBO of the evidence log P ( x ) as described above.

Instead of using KL divergence, the authors in [8] suggest an alternative of χ 2 -divergence to measure the closeness. As a result, they derived an upper bound of the evidence, known as CUBO: log P ( x ) ≤

1 2 log E v ∼ q ( ·| x ) [ ( p ( x , v ) q ( v | x ) ) 2 ] . Similar to the ELBO in Argmax Flow, CUBO can be further simplified under the constraint that the support of q ( ·| x ) is restricted to the region S :

Also, L χ 2 -log P ( x ) = 1 2 log(1 + D χ 2 ( p ( v | x )) || q ( v | x ))) , hence the gap between the ELBO and CUBO approaches 0 as the approximate posterior q ( ·| x ) becomes closer to the true posterior p ( ·| x ) .

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

Flow

Sampleactions

Policy

from Flow Policy

Reject

invalid actions

Actions

Uniformly

sampleone

Policy

update

Invalid

Valid

Action

Actions

Actions

Learning,algorithm

Reward

Action

Environment

Observation

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 as 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 is to model the entire policy, denoted as π ( ·| s ) , which allows us to sample actions and obtain their corresponding probabilities. The policy is parameterized using a 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 a better policy representation, as it is sufficient to require samples from the support set of the policy, represented by a ( i ) ∼ π ( ·| s ) and probability value π ( a ( i ) | s ) > 0 .

'2

a

=

argmaxzk

F-1

n

2K~q(2K|a,s)

1

2

3

Pp(z0ls)

Pp,w(2k∣8)

πp,w(a|s)

Based on this observation, our first contribution is a compact policy representation using Argmax Flow [16], which we refer to as the Flow Policy . Argmax Flow is a state-of-theart discrete normalizing flow model and it has shown great capability of learning categorical data such as in sentence generation. In our context, Flow Policy will output the action (a sample of the flow policy) and its probability, 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 normalizing flow model transforms a base distribution, given by a random variable z 0 , to the desired distribution. In our approach, the desired distribution is the distribution given by policy π . We adapt Argmax Flow approach for learning a stochastic policy.

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) [12] and Diffusion Models [15, 29]. GANs, categorized as implicit generative models, do not allow for an estimation of the data's log probability. While prior research has successfully devised an approach for exact log probability estimation with Diffusion Models, these

Algorithm 1: ELBO Optimization

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, as we find the simple flow model is sufficient to learn a good policy. So normalizing flow is a natural choice to construct our policy.

We demonstrate how to construct the flow policy within the A2C framework. An illustration is provided in Figure 2. To define the policy, π ( a | s ) , we first encode the state that is then used to define the base distribution z 0 . In accordance with standard practice, we select z 0 as a Gaussian distribution with parameters µ, σ defined by the state encoder (a state-dependent neural network): z 0 = µ ρ ( s ) + σ ρ ( s ) · ϵ where ϵ ∼ N (0 , 1) . We write this as z 0 ∼ E ρ ( s ) , where ρ denotes the weights of the state encoder. We then apply a series of invertible transformations given as functions f k that define the flow, followed by the argmax operation (as defined in background on Argmax flow). Consequently, the final sampled action is given by a = argmax( f K ◦ . . . ◦ f 1 ( z 0 )) . Each f k is an invertible neural network [26] and we use F w = f w,K ◦ . . . ◦ f w, 1 to denote the composed function, where w,k is the weight of the network representing f k . Thus, sampled action a = argmax( F w ( z 0 )) for z 0 ∼ E ρ ( s ) . We use shorthand θ = ( ρ, w ) when needed. Also, we use z k = f w,k ( z k -1 ) to denote the output of f w,k and p ρ ( z 0 | s ) to denote the probability density of z 0 ∼ N ( µ ρ , σ ρ ) .

We use a variational distribution q ψ for the reverse mapping from a discrete action (conditional on state s ) to a continuous output of F w ( z 0 ) . The corresponding estimate of the log probability of the sampled action a , denoted by ˆ l π , is the evidence lower bound ELBO, computed as follows:

To ensure that our approximation closely approximates the evidence, we optimize ELBO progressively, following the training scheme of Argmax flow, as shown in the subroutine Algorithm 1 used within the outer A2C framework in Algorithm 2. The target distribution in this subroutine is given by actions a ( j ) sampled from the current policy (line 6); thus, this subroutine aims to update the flow network components ( ρ, w ) to make the distribution of overall policy be closer to this target distribution and to improve the posterior q ψ 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 ( j ) back through the invertible flow layers (using q ψ for the discrete to 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 1; 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 π ( a | s ) , thus improving the convergence of policy gradient learning. Moreover, to tackle the challenge of aforementioned complex constraints, our algorithm proceeds by sampling a number of actions from

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 ˆ l π is a biased estimator of log π ( a | s ) , but it is known in literature [5] that ˆ l π is a consistent estimator (i.e., converges to log π ( a | s ) in the limit) and there are generative methods based on this consistency property, such as [33]. However, we aim to use ˆ l π in the policy gradient and stochastic gradient descent typically uses an 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 π ( a | s ) using the following CUBO:

We then use a weighted average of the upper and lower bounds as a low-bias estimate of log π ( a | s ) , denoted by ̂ log p θ,ψ = α ˆ l π + (1 -α ) ˆ l u π where α is a hyperparameter. We call this the sandwich estimator of log probability. We observed in our experiments that an adaptive α ( ˆ l π , ˆ l u π ) as a function of the two bounds provides better results than a static α = 1 2 (see Appendix B.1 for more details).

Constraints: Since the agent only has oracle access to the constraints, existing safe exploration approaches [20, 24] are not directly applicable to 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-dependent constraints in the full observation setting, we propose a new policy gradient algorithm called Invalid Action Rejection Advantage Actor-Critic (IAR-A2C). IAR-A2C operates by sampling a 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 I ( s ) be the set of invalid actions. We sample an action from π θ ( a | s ) (recall θ = ( ρ, w ) when using Flow Policy), rejecting any invalid action till a valid action is obtained. Clearly, the effective policy π ′ induced by this rejection produces only valid actions. In fact, by renormalizing we obtain π ′ θ ( a | s ) = π θ ( a | s ) ∑ π θ ( a i | s ) for a ∈ C ( s ) .

∈C

(

s

)

For the purpose of policy gradient, we need to obtain the gradient of the long term reward with π ′ : J ( θ ) = E s ∼ b 0 [ V ( s ; π ′ )] . We show that:

Theorem 1 With π ′ defined as above,

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 1. Then, given the network π θ , the empirical estimate of the first gradient term is readily obtainable from trajectory samples of π ′ . To estimate the second term, for every state s in the trajectories, we sample S actions from π θ ( a | s ) and reject any a ∈ C I ( s ) to get l ≤ S valid actions a 1 , ..., a l . We then apply Lemma 1:

Lemma 1 1 l ∑ j ∈ [ l ] ∇ θ log π θ ( a j | s ) and l S are unbiased estimates of ∑ a ∈C ( s ) ∇ θ π θ ( a | s ) and ∑ a ∈C ( s ) π θ ( a | s ) respectively.

The full approach is shown in Algorithm 2, with the changes from standard A2C highlighted in red. Recall that, θ = ( ρ, w ) is the parameter of the flow-based policy network, ϕ is the parameter of the critic, and ψ is the parameter of the variational distribution network q ψ for ELBO.

4 Related Work

Large discrete action space. Dulac-Arnold et al. [10] attempt to solve this problem by embedding discrete actions into a continuous space; however, this approach does not work for our unordered a i

Algorithm 2: IAR-A2C

discrete space as we reveal by thorough comparison in our experiments. 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 [30] apply this approach to 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, 35] 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 [34, 36] 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.

Safe actions. Our approach falls into the field of constrained action space in RL. Situated in the broader safe RL literature [11, 14], these works aim to constrain actions at 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 [11] cannot solve our problem. Also, our state dependent constraints if written as 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 as mathematical inequalities and even then cannot handle state dependent constraints. Although action masking [17] does offer a solution for state-dependent 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 differentiating through an optimization; these methods are 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 et al. [18] proposed to improve the control stability by a novel normalizing-flow 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 a discrete flow, extending the scope of flow policies into discrete action spaces.

5 Experiments

Through our experiments, we aim to address two main research questions related to 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 our approach.

We first describe our set-up. We evaluate IAR-A2C against prior works across a diverse set of environments, including lowdimensional discrete control tasks such as CartPole and Acrobot [3], the visually challenging Pistonball task [31] with high-dimensional image inputs and an extremely large action space (upto 59 , 049 categories), and an emergency resource allocation simulator in a city, referred to as Emergency Resource Allocation (ERA). CartPole and Acrobot are well-established environments. Pistonball is also a standard environment where a series of pistons on a line needs to move (up, down, or stay) to move a ball from right to left (Figure 3). While the Pistonball environment was originally designed for a multi-agent setting with each piston controlled by a distinct agent, we reconfig-

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 5 , v2 with 3 8 , and v3 with 3 10 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 a 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 to them as quickly as possible. Moreover, we explore a partially observable scenario (with no constraints) in which the optimal allocation is randomized, thus, the next node for a resource is 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 7 3 actions, v3 with 8 3 actions, v4 with 9 3 actions, and v5 with 10 3 actions in constrained scenario. The results for ERA-v1 are in Appendix C.

Benchmarks : In the unconstrained setting, we compare our approach to Wol-DDPG [10], A2C [22], factored representation (Factored, discussed in Section 4), 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 our 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 is 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, we 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 C.

1e2

Acrobot-v1

1e2

Pistonball-v1

le2

Pistonball-v2

le2

Pistonball-v3

1.00

1.00

1.00

Retur

tur

0.75

E

n

0.75

tur

0.75

-2

R

R

verage

e

e

0.50

Average

e

0.50

Average

e

0.50

一

3

Aver

0.25

0.25

0.25

I

0.00

0.00

0.00

4

0.0

0.5

1.0

0.0

0.5

1.0

0.0

0.5

1.0

0.0

0.5

1.0

Timesteps

1e6

Timesteps

1e6

Timesteps

1e6

Timesteps

1e6

Ours

A2C

Wol-DDPG

Factored

AR

lel

ERA-V2

lel

ERA-V3

lel

ERA-V4

lel

ERA-v5

Return

-2

Return

2

Return

2

Return

2

Average

-4

Average！

-4

Average

-4

Average

-4

-6

-6

-6

-6

-8

-8

-8

-8

0

2

4

6

0

2

4

6

0

2

4

6

0

2

4

6

Timesteps

1e5

Timesteps

1e5

Timesteps

1e5

Timesteps

1e5

Ours

Wol-DDPG

MASK

AR+IAR

The results of Acrobot demonstrate that the flow-based policy exhibits comparable performance to the optimal policy (A2C), which also highlights the effectiveness of our sandwich estimator for log π ( a | s ) . In more challenging environments with highdimensional 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 designed for large discrete action spaces. We hypothesize that Wol-DDPG might function properly if the actions are discrete yet ordered, as demonstrated in the quantized continuous control task presented in the Wol-DDPG paper [10].

On ERA-Partial, we demonstrate the known fact that the optimal policy in environments with partial observability may be a stochastic one [28]. We compare with the factored approach (Figure 6). In this environment, the optimal stochastic policy

1e2

ERA-Partial

4

Return

Average

2

0

Ours

Factored

0.0

0.5

1.0

1.5

Timesteps

1e6

is a joint distribution that is not factorizable into a product of 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 a 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 5. We observe that our approach demonstrates better or comparable performance compared to the benchmarks. AR is 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 it 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 D.

Also, action masking achieves its performance by querying all actions of the entire action space, whereas our approach only requires querying a batch of actions, which is substantially smaller(e.g., 64 versus 1 , 000 for ERA-v5). Thus, while Figure 5 shows IAR-A2C taking more iteration for convergence, the same figure when drawn with the x-axis as wall clock time in Figure 7 (shown only for v4 here, others are in appendix C) shows an order of magnitude faster convergence in wall clock time. Another critical property of our approach is the guaranteed absence of constraint violations, similar to the action masking method. However, while action masking demands the full knowledge of the validity of all actions, our method merely requires the validity of the sampled actions within a 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 C.

5.3 Ablation Study

(a)

(b)

ERA-V4

lel

ERA-V4

1.0

S

Full model

rate:

0.8

eturn

w/o flow-q

-4

0.6

ion

e

0.4

vera

g

9-

0.2

Full model

Val

w/osandwich est

-8

0.0

0

2

4

6

0

2

4

6

Timesteps

1e5

Timesteps

1e5

lel

ERA-V4

Return

一

4

Average

MWMN

-6

MASK

Ours

-8

0

20

40

Wall time (h)

We conduct ablation studies of IARA2C on ERA-v4 to investigate the effect of various choices of modules and estimators in our approach.

Policy Gradient : We compare the performance of approaches using the policy gradient provided by Theorem 1 (gradient correction) and the original policy gradient of A2C (standard policy gradient), while still applying invalid action rejection. We observe in Figure 8a that the number of valid actions in the batch decreases rapidly, and the program may crash if no valid actions are available.

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 to a reduction in valid actions in Figure 8a and unstable learning as a consequence, similar to the observation when using the standard policy gradient.

Posterior Type : The posterior q ( z | a, s ) can be modeled by a conditional Gaussian or a normalizing flow. In our experiments, we discover that modelling q with a flow posterior is crucial for learning, as it can approximate the true posterior more effectively than a Gaussian, as seen in Figure 8b.

6 Conclusion and Limitations

We introduced a 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 of actions is also a promising research direction for improved performance.

References

Abhinav Bhatia, Pradeep Varakantham, and Akshat Kumar. Resource constrained deep reinforcement learning. In Proceedings of the International Conference on Automated Planning and Scheduling , volume 29, pages 610-620, 2019.

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. URL https://arxiv.org/abs/1902.04623 .

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. ArXiv preprint , abs/1606.01540, 2016. URL https://arxiv.org/abs/1606.01540 .

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.

Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In Yoshua Bengio and Yann LeCun, editors, 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings , 2016. URL http://arxiv.org/abs/1509.00519 .

Jie Chen and Ronny Luss. Stochastic gradient descent with biased but consistent gradient estimators. ArXiv preprint , abs/1807.11880, 2018. URL https://arxiv.org/abs/1807. 11880 .

Olivier Delalleau, Maxim Peter, Eloi Alonso, and Adrien Logut. Discrete and continuous action representation for practical rl in video games. ArXiv preprint , abs/1912.11077, 2019. URL https://arxiv.org/abs/1912.11077 .

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.

Priya L. Donti, David Rolnick, and J. Zico Kolter. DC3: A learning method for optimization with hard constraints. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https://openreview. net/forum?id=V1ZHVxJ6dSS .

Gabriel Dulac-Arnold, Richard Evans, Hado van 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 https://arxiv.org/abs/1512.07679 .

Javier Garcıa and Fernando Fernández. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research , 16(1):1437-1480, 2015.

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(11):139-144, 2020.

Sergei Pavlovich Grachev, Aleksei Aleksandrovich Zhilyaev, Vladimir Borisovich Laryukhin, Dmitrii Evgen'evich Novichkov, Vladimir Andreevich Galuzin, Elena V Simonova, IV Maiyorov, 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.

Shangding Gu, Long Yang, Yali Du, Guang Chen, Florian Walter, Jun Wang, Yaodong Yang, and Alois Knoll. A review of safe reinforcement learning: Methods, theory and applications. ArXiv preprint , abs/2205.10330, 2022. URL https://arxiv.org/abs/2205.10330 .

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and HsuanTien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 612, 2020, virtual , 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html .

Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. In Marc'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 12454-12465, 2021. URL https://proceedings.neurips.cc/paper/2021/hash/ 67d96d458abdef21792e6d8e590244e7-Abstract.html .

Shengyi Huang and Santiago Ontañón. A closer look at invalid action masking in policy gradient algorithms. In The International FLAIRS Conference Proceedings , volume 35, 2022.

Shahbaz Abdul Khader, Hang Yin, Pietro Falco, and Danica Kragic. Learning stable normalizingflow control for robotic manipulation. In 2021 IEEE International Conference on Robotics and Automation (ICRA) , pages 1644-1650. IEEE, 2021.

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.

Jyun-Li Lin, Wei Hung, Shang-Hsuan Yang, Ping-Chun Hsieh, and Xi Liu. Escaping from zero gradient: Revisiting action-constrained reinforcement learning via frank-wolfe policy optimization. In Cassio P. de Campos, Marloes H. Maathuis, and Erik Quaeghebeur, editors, Proceedings of the Thirty-Seventh Conference on Uncertainty in Artificial Intelligence, UAI 2021, Virtual Event, 27-30 July 2021 , volume 161 of Proceedings of Machine Learning Research , pages 397407. AUAI Press, 2021. URL https://proceedings.mlr.press/v161/lin21b.html .

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.

Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement 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. JMLR.org, 2016. URL http://proceedings.mlr.press/v48/mniha16.html .

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 management. Accident Analysis & Prevention , 165:106501, 2022.

Tu-Hoa Pham, Giovanni De Magistris, and Ryuki Tachibana. Optlayer-practical constrained optimization for deep reinforcement learning in the real world. In 2018 IEEE International Conference on Robotics and Automation (ICRA) , pages 6236-6243. IEEE, 2018.

Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research , 22(268):1-8, 2021. URL http://jmlr.org/papers/v22/ 20-1364.html .

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, ICML 2015, Lille, France, 6-11 July 2015 , volume 37 of JMLR Workshop and Conference Proceedings , pages 1530-1538. JMLR.org, 2015. URL http://proceedings. mlr.press/v37/rezende15.html .

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 ThirtySecond Innovative Applications of Artificial Intelligence Conference, February 7-12, 2020 , pages 2226-2235. AAAI Press, 2020. URL https://aaai.org/ojs/index.php/AAAI/ article/view/5599 .

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.

Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 37, 2021 . OpenReview.net, 2021. URL https://openreview.net/forum?id=PxTIG12RRHS .

Yunhao Tang and Shipra Agrawal. Discretizing continuous action space for on-policy optimization. In The Thirty-Fourth AAAI Conference on 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://aaai.org/ojs/index.php/AAAI/article/view/6059 .

Justin K. Terry, Benjamin Black, Nathaniel Grammel, Mario Jayakumar, Ananth Hari, Ryan Sullivan, Luis S. Santos, Clemens Dieffendahl, Caroline Horsch, Rodrigo Perez-Vicente, Niall L. Williams, Yashas Lokesh, and Praveen Ravi. Pettingzoo: Gym for multi-agent reinforcement learning. In Marc'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 https://proceedings.neurips.cc/ paper/2021/hash/7ed2d3454c5eea71148b11d0c25104ff-Abstract.html .

Ivan B Vermeulen, Sander M Bohte, Sylvia G Elkhuizen, Han Lameris, Piet JM Bakker, and Han La Poutré. Adaptive resource allocation for efficient patient scheduling. Artificial intelligence in medicine , 46(1):67-80, 2009.

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.org/abs/1906.02771 .

Mason Wright, Yongzhao Wang, and Michael P Wellman. Iterated deep reinforcement learning in games: History-aware training for improved stability. In EC , pages 617-636, 2019.

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.

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.00589, 2018. URL https://arxiv.org/abs/1806.00589 .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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\nWEAKNESSES:\nFirst 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\nQUESTIONS:\nSome 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.\n\nLIMITATIONS:\nIn 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).", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThere 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.\n\nQUESTIONS:\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\nLIMITATIONS:\nRefer 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe problem that the submission investigates is both important and underexplored. The submission's approach seems reasonable and principled.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n> 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\nLIMITATIONS:\n> Have the authors adequately addressed the limitations\n\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nOverall 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.\n\nWEAKNESSES:\nOverall 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nOverall, 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}