SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Janaka Chathuranga Brahmanage, Jiajing Ling, Akshit Kumar
{janakat.2022, jjling.2018}@phdcs.smu.edu.sg, akshatkumar@smu.edu.sg
Action-constrained reinforcement learning (ACRL) is a popular approach for solving safety-critical and resource-allocation related decision making problems. A major challenge in ACRL is to ensure agent taking a valid action satisfying constraints in each RL step. Commonly used approach of using a projection layer on top of the policy network requires solving an optimization program which can result in longer training time, slow convergence, and zero gradient problem. To address this, first we use a normalizing flow model to learn an invertible, differentiable mapping between the feasible action space and the support of a simple distribution on a latent variable, such as Gaussian. Second, learning the flow model requires sampling from the feasible action space, which is also challenging. We develop multiple methods, based on Hamiltonian Monte-Carlo and probabilistic sentential decision diagrams for such action sampling for convex and non-convex constraints. Third, we integrate the learned normalizing flow with the DDPG algorithm. By design, a well-trained normalizing flow will transform policy output into a valid action without requiring an optimization solver. Empirically, our approach results in significantly fewer constraint violations (upto an order-of-magnitude for several instances) and is multiple times faster on a variety of continuous control tasks.
Action-constrained reinforcement learning (ACRL), where an agent's action taken in each RL step should satisfy specified constraints, has been applied to solve a diverse range of real-world applications [18]. One example is the resource allocation problem in supply-demand matching [4]. The allocation of resources must satisfy constraints such as total assigned resources should be within a specific upper and lower limit. In robotics, kinematics constraints (e.g., limits on velocity, acceleration, torque among others) can be modeled effectively using action constraints [1, 13, 26, 32].
One of the main challenges in solving ACRL is to ensure that during training and evaluation, all the action constraints are satisfied at each time step, while simultaneously improving the policy. Lagrangian Relaxation (LR) is a popular approach for solving constrained RL problems [30], however directly applying LR to ACRL is impractical due to the difficulty in defining cost functions  $(c(s,a))$  that penalize infeasible actions for all  $(s,a)$ , and the inability to guarantee zero constraint violation during training and policy execution. Also, there are studies addressing ACRL with discrete action spaces [19, 20]. As they use propositional logic, it is not clear how to extend them to continuous action spaces. Recently proposed approaches aim to satisfy continuous action constraints using an action projection step. One natural approach is to add a differentiable projection layer at the end of the policy network to satisfy the constraints [2, 4, 7, 26]. This projection layer projects the unconstrained policy action onto the feasible action space by solving a quadratic program (QP). However, this approach has two primary limitations. Firstly, solving a QP at each RL step can be
computationally expensive, particularly when dealing with non-convex action constraints or large action spaces. Secondly, the coupling between the policy network and projection layer can lead to a potential issue of zero gradient during end-to-end training, which can undermine the effectiveness of the projection layer approach, especially in the early training stages, as highlighted by recent studies [18]. This issue occurs when the policy output is far outside the feasible action space, thereby any small change in the policy parameters does not result in any change in the projected action.
To overcome the zero gradient issue, an ACRL algorithm has been proposed in [18], which decouples policy gradients from the projection layer. The policy parameters are updated by following the direction found by the Frank-Wolf method [10], which is within the feasible action space. However, this approach can also be computationally expensive, as solving a QP is still required during the update of policy parameters and during policy execution. To summarize, while there are several proposed approaches for solving ACRL with continuous action space, all of them share the common drawback of requiring QP solving during either training or action execution, or both. This can cause a significant reduction in training speed, which we also validate empirically. Moreover, some of these approaches also suffer from the zero gradient issue, which can make training sample inefficient. Therefore, there is need for more efficient methods to tackle ACRL with continuous action spaces.
Often, action constraints are specified analytically using features from state and actions (e.g.,  $x^{2} + y^{2} < = 1$  where  $x, y$  can be features). Our key idea is to exploit given action constraints as domain knowledge and develop an effective, compact representation of all valid actions that can be integrated with RL algorithms. To find such a representation, generative models provide an attractive solution since they can generate valid actions by learning from a finite number of valid actions (which can be provided by sampling from the valid action space). This allows for a compact feasible action space representation even for continuous action spaces. Furthermore, generative models can be easily integrated with RL algorithms for end-to-end training, as shown in previous work [22, 33], even though they are not specifically solving ACRL problems. The idea of learning the representation of all valid actions also shares a connection with offline RL, where invalid actions are often treated as out-of-distribution (OOD) actions [12, 35]. However, a notable difference between our work and offline RL lies in the availability of a dataset comprising valid actions. In offline RL, such a dataset is assumed to be available (e.g., collected from a random policy or an expert). In ACRL, collecting data even from a random policy is not straightforward as random policy must also select actions uniformly from the feasible action space, which itself is the key problem we address in this work.
- We utilize the normalizing flows, a type of generative model, to learn a differentiable, invertible mapping between the data distribution of sampled valid actions and a simple latent distribution (e.g., Gaussian, uniform distribution) [9, 27]. When a member of the latent distribution is given, the flow model can transform it into an element that conforms to the distribution of valid actions, assuming the model is well-trained. Compared to other generative models, such as Generative Adversarial Networks (GANs) and Variational Autoencoders (VAEs), the normalizing flow model is known to be more efficient at completing the task of data generation [9].
- Sampling uniformly from the feasible action space to train the flow model is challenging. For example, valid action space may be characterized using a set of constraints over binary variables. We therefore develop multiple methods, based on Hamiltonian Monte-Carlo [3] and probabilistic decision diagrams (PSDD) [16, 28], for valid action sampling for (non)convex constraints. These methods are significantly more sample efficient than the standard rejection sampling.
- We propose a simple and easy to implement method to integrate the normalizing flow model with deep RL algorithms such as DDPG [17]. Our approach is general, requires change only in the last layer of the policy network, and can be integrated with other deep RL methods. In our method, the modified policy network outputs an element of the latent distribution used in the normalizing flow model, which is then transformed into a valid action via the mapping provided by the normalizing flow. We also show gradients can be propagated through the normalizing flow to improve the policy. By integrating the normalizing flow with deep RL algorithms, we can avoid the zero gradient issue since there is no projection of invalid action onto the feasible action space. We show empirically that the flow model can be trained well on a variety of action constraints used in the literature, therefore, probability of constraint violation occurring remains low. Empirically, our approach results in significantly fewer constraint violations (upto 10x less), and is multiple times faster on a variety of continuous control tasks than the previous best method [18].
Action-constrained Markov Decision Process We consider a Markov decision process (MDP) model, which is defined by a tuple  $(S,A,p,r,\gamma ,b_{0})$ , where  $S$  is the set of possible states,  $A$  is the set of possible actions that an agent can take,  $p(s_{t + 1}|s_t,a_t)$  is the probability distribution over the next state  $s_{t + 1}$  given the current state  $s_t$  and action  $a_{t}$ ,  $r(s_{t},a_{t})$  is the reward function that determines the reward received by the agent after taking action  $a$  in state  $s$  at time step  $t$ ,  $\gamma \in [0,1)$  is the discount factor, and  $b_{0}$  is the initial state distribution. We assume the action space  $A$  is continuous, and focus on deterministic policies. Specifically, let  $\mu_{\theta}(\cdot)$  denote a deterministic policy parameterized by  $\theta$ ; action taken in state  $s$  is given as  $a = \mu_{\theta}(s)$ . Under policy  $\mu_{\theta}$ , the state value function  $(V)$  and the state-action value function  $(Q)$  are defined as follows respectively.
$$
V ^ {\mu_ {\theta}} (s) = \mathbb {E} [ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r \left(s _ {t}, a _ {t}\right) | s _ {0} = s; \mu_ {\theta} ], Q ^ {\mu_ {\theta}} (s, a) = \mathbb {E} [ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r \left(s _ {t}, a _ {t}\right) | s _ {0} = s, a _ {0} = a; \mu_ {\theta} ] \tag {1}
$$
An action-constrained MDP extends the standard MDP by incorporating explicit action constraints that determine a valid action set  $\mathcal{C}(s) \subseteq A$  for each state  $s \in S$ . In other words, the agent can only choose an action from the set of valid actions at each time step. The goal of an action-constrained MDP is to find a policy that maximizes the expected discounted total reward while ensuring that all chosen actions are valid with respect to the constraints. Formally, we have:
$$
\max  _ {\theta} J \left(\mu_ {\theta}\right) = \mathbb {E} _ {s \sim b _ {0}} \left[ V ^ {\mu_ {\theta}} (s) \right] \quad \text {s . t .} a _ {t} \in \mathcal {C} \left(s _ {t}\right) \quad \forall t \tag {2}
$$
We consider a RL setting where transition and reward functions are not known. The agent learns a policy by interacting with the environment and using collected experiences  $(s_t, a_t, r_t, s_{t+1})$ .
Deep Deterministic Policy Gradient Deep Deterministic Policy Gradient (DDPG) is a RL algorithm specifically designed to handle continuous action spaces [17]. The algorithm combines deterministic policy gradient [29] with deep  $Q$ -learning [25]. To solve an unconstrained RL problem, DDPG applies stochastic gradient ascent to update policy parameters  $\theta$  in the direction of the gradient  $\nabla_{\theta}J(\mu_{\theta})$ . The deterministic policy gradient theorem is used to compute the policy gradient as:
$$
\nabla_ {\theta} J (\mu_ {\theta}) = \mathbb {E} _ {s \sim \mathcal {B}} [ \nabla_ {a} Q ^ {\mu_ {\theta}} (s, a; \phi) \nabla_ {\theta} \mu_ {\theta} (s) | _ {a = \mu_ {\theta} (s)} ] \tag {3}
$$
The state-action value function  $Q(s, a; \phi)$  is approximated using a deep neural network, and the parameters  $\phi$  of this neural network are updated by minimizing the following loss function:
$$
\mathbb {E} _ {(s, a, s ^ {\prime}, r) \sim \mathcal {B}} \left[ \left(r + \gamma Q \left(s ^ {\prime}, \mu_ {\theta^ {\prime}} \left(s ^ {\prime}\right); \phi^ {\prime}\right) - Q (s, a; \phi)\right) ^ {2} \right] \tag {4}
$$
where  $(s, a, s', r)$  is a transition sample;  $\mathcal{B}$  is the replay buffer that stores transition samples;  $\theta'$  and  $\phi'$  are the parameters of the target policy and the target state-action value function respectively.
Normalizing Flows A normalizing flow model is a type of generative model that transforms a simple distribution, such as a Gaussian or uniform distribution, into a complex one by applying a sequence of invertible transformation functions [9, 27]. Let  $f = f_{K} \circ f_{K-1} \circ \ldots \circ f_{1} : \mathbb{R}^{D} \mapsto \mathbb{R}^{D}$  denote a normalizing flow model, where each  $f_{i} : \mathbb{R}^{D} \mapsto \mathbb{R}^{D}$ ,  $i = 1 : K$ , is an invertible transformation function. Starting from an initial sample from a base distribution (or prior distribution),  $z_{0} \sim p(z^{0})$ , the transformed sample from the model is  $x = f_{K} \circ f_{K-1} \circ \ldots \circ f_{1}(z^{0})$ . Each  $f_{i}$  takes input  $z^{i-1}$  and outputs  $z^{i}$ , and  $x = z^{K}$ . Given a training dataset  $\mathcal{D}$ , the mapping function  $f$  is learned by maximizing the log-likelihood of the data, which is defined as  $\log p(\mathcal{D}) = \sum_{x \in \mathcal{D}} \log p(x)$ . The log probability  $\log p(x)$  is computed by repeatedly applying the change of variables theorem, and is expressed as:
$$
\log p (x) = \log p \left(z ^ {0}\right) - \sum_ {i = 1} ^ {K} \log \left| \det  \frac {\partial f _ {i} \left(z ^ {i - 1}\right)}{\partial \left(z ^ {i - 1}\right) ^ {T}} \right| \tag {5}
$$
3 Normalizing Flow for ACRL
3.1 Learning Valid Action Mapping Using Normalizing Flows
We employ normalizing flows to establish an invertible mapping between the support of a simple distribution and the space of valid actions. While various transformation functions can be utilized
for implementing normalizing flows, we specifically focus on the conditional RealNVP model [34] due to its suitability for the general ACRL setting, where the set of valid actions is dependent on the state variable. The conditional RealNVP extends the original RealNVP [9] by incorporating the conditioning variable in both the prior distribution and the transformation functions. These transformation functions, which are implemented as affine coupling layers, possess the advantageous properties of enabling efficient forward propagation during model learning and efficient backward propagation for sample generation [9].
Given a dataset  $\mathcal{D}$  consisting of valid actions (details in Section 3.2), we train the invertible mapping  $f_{\psi}$  parameterized by  $\psi$  to capture the relationship between the support of a uniform distribution  $p$  and the elements in  $\mathcal{D}$ . Learning this mapping provides an easy way to generate new valid actions. The policy network outputs an element from the domain of the uniform distribution, which is then mapped to a valid action using the learned flow. This process of generating actions is much simpler than using a math program for action projection, which can result in the zero gradient issue.
The flow learning process involves maximizing the log-likelihood of the data, following the methodology presented in [34]. In contrast to using a Gaussian distribution, we opt for a uniform distribution as it demonstrated better empirical performance when combined with the DDPG algorithm. Once the normalizing flow model is learned, during backward propagation with conditioning variables, each bijective function  $f_{\psi_i}$ , where  $i = 1,\dots ,K$ , takes inputs of  $z^{i - 1}\in \mathbb{R}^D$  and a conditioning variable  $y\in \mathbb{R}^{D_y}$  (representing state features) and produces an output of  $x^{i}\in \mathbb{R}^{D}$  as:
$$
x _ {1: d} ^ {i} = z _ {1: d} ^ {i - 1} \tag {6}
$$
$$
x _ {d + 1: D} ^ {i} = \left(z _ {d + 1: D} ^ {i - 1} - t \left(z _ {1: d} ^ {i - 1}, y\right)\right) \odot \exp \left(- k \left(z _ {1: d} ^ {i - 1}, y\right)\right) \tag {7}
$$
where  $d < D$  is the index that splits dimensions of input  $x^i$  into two parts;  $k$  and  $t$  are scale and translation functions that map from  $\mathbb{R}^{d + D_y}\mapsto \mathbb{R}^{D - d}$ . These two functions are modeled using neural networks.  $\odot$  denotes the element-wise product (Hadamard product). Therefore, given a random member of the uniform distribution  $z^0\sim p(z^0)$  and a conditioning variable  $y$ , the mapped element is:
$$
x = f _ {\psi_ {K}} \circ f _ {\psi_ {K - 1}} \circ \dots \circ f _ {\psi_ {1}} (z ^ {0}, y) \tag {8}
$$
Note that the transformed output  $x$  obtained through the normalizing flow model may not precisely match any specific element from the dataset  $\mathcal{D}$ . However, when the mapping  $f_{\psi}$  is effectively trained, the resulting  $x$  should exhibit similar characteristics to the elements in  $\mathcal{D}$ , thereby satisfying the constraints with high probability, as we also observed empirically. In essence, the combination of the uniform distribution and the learned mapping serves as an efficient representation of all valid actions, providing a means to generate actions that adhere to the constraints.
Unlike other generative models, normalizing flows provide the ability to measure the recall rate in learning the valid action mapping due to their invertible bijective transformations. The recall rate (we call it "coverage") indicates the fraction of valid actions that can be generated from the latent space. Given a conditioning variable  $y$ , let  $\mathcal{C}(y)$  denote the set of valid actions that are uniformly distributed in the feasible region. The recall can be computed as follows.
$$
\operatorname {r e c a l l} (y) = \frac {\sum_ {x \in \mathcal {C} (y)} \mathbb {I} _ {\operatorname {d o m} _ {f _ {\psi}}} f _ {\psi} ^ {- 1} (x , y)}{| \mathcal {C} (y) |} \tag {9}
$$
where  $f_{\psi}^{-1}$  is the inverse transformation function of the normalizing flows model  $f_{\psi}$ . Recall is useful to characterize the learned generative model. If the recall is low, it would limit the feasible constrained space over which ACRL algorithm optimizes the policy. This can result in lower solution quality.
The Mollified Uniform Prior Most recent studies of normalizing flows employ a standard Gaussian distribution as the latent space distribution [9, 15, 27], which spans the entire real number space  $\mathbb{R}^D$ . However, through empirical observation, we noticed that with Gaussian latent distribution, the feasible action region tends to be mapped closer to zero in the latent space, where the probability of the standard Gaussian is higher. On the other hand, points located far away from the center of the Gaussian (several standard deviations) often get mapped into the infeasible action space. Consequently, when coupling the DDPG policy network with the flow, it resulted in high constraint violations as the policy network output can be any number from the support of the Gaussian, and not just limited to the high probability region.
(c)
Figure 1: (a) An SDD representing the PB constraint  $A \cdot 2^{1} + B \cdot 2^{0} + C \cdot 2^{1} + D \cdot 2^{0} \leq 2$ ; (b) A PSDD; (c) Samples generated with HMC
To address this issue, we decided to use a uniform distribution from  $[-1, 1]^D$  as the latent distribution instead of the standard Gaussian. In the uniform distribution, there is no high or low probability region unlike the standard Gaussian. However, the uniform distribution is not differentiable, therefore we employ a modified version called the Mollified uniform distribution [11]. In order to train the normalizing flow, we require the probability density of the prior distribution as in Eq. 5. Computing the probability density of a multi-dimensional mollified distribution can be challenging since it requires integrating over all dimensions [11]. However, in the case of a single dimension, a sample drawn from a mollified uniform distribution can be seen as a sample from a uniform distribution with added Gaussian noise. In other words, a sample  $x$  is obtained as  $x = x_0 + \delta$ , where  $x_0$  follows a uniform distribution in the range  $[-1, 1]$ , and  $\delta$  follows a Gaussian distribution with mean 0 and standard deviation  $\sigma$ .
To calculate the overall probability density, we treat each dimension individually, compute their probability densities, and take their product as the overall probability. The probability density in a single dimension can be easily calculated using the following formula:  $p(x) = \int_{-1}^{1}\frac{1}{\sqrt{2\pi\sigma^2}} e^{-\frac{1}{2\sigma^2} (x - x_0)^2}dx_0 = \Phi \left(\frac{1 - x}{\sigma}\right) - \Phi \left(\frac{-1 - x}{\sigma}\right)$ . In this formula,  $\Phi$  represents the cumulative density function of the standard Gaussian distribution.
3.2 Generating Random Samples from the Valid Action Space
To effectively train the normalizing flow model, we need to generate samples from the set of valid actions. Even with the known analytical form of action constraints this is a challenging problem. Standard rejection sampling is highly inefficient to sample from the complex constraints that arise in ACRL benchmarks. Therefore, we use two methods described next.
Hamiltonian Monte-Carlo (HMC) HMC is a Markov chain Monte Carlo method, which utilizes energy conservation to effectively sample a target distribution [3]. It utilizes Hamiltonian-dynamics to simulate the movement of a particle in a continuous space based on a potential energy function. It uses leap-frog algorithm with the first order integrator, to estimate the next position of the particle. In our work, we need to generate samples in the feasible region. Therefore we define energy function to be a fixed value 0 in the feasible region. Then we bound the feasible region with a hard wall of infinity potential by setting the potential to  $\infty$  in the infeasible region. Figure 1(c) depicts the distribution of feasible action samples generated by HMC in the action constraint of the Reacher Environment ( $a_1^2 + a_2^2 \leq 0.05$ ). It presents a contour map of the sample density in Figure 1(c)-(ii), as well as their projection onto each dimension in Figure 1(c)-(i) and Figure 1(c)-(iii) respectively.
Probabilistic Sentential Decision Diagrams In many real-world applications such as resource allocation, the constraints are defined using a set of linear equalities/inequalities over integer variables. However, sampling valid actions using HMC that satisfy such constraints is very challenging due to the presence of equality constraints. To address this challenge, we propose to use a Probabilistic Sentential Decision Diagram (PSDD) [16] to encode a uniform distribution over all the valid actions and then sample actions from the PSDD. For this setting, we assume that all the constraints are linear (in)equalities. We describe the process of constructing a PSDD next. Consider an inequality constraint  $x + y \leq 2$ . We follow the following steps: