SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \operatorname * {P r} (\hat {Z} _ {t} ^ {\mathrm {t o t}} = z) := \frac {1}{\mathcal {N} _ {t}} (1 - \lambda) \sum_ {i = 0} ^ {\infty} \lambda^ {i} \prod_ {j = t + 1} ^ {t + i} \frac {\pi (a _ {j} | s _ {j})}{\mu (a _ {j} | s _ {j})} \sum_ {m = 1} ^ {M} \frac {1}{M} \delta_ {\hat {Z} _ {R, m} ^ {(i + 1)} (s _ {t}, a _ {t})} (z) \\ = \frac {1 - \lambda}{\mathcal {N} _ {t}} \left(\sum_ {m = 1} ^ {M} \frac {1}{M} \delta_ {\hat {Z} _ {t, m} ^ {(1)}} (z) + \lambda \frac {\pi \left(a _ {t + 1} \mid s _ {t + 1}\right)}{\mu \left(a _ {t + 1} \mid s _ {t + 1}\right)} \sum_ {i = 0} ^ {\infty} \lambda^ {i} \prod_ {j = t + 2} ^ {t + 1 + i} \frac {\pi \left(a _ {j} \mid s _ {j}\right)}{\mu \left(a _ {j} \mid s _ {j}\right)} \sum_ {m = 1} ^ {M} \frac {1}{M} \delta_ {\hat {Z} _ {R, m} ^ {(i + 2)} \left(s _ {t}, a _ {t}\right)} (z)\right) \tag {12} \\ = \frac {1 - \lambda}{\mathcal {N} _ {t}} \underbrace {\sum_ {m = 1} ^ {M} \frac {1}{M} \delta_ {\hat {Z} _ {t , m} ^ {(1)}} (z)} _ {\text {(a)}} + \lambda \frac {\mathcal {N} _ {t + 1}}{\mathcal {N} _ {t}} \frac {\pi (a _ {t + 1} | s _ {t + 1})}{\mu (a _ {t + 1} | s _ {t + 1})} \underbrace {\operatorname * {P r} (R _ {t} + \gamma \hat {Z} _ {t + 1} ^ {\mathrm {t o t}} = z)} _ {\text {(b)}}, \\ \end{array}
$$
where  $\mathcal{N}_t$  is a normalization factor. If the target for time step  $t + 1$  is obtained, the target distribution for time step  $t$  becomes the weighted sum of (a) the current one-step TD target and (b) the shifted
Figure 2: Procedure for constructing the target distribution. First, multiply the target at  $t + 1$  step by  $\gamma$  and add  $R_{t}$ . Next, weight-combine the shifted previous target and one-step target at  $t$  step and restore the CDF of the combined target. The CDF can be restored by sorting the positions of the atoms and then accumulating the weights at each atom position. Finally, the projected target can be obtained by finding the positions of the atoms corresponding to  $M'$  quantiles in the CDF. Using the projected target, the target at  $t - 1$  step can be found recursively.
previous target distribution, so it can be obtained recursively, as shown in (12). Since the definition requires infinite sums, the recursive way is more practical for computing the target. Nevertheless, to obtain the target in that recursive way, we need to store all quantile positions and weights for all time steps, which is not memory-efficient. Therefore, we propose to project the target distribution into a quantile distribution with a specific number of atoms,  $M'$  (we set  $M' > M$  to reduce information loss). The overall process to get the  $\mathrm{TD}(\lambda)$  target distribution is illustrated in Figure 2, and the pseudocode is given in Appendix A.5. Furthermore, we can show that a distribution trained with the proposed target converges to the distribution of  $Z_R^\pi$ .
Theorem 3.2. Let define a distributional operator  $\mathcal{T}_{\lambda}^{\mu, \pi}$ , whose probability density function is:
$$
\Pr \left(\mathcal {T} _ {\lambda} ^ {\mu , \pi} Z (s, a) = z\right) \propto
$$
$$
\sum_ {i = 0} ^ {\infty} \mathbb {E} _ {\mu} \left[ \lambda^ {i} \prod_ {j = 1} ^ {i} \frac {\pi \left(a _ {j} \mid s _ {j}\right)}{\mu \left(a _ {j} \mid s _ {j}\right) {} _ {a ^ {\prime}} \sim \pi (\cdot | s _ {i + 1})} \mathbb {E} _ {a ^ {\prime} \sim \pi (\cdot | s _ {i + 1})} \left[ \Pr \left(\sum_ {t = 0} ^ {i} \gamma^ {t} R _ {t} + \gamma^ {i + 1} Z \left(s _ {i + 1}, a ^ {\prime}\right) = z\right) \right] \mid s _ {0} = s, a _ {0} = a \right]. \tag {13}
$$
Then, a sequence,  $Z_{k + 1}(s,a) = \mathcal{T}_{\lambda}^{\mu ,\pi}Z_{k}(s,a)\forall (s,a)$ , converges to  $Z_R^\pi$
The  $\mathrm{TD}(\lambda)$  target is a quantile distribution version of the distributional operator  $\mathcal{T}_{\lambda}^{\mu,\pi}$  in Theorem 3.2. Consequently, a distribution updated by minimizing the quantile regression loss with the  $\mathrm{TD}(\lambda)$  target converges to the distribution of  $Z_R^\pi$  if the number of atoms is infinite, according to Theorem 3.2. The proof of Theorem 3.2 is provided in Appendix A.4. After calculating the target distribution for all time steps, the critic can be trained to reduce the quantile regression loss with the target distribution. To provide more insight, we experiment with a toy example in Appendix A.6, and the results show that the proposed target distribution can trade off bias and variance through the trace-decay  $\lambda$ .
3.3 Safe Distributional Actor-Critic
Finally, we describe the proposed method, safe distributional actor-critic (SDAC). After collecting trajectories, the policy is updated by solving (5), which can be solved through a line search (for more detail, see Appendix A.8). The cost value and the cost square value functions in (4) can be obtained using the distributional critics as follows:
$$
Q _ {C} ^ {\pi} (s, a) = \int_ {- \infty} ^ {\infty} z \Pr \left(Z _ {C} ^ {\pi} (s, a) = z\right) d z \approx \frac {1}{M} \sum_ {m = 1} ^ {M} \theta_ {m} (s, a), \tag {14}
$$
$$
S _ {C} ^ {\pi} (s, a) = \int_ {- \infty} ^ {\infty} z ^ {2} \Pr (Z _ {C} ^ {\pi} (s, a) = z) d z \approx \frac {1}{M} \sum_ {m = 1} ^ {M} \theta_ {m} (s, a) ^ {2}.
$$
If a solution of (5) does not exist, the policy is projected into a feasible region through the proposed gradient integration method. The critics can also be updated by the regression loss (7) between the target distribution obtained from (12). The proposed method is summarized in Algorithm 1.
Safe Reinforcement Learning. Garcia and Fernandez [2015] and Gu et al. [2022] have researched and categorized safe RL methodologies from various perspectives. In this paper, we introduce safe RL
Algorithm 1 Safe Distributional Actor-Critic
Input: Policy network  $\pi_{\psi}$ , reward and cost critical networks  $Z_{R,\theta}^{\pi}$ ,  $Z_{C_k,\theta}^{\pi}$ , and replay buffer  $\mathcal{D}$ .
Initialize network parameters  $\psi, \theta$ , and replay buffer  $\mathcal{D}$ .
for epochs  $= 1$  to  $E$  do
Sample  $a_{t}\sim \pi_{\psi}(\cdot |s_{t})$  and get  $s_{t + 1},r_t = R(s_t,a_t,s_{t + 1})$  , and  $c_{k,t} = C_k(s_t,a_t,s_{t + 1})\forall k.$
Store  $(s_t, a_t, \pi_{\psi}(a_t | s_t), r_t, c_{\{1, \dots, K\}, t}, s_{t+1})$  in  $\mathcal{D}$ .
Calculate the TD(λ) target distribution (Section 3.2) using  $\mathcal{D}$  and update the critics to minimize the quantile loss defined in (7).
Calculate the surrogates for the objective and constraints defined in (37) using  $\mathcal{D}$ .
Update the policy by solving (5), but if (5) has no solution, take a recovery step (Section 3.1).
methods depending on how to update policies to reflect safety constraints. First, trust region-based safe RL methods [Achiam et al., 2017, Yang et al., 2020, Kim and Oh, 2022a] find policy update directions by approximating the safe RL problem as a linear-quadratic constrained linear program and update policies through a line search. Yang et al. [2020] also employ projection to meet a constraint; however, their method is limited to a single constraint and does not show to satisfy the constraint for the infeasible starting case. Second, Lagrangian-based methods [Stooke et al., 2020, Yang et al., 2021, Liu et al., 2020] convert the safe RL problem to a dual problem and update the policy and dual variables simultaneously. Last, expectation-maximization (EM) based methods [Liu et al., 2022, Zhang et al., 2022] find non-parametric policy distributions by solving the safe RL problem in E-steps and fit parametric policies to the found non-parametric distributions in M-steps. Also, there are other ways to reflect safety other than policy updates. Qin et al. [2021], Lee et al. [2022] find optimal state or state-action distributions that satisfy constraints, and Bharadhwaj et al. [2021], Thananjeyan et al. [2021] reflect safety during exploration by executing only safe action candidates. In the experiments, only the safe RL methods of the policy update approach are compared with the proposed method.
Distributional TD( $\lambda$ ). TD( $\lambda$ ) [Precup et al., 2000] can be extended to the distributional critic to trade off bias-variance. Gruslys et al. [2018] have proposed a method to obtain target distributions by mixing  $n$ -step distributions, but the method is applicable only in discrete action spaces. Nam et al. [2021] have proposed a method to obtain target distributions using sampling to apply to continuous action spaces, but this is only for on-policy settings. A method proposed by Tang et al. [2022] updates the critics using newly defined distributional TD errors rather than target distributions. This method is applicable for off-policy settings but has the disadvantage that memory usage increases linearly with the number of TD error steps. In contrast to these methods, the proposed method is memory-efficient and applicable for continuous action spaces under off-policy settings.
Gradient Integration. The proposed feasibility handling method utilizes a gradient integration method, which is widely used in multi-task learning (MTL). The gradient integration method finds a single gradient to improve all tasks by using gradients of all tasks. Yu et al. [2020] have proposed a projection-based gradient integration method, which is guaranteed to converge Pareto-stationary sets. A method proposed by Liu et al. [2021] can reflect user preference, and Navon et al. [2022] proposed a gradient-scale invariant method to prevent the training process from being biased by a few tasks. The proposed method can be viewed as a mixture of projection and scale-invariant methods as gradients are clipped and projected onto a trust region.
We evaluate the safety performance of the proposed method in single- and multi-constrained robotic tasks. For single constraints, the agent performs four tasks provided by Safety Gym [Ray et al., 2019], and for multi-constraints, it performs bipedal and quadrupedal locomotion tasks.
5.1 Safety Gym
Tasks. We employ two robots, point and car, to perform goal and button tasks in the Safety Gym. The goal task is to control a robot toward a randomly spawned goal without passing through hazard
Figure 3: Safety Gym task results. The cost rates show the cost sums divided by the episode length, and the dashed black lines indicate the threshold of the constraint. All methods are trained with five random seeds.
regions. The button task is to click a randomly designated button using a robot, where not only hazard regions but also dynamic obstacles exist. Agents get a cost when touching undesignated buttons and obstacles or entering hazard regions. There is only one constraint for the Safety Gym tasks, and it is defined using (4) with the sum of costs. Constraint violations (CVs) are counted when the cost sum exceeds the threshold. For more details, see Appendix B.
Baselines. Safe RL methods based on various types of policy updates are used as baselines. For the trust region-based method, we use constrained policy optimization (CPO) [Achiam et al., 2017] and off-policy trust-region CVaR (OffTRC) [Kim and Oh, 2022a], which extend the CPO to an off-policy and mean-standard constrained version. For the Lagrangian-based method, distributional worst-case soft actor-critic (WCSAC) [Yang et al., 2022] is used, and constrained variational policy optimization (CVPO) [Liu et al., 2022] based on the EM method is used. Specifically, WCSAC, OffTRC, and the proposed method, SDAC, use the risk-averse constraints, so we experiment with those for  $\alpha = 0.25$  and 1.0 (when  $\alpha = 1.0$ , the constraint is identical to the risk-neutral constraint).
Results. The graph of the final reward sum, cost rate, and the total number of CVs are shown in Figure 3a, and the training curves are provided in Appendix C.1. We can interpret the results as good if the reward sum is high and the cost rate and total CVs are low. SDAC with  $\alpha = 0.25$ , risk-averse constraint situations, satisfies the constraints in all tasks and shows an average of 1.78 times fewer total CVs than the second-best algorithm. Nevertheless, since the reward sums are also in the middle or upper ranks, its safety performance is of high quality. SDAC with  $\alpha = 1.0$ , risk-neutral constraint situations, shows that the cost rates are almost the same as the thresholds except for the car button. In the case of the car button, the constraint is not satisfied, but by setting  $\alpha = 0.25$ , SDAC can achieve the lowest total CVs and the highest reward sum compared to the other methods. As for the reward sum, SDAC is the highest in the point goal and car button, and WCSAC is the highest in the rest. However, WCSAC seems to lose the risk-averse properties seeing that the cost rates do not change significantly according to  $\alpha$ . This is because WCSAC does not define constraints as risk measures of cost returns but as expectations of risk measures [Yang et al., 2022]. OffTRC has lower safety performance than SDAC in most cases because, unlike SDAC, it does not use distributional critics. Finally, CVPO and CPO are on-policy methods, so they are less efficient than the other methods.
5.2 Locomotion Tasks
Tasks. The locomotion tasks are to train robots to follow  $xy$ -directional linear and  $z$ -directional angular velocity commands. Mini-Cheetah from MIT [Katz et al., 2019] and Laikago from Unitree [Wang, 2018] are used for quadrupedal robots, and Cassie from Agility Robotics [Xie et al., 2018] is used for a bipedal robot. In order to perform the locomotion tasks, robots should keep balancing, standing, and stamping their feet so that they can move in any direction. Therefore, we define three constraints. The first is to keep the balance so that the body angle does not deviate from zero, and the second is to keep the height of CoM above a threshold. The third is to match the current foot
(a) Training curves of the Cassie task.
(b) Training curves of the Laikago task.
(c) Training curves of the Mini-Cheetah task.
(d) Training curves of the naive and proposed methods for the Cassie task.
Figure 4: Locomotion task results. The black dashed lines indicate the thresholds, and the dotted lines in (d) represent the thresholds  $+0.025$ . The shaded area represents half of the standard deviation, and all methods are trained with five random seeds.
contact state with a predefined foot contact timing. The reward is defined as the negative  $l^2$ -norm of the difference between the command and the current velocity. CVs are counted when the sum of at least one cost rate exceeds the threshold. For more details, see Appendix B.
Baselines. The baseline methods are identical to the Safety Gym tasks, and CVPO is excluded because it is technically challenging to scale to multiple constraint settings. We set  $\alpha$  to 1 for the risk-averse constrained methods (OffTRC, WCSAC, and SDAC) to focus on measuring multi-constraint handling performance.
Results. Figure 4.(a-c) presents the training curves. SDAC shows the highest reward sums and the lowest total CVs in all tasks. In particular, the number of steps required to satisfy all constraints is 1.93 times fewer than the second-best algorithm on average. Trust region methods (OffTRC, CPO) stably satisfy constraints, but they are not efficient since they handle constraints by the naive approach. WCSAC, a Lagrangian method, fails to keep the constraints and shows the lowest reward sums. This is because the Lagrange multipliers can hinder the training stability due to the concurrent update with policy [Stooke et al., 2020].
We conduct ablation studies to show whether the proposed target distribution lowers the estimation bias and whether the proposed gradient integration quickly converges to the feasibility condition. In Figure 3b, the number of CVs is reduced as  $\lambda$  increases, which means that the bias of constraint estimation decreases. However, the score also decreases due to large variance, showing that  $\lambda$  can adjust the bias-variance tradeoff. In Figure 4d, the proposed gradient integration method is compared with the naive approach, which minimizes the constraints in order from the first to the third constraint, as described in Section 3.1. The proposed method reaches the feasibility condition faster than the naive approach and shows stable training curves because it reflects all constraints concurrently. Additionally, we analyze the distributional critics in Appendix C.2 and the hyperparameters, such as the trust region size, in Appendix C.3. Furthermore, we analyze the sensitivity of traditional RL algorithms to reward configuration in Appendix D, emphasizing the advantage of safe RL that does not require reward tuning.
A limitation of the proposed method is that the computational complexity of the gradient integration is proportional to the square of the number of constraints, whose qualitative analysis is presented in Appendix E.1. Also, we conducted quantitative analyses in Appendix E.2 by measuring wall clock training time. In the mini-cheetah task which has three constraints, the training time of SDAC is the third fastest among the four safe RL algorithms. Gradient integration is not applied when the policy satisfies constraints, so it may not constitute a significant proportion of training time. However, its influence can be dominant as the number of constraints increases. In order to resolve this limitation, the calculation can be speeded up by stochastically selecting a subset of constraints [Liu et al., 2021], or by reducing the frequency of policy updates [Navon et al., 2022]. The other limitation is that the mean-standard defined in (2) is not a coherent risk measure. As a result, mean-standard constraints can be served as reducing uncertainty rather than risk, although we experimentally showed that constraint violations are efficiently reduced. To resolve this, we can use the CVaR constraint, which can be estimated using an auxiliary variable, as done by Chow et al. [2017]. However, this solution can destabilize the training process due to the auxiliary variable, as observed in experiments of Kim and Oh [2022b]. Hence, a stabilization technique should be developed to employ the CVaR constraint.
We have presented the trust region-based safe distributional RL method, called SDAC. Through the locomotion tasks, it is verified that the proposed method efficiently satisfies multiple constraints using the gradient integration. Moreover, constraints can be stably satisfied in various tasks due to the low-biased distributional critics trained using the proposed target distributions. In addition, the proposed method is analyzed from multiple perspectives through various ablation studies. However, to compensate for the computational complexity, future work plans to devise efficient methods when dealing with large numbers of constraints.
This work was partly supported by Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. 2019-0-01190, [SW Star Lab] Robot Learning: Efficient, Safe, and Socially-Acceptable Machine Learning,  $34\%$ ), Basic Science Research Program through the National Research Foundation of Korea (NRF) funded by the Ministry of Science and ICT (NRF-2022R1A2C2008239, General-Purpose Deep Reinforcement Learning Using Metaverse for Real World Applications,  $33\%$ ), and Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. 2021-0-01341, AI Graduate School Program, CAU,  $33\%$ ).
J. Achiam, D. Held, A. Tamar, and P. Abbeel. Constrained policy optimization. In Proceedings of International Conference on Machine Learning, pages 22-31, 2017.
A. Agarwal, S. M. Kakade, J. D. Lee, and G. Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. The Journal of Machine Learning Research, 22 (1):4431-4506, 2021.
E. Altman. Constrained Markov decision processes, volume 7. CRC Press, 1999.
Q. Bai, A. S. Bedi, M. Agarwal, A. Koppel, and V. Aggarwal. Achieving zero constraint violation for constrained reinforcement learning via primal-dual approach. Proceedings of the AAAI Conference on Artificial Intelligence, 36(4), 2022.
M. G. Bellemare, W. Dabney, and R. Munos. A distributional perspective on reinforcement learning. In Proceedings of International Conference on Machine Learning, pages 449-458, 2017.
M. G. Bellemare, W. Dabney, and M. Rowland. Distributional Reinforcement Learning. MIT Press, 2023. http://www.distributional-rl.org.
T. Xu, Y. Liang, and G. Lan. CRPO: A new approach for safe reinforcement learning with convergence guarantee. In Proceedings of International Conference on Machine Learning, pages 11480-11491, 2021.
Q. Yang, T. D. Simão, S. H. Tindemans, and M. T. J. Span. WCSAC: Worst-case soft actor critic for safety-constrained reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence, 35(12):10639-10646, 2021.
Q. Yang, T. D. Simão, S. H. Tindemans, and M. T. J. Spanan. Safety-constrained reinforcement learning with a distributional safety critic. Machine Learning, 112:859-887, 2022.
T.-Y. Yang, J. Rosca, K. Narasimhan, and P. J. Ramadge. Projection-based constrained policy optimization. In Proceedings of International Conference on Learning Representations, 2020.
T. Yu, S. Kumar, A. Gupta, S. Levine, K. Hausman, and C. Finn. Gradient surgery for multi-task learning. In Advances in Neural Information Processing Systems, pages 5824-5836, 2020.
H. Zhang, Y. Lin, S. Han, S. Wang, and K. Lv. Conservative distributional reinforcement learning with safety constraints. arXiv preprint arXiv:2201.07286, 2022.
J. Zhang and P. Weng. Safe distributional reinforcement learning. In J. Chen, J. Lang, C. Amato, and D. Zhao, editors, Proceedings of International Conference on Distributed Artificial Intelligence, pages 107-128, 2022.
We denote the policy parameter space as  $\Psi \subseteq \mathbb{R}^d$ , the parameter at the  $t$ th iteration as  $\psi_t \in \Psi$ , the Hessian matrix as  $H(\psi_t) = \nabla_{\psi}^2 D_{\mathrm{KL}}(\pi_{\psi_t}||\pi_{\psi})|_{\psi = \psi_t}$ , and the  $k$ th constraint as  $F_k(\psi_t) = F_k(\pi_{\psi_t};\alpha)$ . As we focus on the  $t$ th iteration, the following notations are used for brevity:  $H = H(\psi_t)$  and  $g_k = \nabla F_k(\psi_t)$ . The proposed gradient integration at  $t$ th iteration is defined as the following quadratic program (QP):
$$
g _ {t} = \underset {g} {\operatorname {a r g m i n}} \frac {1}{2} g ^ {T} H g \quad \text {s . t .} g _ {k} ^ {T} g + c _ {k} \leq 0 \text {f o r} \forall k, \tag {15}
$$
where  $c_{k} = \min (\sqrt{2\epsilon g_{k}^{T}H^{-1}g_{k}},F_{k}(\pi_{\psi}) - d_{k} + \zeta)$ . In the remainder of this section, we introduce the assumptions and new definitions, discuss the existence of a solution (15), show the convergence to the feasibility condition for varying step size cases, and provide the proof of Theorem 3.1.
Assumption. 1) Each  $F_{k}$  is differentiable and convex, 2)  $\nabla F_{k}$  is  $L$ -Lipschitz continuous, 3) all eigenvalues of the Hessian matrix  $H(\psi)$  are equal or greater than  $R \in \mathbb{R}_{>0}$  for  $\forall \psi \in \Psi$ , and 4)  $\{\psi | F_{k}(\psi) + \zeta < d_{k} \text{ for } \forall k\} \neq \emptyset$ .
Definition. Using the Cholesky decomposition, the Hessian matrix can be expressed as  $H = B \cdot B^T$  where  $B$  is a lower triangular matrix. By introducing new terms,  $\bar{g}_k \coloneqq B^{-1} g_k$  and  $b_t \coloneqq B^T g_t$ , the following is satisfied:  $g_k^T H^{-1} g_k = ||\bar{g}_k||_2^2$ . Additionally, we define the in-boundary and out-boundary sets as:
$$
\mathrm {I B} _ {k} := \left\{\psi | F _ {k} (\psi) - d _ {k} + \zeta \leq \sqrt {2 \epsilon \nabla F _ {k} (\psi) ^ {T} H ^ {- 1} (\psi) \nabla F _ {k} (\psi)} \right\},
$$
$$
\mathrm {O B} _ {k} := \left\{\psi | F _ {k} (\psi) - d _ {k} + \zeta \geq \sqrt {2 \epsilon \nabla F _ {k} (\psi) ^ {T} H ^ {- 1} (\psi) \nabla F _ {k} (\psi)} \right\}.
$$
The minimum of  $||\bar{g}_k||$  in  $\mathrm{OB}_k$  is denoted as  $m_k$ , and the maximum of  $||\bar{g}_k||$  in  $\mathrm{IB}_k$  is denoted as  $M_k$ . Also,  $\min_k m_k$  and  $\max_k M_k$  are denoted as  $m$  and  $M$ , respectively, and we can say that  $m$  is positive.