Risk-Averse Model Uncertainty for
Distributionally Robust Safe Reinforcement Learning

In this work, we consider RL in the presence of safety constraints. We model this sequential decision making problem as an infinite-horizon, discounted Constrained Markov Decision Process (CMDP) [4] defined by the tuple $(\mathcal{S},\mathcal{A},p,r,c,d_{0},\gamma)$, where $\mathcal{S}$ is the set of states, $\mathcal{A}$ is the set of actions, $p:\mathcal{S}\times\mathcal{A}\rightarrow P(\mathcal{S})$ is the transition model where $P(\mathcal{S})$ represents the space of probability measures over $\mathcal{S}$, $r,c:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ are the reward function and cost function used to define the objective and constraint, respectively, $d_{0}\in P(\mathcal{S})$ is the initial state distribution, and $\gamma$ is the discount rate. We focus on the setting with a single constraint, but all results can be extended to the case of multiple constraints.

We model the agent’s decisions as a stationary policy $\pi:\mathcal{S}\rightarrow P(\mathcal{A})$. For a given CMDP and policy $\pi$, we write the expected total discounted rewards and costs as $J_{p,r}(\pi)=\mathbb{E}_{\tau\sim(\pi,p)}\left[\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})\right]$ and $J_{p,c}(\pi)=\mathbb{E}_{\tau\sim(\pi,p)}\left[\sum_{t=0}^{\infty}\gamma^{t}c(s_{t},a_{t})\right]$, respectively, where $\tau\sim(\pi,p)$ represents a trajectory sampled according to $s_{0}\sim d_{0}$, $a_{t}\sim\pi(\,\cdot\mid s_{t})$, and $s_{t+1}\sim p(\,\cdot\mid s_{t},a_{t})$. The goal of safe RL is to find a policy $\pi$ that maximizes the constrained optimization problem

where $B$ is a safety budget on expected total discounted costs.

We write the corresponding state-action value functions (i.e., Q functions) for a given transition model $p$ and policy $\pi$ as $Q^{\pi}_{p,r}(s,a)$ and $Q^{\pi}_{p,c}(s,a)$, respectively. Off-policy optimization techniques [55, 28] find a policy that maximizes (1) by solving at each iteration the related optimization problem

where $\pi_{k}$ is the current policy and $\mathcal{D}$ is a replay buffer containing data collected in the training environment. Note that $Q^{\pi}_{p,r}(s,a)$ and $Q^{\pi}_{p,c}(s,a)$ are the respective fixed points of the Bellman operators

Rather than focusing on a single CMDP with transition model $p$, we incorporate uncertainty about the transition model by considering a distribution $\mu$ over models. We focus on distributions of the form $\mu=\prod_{(s,a)\in\mathcal{S}\times\mathcal{A}}\mu_{s,a}$, where $\mu_{s,a}$ represents a distribution over transition models $p_{s,a}=p(\,\cdot\mid s,a)\in P(\mathcal{S})$ at a given state-action pair and $\mu$ is the product over all $\mu_{s,a}$. This is known as rectangularity, and is a common assumption in the literature [53, 56, 16, 15, 11]. Note that $\mu_{s,a}\in P(\mathcal{M})$, where we write $\mathcal{M}=P(\mathcal{S})$ to denote model space. Compared to robust RL methods that apply uncertainty sets over transition models, the use of a distribution $\mu$ over transition models is a more informative way to represent model uncertainty that does not require solving for worst-case environments (i.e., does not introduce a minimax formulation).

In order to incorporate robustness to the choice of $\mu$, distributionally robust MDPs [53, 56] consider an ambiguity set $\mathcal{U}=\bigotimes_{(s,a)\in\mathcal{S}\times\mathcal{A}}\mathcal{U}_{s,a}$ of distributions over transition models, where $\mu_{s,a}\in\mathcal{U}_{s,a}\subseteq P(\mathcal{M})$. The goal of distributionally robust RL is to optimize the worst-case average performance across all distributions contained in $\mathcal{U}$. In this work, we will show that a risk-averse perspective towards model uncertainty defined by $\mu$ is equivalent to distributionally robust RL for appropriate choices of ambiguity sets in the objective and constraint of a CMDP. However, our use of risk measures avoids the need to solve for worst-case distributions in $\mathcal{U}$ throughout training.

Consider the probability space $(\mathcal{M},\mathcal{F},\mu_{s,a})$, where $\mathcal{F}$ is a $\sigma$-algebra on $\mathcal{M}$ and $\mu_{s,a}\in P(\mathcal{M})$ defines a probability measure over $\mathcal{M}$. Let $\mathcal{Z}$ be a space of random variables defined on this probability space, and let $\mathcal{Z}^{*}$ be its corresponding dual space. A real-valued risk measure $\rho:\mathcal{Z}\rightarrow\mathbb{R}$ summarizes a random variable as a value on the real line. In this section, we consider cost random variables $Z\in\mathcal{Z}$ where a lower value of $\rho(Z)$ is better. We can define a corresponding risk measure $\rho^{+}$ for reward random variables through an appropriate change in sign, where $\rho^{+}(Z)=-\rho(-Z)$. Risk-sensitive methods typically focus on classes of risk measures with desirable properties [30], such as coherent risk measures [6] and distortion risk measures [50, 17].

A distortion risk measure is coherent if and only if $g$ is concave [52]. In this work, we focus on the class of coherent distortion risk measures. We will leverage properties of coherent risk measures to provide robustness guarantees for our framework, and we will leverage properties of distortion risk measures to propose an efficient, model-free implementation that does not involve minimax optimization. See the Appendix for additional details on the properties of coherent distortion risk measures. Many commonly used risk measures belong to this class, including expectation, conditional value-at-risk (CVaR), and the Wang transform [51] for $\eta\geq 0$ which is defined by the distortion function $g_{\eta}(u)=\Phi(\Phi^{-1}(u)+\eta)$, where $\Phi$ is the standard Normal cumulative distribution function.

The CMDP framework is the most popular approach to safety in RL, and several deep RL algorithms have been developed to solve the constrained optimization problem in (1). These include primal-dual methods that consider the Lagrangian relaxation of (1) [41, 47, 44], algorithms that compute closed-form solutions to related or approximate versions of (1) [3, 28], and direct methods for constraint satisfaction such as the use of barriers [27] or immediate switching between the objective and constraint [55]. All of these approaches are designed to satisfy expected cost constraints for a single CMDP observed during training. In our work, on the other hand, we consider a distribution over possible transition models.

Our work focuses on irreducible uncertainty about the true environment at deployment time, which we refer to as model uncertainty and represent using a distribution $\mu$ over transition models. The most popular approach that incorporates model uncertainty in this way is domain randomization [48, 37], which randomizes across parameter values in a simulator and trains a policy to maximize average performance over this training distribution. This represents a risk-neutral attitude towards model uncertainty, which has been referred to as a soft-robust approach [16]. Distributionally robust MDPs incorporate robustness to the choice of $\mu$ by instead considering a set of distributions [53, 56, 15, 11], but application of this distributionally robust framework has remained limited in deep RL as it leads to a difficult minimax formulation that requires solving for worst-case distributions over transition models.

Robust RL represents an alternative approach to model uncertainty that considers uncertainty sets of transition models [21, 34]. A major drawback of robust RL is the need to calculate worst-case environments during training, which is typically approximated through the use of parametric uncertainty with multiple training environments [39, 31, 32] or a trained adversary that directly intervenes during trajectory rollouts [38, 46, 49]. Unlike these methods, we propose a robust approach to model uncertainty based on a distribution $\mu$ over models, which does not require access to a range of simulated training environments, does not impact data collection during training, and does not involve minimax optimization problems.

In contrast to irreducible model uncertainty, epistemic uncertainty captures estimation error that can be reduced during training through data collection. Epistemic uncertainty has been considered in the estimation of Q functions [35, 36, 9] and learned transition models [13, 25, 22, 40, 7], and has been applied to promote both exploration and safety in a fixed MDP. Finally, risk-sensitive methods typically focus on the aleatoric uncertainty in RL, which refers to the range of stochastic outcomes within a single MDP. Rather than considering the standard expected value objective, they learn risk-sensitive policies over this distribution of possible outcomes in a fixed MDP [43, 12, 45, 24, 26]. Distributional RL [8] trains critics that estimate the full distribution of future returns due to aleatoric uncertainty, and risk measures can be applied to these distributional critics for risk-sensitive learning [14, 29]. We also consider the use of risk measures in our work, but different from standard risk-sensitive RL methods we apply a risk measure over model uncertainty instead of aleatoric uncertainty.

Using the formulation of our RAMU Bellman operators from Lemma 1, we can leverage properties of distortion risk measures to efficiently estimate the results in (5) using sample-based weighted averages of standard Bellman targets. For $n$ transition models $p_{s,a}^{(i)}$, $i=1,\ldots,n$, sampled independently from $\mu_{s,a}$ and sorted according to their corresponding Bellman targets, consider the weights

where $g$ defines the distortion risk measure $\rho$ according to Definition 2. See Figure 1 for the distortion functions and weights associated with the risk measures used in our experiments. Then, from Jones and Zitikis [23] we have that

are consistent estimators of the results in (5), where $\mathcal{T}^{\pi}_{p^{(i)},r}Q(s,a)$ are sorted in ascending order and $\mathcal{T}^{\pi}_{p^{(i)},c}Q(s,a)$ are sorted in descending order. Finally, we can replace $\mathcal{T}^{\pi}_{p^{(i)},r}Q(s,a)$ and $\mathcal{T}^{\pi}_{p^{(i)},c}Q(s,a)$ with the standard unbiased sample-based estimates

where $s^{\prime}\sim p_{s,a}^{(i)}$ and $a^{\prime}\sim\pi(\,\cdot\mid s^{\prime})$. This leads to the sample-based estimates

which we use to train our RAMU Q functions. Note that the estimates in (7) can be computed very efficiently, which is a major benefit of our RAMU framework compared to robust RL methods. Next, we describe how we can sample models $p_{s,a}^{(i)}$, $i=1,\ldots,n$, from $\mu_{s,a}$, and generate state transitions from these models to use in the calculation of our sample-based Bellman targets in (7).

Note that our RAMU framework can be applied using any choice of distribution $\mu$, provided we can sample transition models $p_{s,a}^{(i)}\sim\mu_{s,a}$ and corresponding next states $s^{\prime}\sim p_{s,a}^{(i)}$. In this work, we define the distribution $\mu$ over perturbed versions of a single training environment $p^{\textnormal{train}}$, and we propose a generative approach to sampling transition models and corresponding next states that only requires data collected from $p^{\textnormal{train}}$. By doing so, our RAMU framework achieves robust performance and safety with minimal assumptions on the training process, and can even be applied to settings that require real-world data collection for training.

We consider a latent variable $x\sim X$, and we define a transition model $p_{s,a}(x)$ for every $x\sim X$ that shifts the probability of $s^{\prime}$ under $p_{s,a}^{\textnormal{train}}$ according to a perturbation function $f_{x}:\mathcal{S}\times\mathcal{S}\rightarrow\mathcal{S}$. This perturbation function takes as input a state transition $(s,s^{\prime})$, and outputs a perturbed next state $\tilde{s}^{\prime}$ that depends on the latent variable $x\sim X$. Therefore, a distribution over latent space implicitly defines a distribution $\mu_{s,a}$ over perturbed versions of $p_{s,a}^{\textnormal{train}}$. In order to obtain the next state samples needed to compute the Bellman target estimates in (7), we sample latent variables $x\sim X$ and apply $f_{x}$ to the state transition observed in the training environment. We have that $s^{\prime}\sim p_{s,a}^{\textnormal{train}}$ for data collected in the training environment, so $\tilde{s}^{\prime}=f_{x}(s,s^{\prime})$ represents the corresponding sample from the perturbed transition model $p_{s,a}(x)$.

In our experiments, we consider a simple implementation for the common case where $\mathcal{S}=\mathbb{R}^{d}$. We use uniformly distributed latent variables $x\sim U([-2\epsilon,2\epsilon]^{d})$, and we define the perturbation function as

where all operations are performed per-coordinate. Therefore, the latent variable $x\sim U([-2\epsilon,2\epsilon]^{d})$ can be interpreted as the percentage change in each dimension of a state transition observed in the training environment, where the average magnitude of the percentage change is $\epsilon$. The hyperparameter $\epsilon$ determines the distribution $\mu_{s,a}$ over transition models, where a larger value of $\epsilon$ leads to transition models that vary more significantly from the training environment. The structure of $f_{x}$ provides an intuitive, scale-invariant meaning for the hyperparameter $\epsilon$, which makes it easy to tune in practice. This choice of distribution $\mu_{s,a}$ captures general uncertainty in the training environment, without requiring specific domain knowledge of potential disturbances.

We summarize the implementation of our RAMU framework in Algorithm 1. Given data collected in a single training environment, we can efficiently calculate the sample-based RAMU Bellman targets in (7) by (i) sampling from a latent variable $x\sim X$, (ii) computing the corresponding next state samples $f_{x}(s,s^{\prime})$, and (iii) sorting the standard Bellman estimates that correspond to these sampled transition models. Given the sample-based RAMU Bellman targets, updates of the critics and policy have the same form as in standard deep safe RL algorithms. Therefore, our RAMU framework can be easily combined with many popular safe RL algorithms to incorporate model uncertainty with robustness guarantees, using only a minor change to the estimation of Bellman targets that is efficient to implement in practice.

First, we analyze the impact of our RAMU framework compared to standard safe RL. In both cases, we train policies using data collected from a single training environment, so the only difference comes from our use of risk-averse model uncertainty to learn RAMU Q functions. By evaluating the learned policies in perturbed test environments different from the training environment, we see that our RAMU framework provides robustness in terms of both total rewards and safety. In particular, the risk-averse implementation of our algorithm leads to safety constraint satisfaction in 80% of test environments, compared to only 51% with standard safe RL. In addition, this implementation results in higher total rewards (1.08x) and lower total costs (0.51x), on average. We see in Table 1 that the use of expectations over model uncertainty (i.e., a risk-neutral approach) also improves robustness in both the objective and constraint, on average, compared to standard safe RL. However, we further improve upon the benefits observed in the risk-neutral case by instead applying a risk-averse perspective.

Next, we compare our RAMU framework to domain randomization, a popular approach that also represents model uncertainty using a distribution $\mu$ over models. Note that domain randomization considers parametric uncertainty and has the benefit of training on a range of simulated environments, while our method only collects data from a single training environment. In order to evaluate the importance of domain knowledge for defining the training distribution in domain randomization, we consider two different cases: an in-distribution version that trains on a subset of the perturbed test environments, and an out-of-distribution (OOD) version that randomizes over a different perturbation parameter than the one varied at test time.

The results in Table 1 and Figure 2 show the importance of domain knowledge: in-distribution domain randomization leads to improved robustness compared to standard safe RL and the highest normalized average rewards (1.14x), while the out-of-distribution version provides little benefit. In both cases, however, domain randomization achieves lower levels of safety, on average, than our risk-averse formulation. In fact, we see in Figure 2 that the safety constraint satisfaction of our risk-averse formulation is at least as strong as both versions of domain randomization in 4 out of 5 tasks, despite only training on a single environment with no specific knowledge about the disturbances at test time. This demonstrates the key benefit of our risk-averse approach to model uncertainty.

Finally, we compare our approach to adversarial RL using the action-robust PR-MDP framework [46], which randomly applies worst-case actions a percentage of the time during data collection. Although adversarial RL only collects data from a single training environment, it requires potentially dangerous adversarial interventions during training in order to provide robustness at test time. In order to apply this method to the safe RL setting, we train an adversary to maximize costs and consider a 5% probability of intervention during training. The performance of adversarial RL is typically evaluated without adversarial interventions, which requires a clear distinction between training and testing.

We see in Figure 2 that adversarial RL learns policies that achieve robust safety constraint satisfaction at test time in the Quadruped tasks. Our risk-averse formulation, on the other hand, achieves higher levels of safety in the remaining 3 out of 5 tasks, and similar levels of safety on average. Unlike adversarial RL, our RAMU framework achieves robust safety in a way that (i) does not alter the data collection process, (ii) does not require training an adversary in a minimax formulation, and (iii) does not require different implementations during training and testing. In addition, our use of a distribution over models represents a less conservative approach than adversarial RL, resulting in higher normalized average rewards as shown in Table 1.