SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chen Changyu<sup>1</sup>, Ramesha Karunasena<sup>1</sup>, Thanh Hong Nguyen<sup>2</sup>, Arunesh Sinha<sup>3</sup>, Pradeep Varakantham<sup>1</sup>
Singapore Management University<sup>1</sup>, University of Oregon<sup>2</sup>, Rutgers University<sup>3</sup>, cychen.2020@phdcs.smu.edu.sg, rameshak@smu.edu.sg,
thanhhng@cs.uoregon.edu, arunesh.sinha@rutgers.edu, pradeepv@smu.edu.sg
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<sup>1</sup>.
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
<sup>1</sup>Our implementation is available at https://github.com/cameron-chen/flow-iar.
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\mathrm{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  $\langle S, \mathcal{A}, P, r, \gamma, b_0 \rangle$ , where an agent can be in any state  $s_t \in S$  at a given time  $t$ . The agent takes an action  $a_t \in \mathcal{A}$ , causing the environment to transition to a new state  $s_{t+1}$  with a probability  $P: S \times \mathcal{A} \times S \mapsto [0,1]$ . Subsequently, the agent receives a reward  $r: S \times \mathcal{A} \mapsto \mathbb{R}$ . In the infinite-horizon setting, the discounted factor is  $0 < \gamma < 1$ . The distribution of the initial state is  $b_0$ .
In our work, we focus on a categorical action space,  $\mathcal{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  $\mathcal{C}(s) \subseteq \mathcal{A}$ . There is an oracle to answer whether an action  $a \in \mathcal{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,  $\pi(\cdot | s)$ , which generates a probability distribution over actions for state  $s$  with support only over the valid actions  $\mathcal{C}(s)$  in state  $s$ . We call the set of such stochastic policies as valid policies  $\Pi^{\mathcal{C}}$  given the per state constraint  $\mathcal{C}$ . We aim to maximize the long-term expected rewards over valid policies, that is,  $\max_{\pi \in \Pi^{\mathcal{C}}} J(\pi)$ , where  $J$  is as follows:
$$
J (\pi) = \mathbb {E} _ {s \sim b _ {0}} [ V (s; \pi) ] \text {w h e r e} V (s; \pi) = \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r \left(s _ {t}, a _ {t}\right) | s _ {0} = s; \pi \right] \tag {1}
$$
In addition, we also consider settings with partial observability where the agent observes  $o_t \in \mathcal{O}$  at time  $t$  where the observation arises from the state  $s_t$  with probability  $O: \mathcal{O} \times S \mapsto [0,1]$ . In this case, the optimal policy is a stochastic policy,  $\pi(\cdot|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  $\Pi$  that maximize the long term return, that is,  $\max_{\pi \in \Pi} J(\pi)$ .  $J(\pi)$  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  $\chi^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.
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  $\mathcal{V} = \mathcal{Z} = \mathbb{R}^d$  with densities  $p_V$  and  $p_Z$  respectively, normalizing flows [26] aim to learn a bijective and differentiable transformation  $f:\mathcal{Z}\to \mathcal{V}$ . This deterministic transformation allows us to evaluate the density at any point  $v\in \mathcal{V}$  based on the density of  $z\in \mathcal{Z}$ , as follows:
$$
p _ {V} (\boldsymbol {v}) = p _ {Z} (\boldsymbol {z}) \cdot \left| \det  \frac {\mathrm {d} \boldsymbol {z}}{\mathrm {d} \boldsymbol {v}} \right|, \quad \boldsymbol {v} = f (\boldsymbol {z}) \tag {2}
$$
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  $\pmb{v} \in \mathbb{R}^{D \times M}$  and  $\pmb{x} \in \{1, \dots, M\}^D$ . The argmax operation is interpreted as a surjective flow layer  $\pmb{v} \mapsto \pmb{x}$ , which is deterministic in one direction ( $x_d = \arg \max_m \pmb{v}_d$ , written compactly as  $\pmb{x} = \arg \max_{d} \pmb{v}$ ) and stochastic in the other ( $\pmb{v} \sim q(\cdot | \pmb{x})$ ). With this interpretation, the argmax operation can be considered a probabilistic right-inverse in the latent variable model expressed by:
$$
P (\boldsymbol {x}) = \int P (\boldsymbol {x} | \boldsymbol {v}) p (\boldsymbol {v}) \mathrm {d} \boldsymbol {v}, \quad P (\boldsymbol {x} | \boldsymbol {v}) = \delta (\boldsymbol {x} = \operatorname {a r g m a x} (\boldsymbol {v})) \tag {3}
$$
where  $\operatorname{argmax}$  is applied in the last dimension of  $\pmb{v}$ . In this scenario, the density model  $p(\pmb{v})$  is modeled using a normalizing flow. The learning process involves introducing a variational distribution  $q(\pmb{v}|\pmb{x})$ , which models the probabilistic right-inverse for the  $\arg\max$  surjection, and optimizing the evidence lower bound (ELBO), which is the RHS of the following inequality:
$$
\log P (\boldsymbol {x}) \geq \mathbb {E} _ {\boldsymbol {v} \sim q (\cdot | \boldsymbol {x})} [ \log P (\boldsymbol {x} | \boldsymbol {v}) + \log p (\boldsymbol {v}) - \log q (\boldsymbol {v} | \boldsymbol {x}) ] = \mathbb {E} _ {\boldsymbol {v} \sim q (\cdot | \boldsymbol {x})} [ \log p (\boldsymbol {v}) - \log q (\boldsymbol {v} | \boldsymbol {x}) ] = \mathcal {L}
$$
The last equality holds under the constraint that the support of  $q(\pmb{v}|\pmb{x})$  is enforced to be only over the region  $S = \{\pmb{v} \in \mathbb{R}^{D \times M} : \pmb{x} = \arg \max \pmb{v}\}$  which ensures that  $P(\pmb{x}|\pmb{v}) = 1$ . From standard variational inference results,  $\log P(\pmb{x}) - \mathcal{L} = \mathrm{KL}(q(\pmb{v}|\pmb{x})||p(\pmb{v}|\pmb{x}))$ , which also approaches 0 as the approximate posterior  $q(\pmb{v}|\pmb{x})$  comes closer to the true posterior  $p(\pmb{v}|\pmb{x})$  over the training time.
$\chi^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  $\mathrm{KL}(q||p)$  is employed to measure closeness, where  $q(\boldsymbol{v}|\boldsymbol{x})$  represents a variational family. This approach yields ELBO of the evidence  $\log P(\boldsymbol{x})$  as described above.
Instead of using KL divergence, the authors in [8] suggest an alternative of  $\chi^2$ -divergence to measure the closeness. As a result, they derived an upper bound of the evidence, known as CUBO:  $\log P(\pmb{x}) \leq \frac{1}{2} \log \mathbb{E}_{\pmb{v} \sim q(\cdot | \pmb{x})} \left[ \left( \frac{p(\pmb{x}, \pmb{v})}{q(\pmb{v} | \pmb{x})} \right)^2 \right]$ . Similar to the ELBO in Argmax Flow, CUBO can be further simplified under the constraint that the support of  $q(\cdot | \pmb{x})$  is restricted to the region  $\mathcal{S}$ :
$$
\log P (\boldsymbol {x}) \leq \frac {1}{2} \log \mathbb {E} _ {\boldsymbol {v} \sim q (\cdot | \boldsymbol {x})} \left[ \left(\frac {p (\boldsymbol {v})}{q (\boldsymbol {v} | \boldsymbol {x})}\right) ^ {2} \right] = \mathcal {L} _ {\chi^ {2}} \tag {4}
$$
Also,  $\mathcal{L}_{\chi^2} - \log P(\pmb{x}) = \frac{1}{2}\log (1 + D_{\chi^2}(p(\pmb{v}|\pmb{x}))||q(\pmb{v}|\pmb{x})))$ , hence the gap between the ELBO and CUBO approaches 0 as the approximate posterior  $q(\cdot |\pmb{x})$  becomes closer to the true posterior  $p(\cdot |\pmb{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
Figure 1: Our IAR-A2C framework. At each time step, an 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.
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  $\pi (\cdot |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)}\sim \pi (\cdot |s)$  and probability value  $\pi (a^{(i)}|s) > 0$
Figure 2: Composition of a conditional flow  $p_{\rho, w}(z_K | s)$  and argmax transformation resulting in the policy  $\pi_{\rho, w}(a | s)$ . The flow maps from a base distribution  $p_{\rho}(z_0 | s)$  by using a bijection  $F_w$ . The diagram is adapted from Hoogeboom et al. [16].
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-the-art 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  $\pi$ . 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
1 Input: Invertible flow  $F_{w} = f_{w,k} \circ \ldots \circ f_{w,1}$ , State encoder  $E_{\rho}$ , Posterior  $q_{\psi}$ , rollout  $\tau$   
2 Sample  $States = \{(s^{(i)})\}_{i=1}^{B} \sim \tau$   
3 for  $s \in States$  do  
4 for  $j \gets 1$ : n_elbo_steps do  
5  $\begin{array}{r}z^{(j)} = F_w(z_0^{(j)})\mathrm{,where}z_0^{(j)}\sim E_\rho (s)\\ a^{(j)} = \operatorname {argmax}z^{(j)} \end{array}$   
6  
7 for  $j \gets 1$ : n_elbo_steps do  
8  $\begin{array}{r}z^{\prime (jn)} = \mathrm{threshold}_T(u^{(jn)})\mathrm{,where}\{(u^{(jn)})\}_{n = 1}^N\sim q_\psi (\cdot |a^{(j)},s)\\ z_0^{\prime (jn)} = F_w^{-1}(z^{\prime (jn)}) \end{array}$   
9  
10 Take a gradient ascending step of ELBO w.r.t.  $\rho ,w$  and  $\psi$   
11  $\mathcal{L} = \frac{1}{N}\sum_{i = 1}^{N}\left(\log p_{\rho}(z_0^{'(jn)}|s) - \sum_{k = 1}^{K}\log \left|\det \frac{\partial f_{w,k}}{\partial z_{k - 1}^{'(jn)}}\right| - \log q_{\psi}(z^{'(jn)}|a^{(j)},s)\right)$
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,  $\pi(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  $\mu, \sigma$  defined by the state encoder (a state-dependent neural network):  $z_0 = \mu_{\rho}(s) + \sigma_{\rho}(s) \cdot \epsilon$  where  $\epsilon \sim \mathcal{N}(0,1)$ . We write this as  $z_0 \sim E_{\rho}(s)$ , where  $\rho$  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 = \operatorname{argmax}\left(f_K \circ \ldots \circ f_1(z_0)\right)$ . Each  $f_k$  is an invertible neural network [26] and we use  $F_w = f_{w,K} \circ \ldots \circ f_{w,1}$  to denote the composed function, where  $w, k$  is the weight of the network representing  $f_k$ . Thus, sampled action  $a = \operatorname{argmax}\left(F_w(z_0)\right)$  for  $z_0 \sim E_{\rho}(s)$ . We use shorthand  $\theta = (\rho, w)$  when needed. Also, we use  $z_k = f_{w,k}(z_{k-1})$  to denote the output of  $f_{w,k}$  and  $p_{\rho}(z_0|s)$  to denote the probability density of  $z_0 \sim \mathcal{N}(\mu_{\rho}, \sigma_{\rho})$ .
We use a variational distribution  $q_{\psi}$  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  $\hat{l}_{\pi}$ , is the evidence lower bound ELBO, computed as follows:
$$
\hat {l} _ {\pi} (a | s) = \mathcal {L} = \mathbb {E} _ {z _ {K} \sim q _ {\psi} (\cdot | a, s)} \left[ \log p _ {\rho} (z _ {0} | s) - \sum_ {k = 1} ^ {K} \left(\log \left| \det  \frac {\partial f _ {k}}{\partial z _ {k - 1}} \right|\right) - \log q _ {\psi} (z _ {K} | a, s) \right] \tag {5}
$$
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  $(\rho, w)$  to make the distribution of overall policy be closer to this target distribution and to improve the posterior  $q_{\psi}$  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_{\psi}$  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 \pi (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  $\hat{l}_{\pi}$  is a biased estimator of  $\log \pi (a|s)$ , but it is known in literature [5] that  $\hat{l}_{\pi}$  is a consistent estimator (i.e., converges to  $\log \pi (a|s)$  in the limit) and there are generative methods based on this consistency property, such as [33]. However, we aim to use  $\hat{l}_{\pi}$  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 \pi (a|s)$  using the following CUBO:
$$
\hat {l} _ {\pi} ^ {u} (a | s) = \mathcal {L} _ {\chi^ {2}} = \frac {1}{2} \log \mathbb {E} _ {z _ {K} \sim q _ {\psi} (z _ {K} | a, s)} \left[ \left(\frac {p _ {\rho} (z _ {0} | s)}{\prod_ {k = 1} ^ {K} \left| \det  \frac {\partial f _ {w , k}}{\partial z _ {k - 1}} \right| q _ {\psi} (z _ {K} | a , s)}\right) ^ {2} \right] \tag {6}
$$