SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We then use a weighted average of the upper and lower bounds as a low-bias estimate of  $\log \pi (a|s)$ , denoted by  $\widehat{\log p_{\theta,\psi}} = \alpha \hat{l}_{\pi} + (1 - \alpha) \hat{l}_{\pi}^{u}$  where  $\alpha$  is a hyperparameter. We call this the sandwich estimator of log probability. We observed in our experiments that an adaptive  $\alpha (\hat{l}_{\pi}, \hat{l}_{\pi}^{u})$  as a function of the two bounds provides better results than a static  $\alpha = \frac{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  $\mathcal{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  $\mathcal{C}(s)$  is the set of valid actions in state  $s$  and let  $\mathcal{C}^J(s)$  be the set of invalid actions. We sample an action from  $\pi_{\theta}(a|s)$  (recall  $\theta = (\rho, w)$  when using Flow Policy), rejecting any invalid action till a valid action is obtained. Clearly, the effective policy  $\pi'$  induced by this rejection produces only valid actions. In fact, by renormalizing we obtain  $\pi_{\theta}'(a|s) = \frac{\pi_{\theta}(a|s)}{\sum_{a_i \in \mathcal{C}(s)} \pi_{\theta}(a_i|s)}$  for  $a \in \mathcal{C}(s)$ . For the purpose of policy gradient, we need to obtain the gradient of the long term reward with  $\pi': J(\theta) = \mathbb{E}_{s \sim b_0}[V(s; \pi')]$ . We show that:
Theorem 1 With  $\pi^{\prime}$  defined as above,
$$
\nabla_ {\theta} J (\theta) = \mathbb {E} _ {\pi^ {\prime}} \left[ Q ^ {\pi^ {\prime}} (s, a) \nabla_ {\theta} \log \pi_ {\theta} (a | s) - Q ^ {\pi^ {\prime}} (s, a) \frac {\sum_ {a \in \mathcal {C} (s)} \nabla_ {\theta} \pi_ {\theta} (a | s)}{\sum_ {a \in \mathcal {C} (s)} \pi_ {\theta} (a | s)} \right] \tag {7}
$$
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  $\pi_{\theta}$ , the empirical estimate of the first gradient term is readily obtainable from trajectory samples of  $\pi'$ . To estimate the second term, for every state  $s$  in the trajectories, we sample  $S$  actions from  $\pi_{\theta}(a|s)$  and reject any  $a \in \mathcal{C}^I(s)$  to get  $l \leq S$  valid actions  $a_1, \ldots, a_l$ . We then apply Lemma 1:
Lemma 1  $\frac{1}{l}\sum_{j\in [l]}\nabla_{\theta}\log \pi_{\theta}(a_j|s)$  and  $\frac{l}{S}$  are unbiased estimates of  $\sum_{a\in \mathcal{C}(s)}\nabla_{\theta}\pi_{\theta}(a|s)$  and  $\sum_{a\in \mathcal{C}(s)}\pi_{\theta}(a|s)$  respectively.
The full approach is shown in Algorithm 2, with the changes from standard A2C highlighted in red. Recall that,  $\theta = (\rho ,w)$  is the parameter of the flow-based policy network,  $\phi$  is the parameter of the critic, and  $\psi$  is the parameter of the variational distribution network  $q_{\psi}$  for ELBO.
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
Algorithm 2: IAR-A2C
1 Initialize step count  $t\gets 1$  , Initialize episode counter  $E\gets 1$    
2 repeat   
3 Reset gradients:  $d(\theta ,\psi)\gets 0,d\phi \gets 0$  , Initialize experience set  $\tau$ $t_{start} = t$  , get state  $s_t$    
4 repeat   
5 Execute  $a_{t}$  according to policy  $\pi_{\theta}^{\prime}(a_{t}|s_{t})$  // Defined in Section 3.2   
6 Get a number of valid and sampled actions  $l_{t}$  and  $S_{t}$    
7 Receive reward  $r_t$  and new state  $s_{t + 1}$    
8 Add new experience  $\tau \leftarrow \tau \cup \{s_t,a_t,s_{t + 1},r_t\}$  , and update  $t\gets t + 1$    
9 until terminal  $s_t$  or  $t - t_{start} = t_{max}$    
10  $R = \left\{ \begin{array}{ll}0 & \text{for terminal } s_t\\ V_{\phi}(s_t) & \text{for non-terminal } s_t / / \text{Bootstrap from last state} \end{array} \right.$    
11 for  $i\in \{t - 1,\dots,t_{start}\}$  do   
12  $R\gets r_i + \gamma R$    
13  $\widehat{\log p}_{\theta ,\psi}(a_i|s_i) = \alpha \hat{l}_\pi (a_i|s_i) + (1 - \alpha)\hat{l}_\pi^u (a_i|s_i)$  // Sandwich estimator   
14 Accumulate gradients w.r.t.  $(\theta ,\psi)\colon d(\theta ,\psi)\gets$ $d(\theta ,\psi) + (R - V_{\phi}(s_{i}))\left(\widehat{\nabla_{\theta,\psi}}\widehat{\log p_{\theta,\psi}}(a_{i}|s_{i}) - \frac{S_{j}}{l_{j}^{2}}\sum_{j\in [l]}\nabla_{\theta ,\psi}\widehat{\log p_{\theta,\psi}}(a_{ij}|s_i)\right)$    
15 Accumulate gradients w.r.t.  $\phi$  ..  $d\phi \gets d\phi +\partial (R - V_{\phi}(s_i))^{2} / \partial \phi$    
16 Perform update of  $\theta$  using  $d\theta$  , of  $\psi$  using  $d\psi$  and of  $\phi$  using  $d\phi$    
17 Execute ELBO updating by running Algorithm 1 with  $\tau$  as input   
18  $E\gets E + 1$    
19 until  $E > E_{max}$
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.
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 low-dimensional 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-
Figure 3: Pistonball, goal is to move the ball to the left border by operating the pistons.
ure this as a single-agent control of all pistons. This modification presents a challenging task for the central controller as the action space is exponentially large. We show results for Acrobot and three versions of Pistonball: v1 with  $3^{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.
Figure 4: Learning curves over time in (a) Acrobot and (b) Pistonball-v[1-3] (without constraints). All approaches are trained with 5 seeds, except for Wol-DDPG, trained with 3 seeds due to its expensive training cost on Pistonball.
Figure 5: Learning curves in ERA-v[2-5] (with constraints). All settings are trained with 5 seeds.
The results of Acrobot demonstrate that the flow-based policy exhibits comparable performance to the optimal policy (A2C), which also highlights the effectiveness of our sandwich estimator for  $\log \pi(a|s)$ . In more challenging environments with high-dimensional image observations and extremely large action spaces, our model has comparable performance to Factored and AR, while significantly outperforming A2C and Wol-DDDPG. Interestingly, we observe that Wol-DDDPG struggles to learn even in the simplest Pistonball environment, despite being designed for large discrete action spaces. We hypothesize that Wol-DDDPG might function properly if the actions are discrete yet ordered, as demonstrated in the quantized continuous control task presented in the Wol-DDDPG 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
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.
Figure 6: Learning curves in ERA-Partial. Our approach obtains higher return than Factored approach.
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.
Figure 7: Learning curves for wall clock time. Our approach converges much faster than action masking.
(a)
Figure 8: (a) Ablation of gradient correction and sandwich estimator (b) Ablation of posterior type.
We conduct ablation studies of IAR-A2C 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.
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.
This research is supported by the National Research Foundation Singapore under its AI Singapore Programme (Award Number: AISG2-RP-2020-016). Dr. Nguyen was supported by grant W911NF20-1-0344 from the US Army Research Office.
[1] Abhinav Bhatia, Pradeep Varakantham, and Akshit Kumar. Resource constrained deep reinforcement learning. In Proceedings of the International Conference on Automated Planning and Scheduling, volume 29, pages 610-620, 2019.
[2] Steven Bohez, Abbas Abdolmaleki, Michael Neunert, Jonas Buchli, Nicolas Heess, and Raia Hadsell. Value constrained model-free continuous control. ArXiv preprint, abs/1902.04623, 2019. URL https://arxiv.org/abs/1902.04623.
[3] 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.
[4] 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.
[5] 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.
[6] 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.
[7] Olivier Delalleau, Maxim Peter, Eloi Alonso, and Adrien Logut. Discrete and continuous action representation for practical rl in video games. ArXiv preprint, abs/1912.11077, 2019. URL https://arxiv.org/abs/1912.11077.
[8] Adji Bousso Dieng, Dustin Tran, Rajesh Ranganath, John Paisley, and David Blei. Variational inference via chi upper bound minimization. Advances in Neural Information Processing Systems, 30, 2017.
[9] Priya 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.
[10] 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.
[11] Javier Garcia and Fernando Fernández. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research, 16(1):1437-1480, 2015.
[12] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139–144, 2020.
[13] 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.