SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Marc Rigter, Bruno Lacerda, Nick Hawes
Correspondence: marcrigter@gmail.com
Offline reinforcement learning (RL) is suitable for safety-critical domains where online exploration is not feasible. In such domains, decision-making should take into consideration the risk of catastrophic outcomes. In other words, decision-making should be risk-averse. An additional challenge of offline RL is avoiding distributional shift, i.e. ensuring that state-action pairs visited by the policy remain near those in the dataset. Previous offline RL algorithms that consider risk combine offline RL techniques (to avoid distributional shift), with risk-sensitive RL algorithms (to achieve risk-aversion). In this work, we propose risk-aversion as a mechanism to jointly address both of these issues. We propose a model-based approach, and use an ensemble of models to estimate epistemic uncertainty, in addition to aleatoric uncertainty. We train a policy that is risk-averse, and avoids high uncertainty actions. Risk-aversion to epistemic uncertainty prevents distributional shift, as areas not covered by the dataset have high epistemic uncertainty. Risk-aversion to aleatoric uncertainty discourages actions that are risky due to environment stochasticity. Thus, by considering epistemic uncertainty via a model ensemble and introducing risk-aversion, our algorithm (1R2R) avoids distributional shift in addition to achieving risk-aversion to aleatoric risk. Our experiments show that 1R2R achieves strong performance on deterministic benchmarks, and outperforms existing approaches for risk-sensitive objectives in stochastic domains.
In safety-critical applications, online reinforcement learning (RL) is impractical due to the requirement for extensive random exploration that may be dangerous or costly. To alleviate this issue, offline RL aims to find the best possible policy using only pre-collected data. In safety-critical settings we often want decision-making to be risk-averse, and take into consideration variability in the performance between each episode. However, the vast majority of offline RL research considers the expected value objective, and therefore disregards such variability.
A core issue in offline RL is avoiding distributional shift: to obtain a performant policy, we must ensure that the state-action pairs visited by the policy remain near those covered by the dataset. Another way that we can view this is that we need the policy to be averse to epistemic uncertainty (uncertainty stemming from a lack of data). In regions not covered by the dataset, epistemic uncertainty is high and there will be large errors in any learned model or value function. A naively optimised policy may exploit these errors, resulting in erroneous action choices.
In risk-sensitive RL, the focus is typically on finding policies that are risk-averse to environment stochasticity, i.e. aleatoric uncertainty. Previous approaches to risk-sensitive offline RL [47, 68]
combine two techniques: offline RL techniques, such as conservative value function updates [47] or behaviour cloning constraints [68], to avoid distributional shift; and risk-sensitive RL algorithms, to achieve risk-aversion. Therefore, the issues of epistemic uncertainty and aleatoric uncertainty are handled separately, using different techniques. In contrast, we propose to avoid epistemic uncertainty due to distributional shift, as well as aleatoric uncertainty due to environment stochasticity, by computing a policy that is risk-averse to both sources of uncertainty. This enables us to obtain a performant and risk-averse offline RL algorithm that is simpler than existing approaches.
In this work we propose 1R2R (Figure 1), a model-based algorithm for risk-averse offline RL. The intuition behind our approach is that the policy should avoid the risk of a bad outcome, regardless of whether this risk is due to high uncertainty in the model (epistemic uncertainty) or due to high stochasticity in the environment (aleatoric uncertainty). To represent epistemic uncertainty, we learn an ensemble of models. To optimise a risk-averse policy, we penalise taking actions that are predicted to have highly variable outcomes. In out-of-distribution regions, the disagreement between members of the ensemble is high, producing high variability in the model predictions. Therefore, the risk-averse policy is trained to avoid out-of-distribution regions. Risk-aversion also ensures that the policy avoids
Figure 1: Illustration of the 1R2R algorithm.
risk due to highly stochastic outcomes, i.e. avoids aleatoric risk. Our work demonstrates that incorporating risk-aversion alone is a promising approach to offline RL, and proposes an algorithm that is simpler than existing approaches to risk-averse offline RL [47, 68]. Our main contributions are (a) proposing risk-aversion towards epistemic uncertainty as a mechanism to address distributional shift in offline RL, and (b) introducing the first model-based algorithm for risk-averse offline RL, which jointly addresses epistemic and aleatoric uncertainty.
We evaluate our approach in several scenarios. Our experiments show that our algorithm achieves strong performance relative to state-of-the-art baselines on deterministic benchmarks [21], and outperforms existing approaches for risk-sensitive objectives in stochastic domains.
Offline RL: Offline RL addresses the problem of learning policies from fixed datasets. Most works on offline RL are risk-neutral, meaning that they optimise the expected value [53]. A key issue in offline RL is avoiding distributional shift so that the states visited by the learnt policy remain within the distribution covered by the dataset [42]. Approaches to model-free offline RL include: constraining the learnt policy to be similar to the behaviour policy [23, 38, 71, 22], conservative value function updates [11, 37, 40, 72], importance sampling algorithms [43, 51], or using Q-ensembles for more robust value estimates [1, 3, 73].
Model-based approaches learn a model of the environment and generate synthetic data from that model [63] to optimise a policy via planning [5] or RL algorithms [76, 75]. By training a policy on additional synthetic data, model-based approaches have the potential for broader generalisation [7]. Approaches to preventing model exploitation include constraining the trained policy to be similar to the behaviour policy [64] or applying conservative value function updates [75] in the same fashion as model-free approaches. Another approach is to modify the learnt MDP by either applying reward penalties for state-action pairs with high uncertainty [35, 44, 74, 76], or modifying the transition model to produce pessimistic synthetic transitions [29, 58]. Our work is similar in spirit to [29, 58], but these existing works do not consider risk.
Risk-Sensitive RL: Many works have argued for risk-sensitive algorithms in safety-critical domains [25, 48]. Risk-sensitive algorithms for MDPs include approaches that adapt policy gradients to risk objectives [12, 66, 65, 67]. More recently, a number of works have utilised distributional RL [9, 50], which learns a distributional value function that predicts the full return distribution. Knowledge of the return distribution can then be used to optimise a risk-sensitive criterion [15, 33, 46].
The main focus of research on risk-sensitive online RL is risk due to environment stochasticity [19] (i.e. aleatoric uncertainty). However, there are a number of works that explicitly address risk due to epistemic uncertainty [14, 16, 19, 18, 57]. To our knowledge, this is the first work to consider risk-aversion to epistemic uncertainty as a means to mitigate distributional shift in offline RL.
Risk-Sensitive Offline RL: To our knowledge, there are two previous algorithms for risk-sensitive offline RL: ORAAC [68] and CODAC [47]. To avoid distributional shift, ORAAC utilises policy constraints, while CODAC employs pessimistic value function updates to penalise out-of-distribution state-action pairs. Both algorithms learn a distributional value function [9] which is then used to optimise a risk-averse policy. Thus, both of these works are model-free, and combine techniques from offline RL and distributional RL to address the issues of distributional shift and risk-sensitivity, respectively. In this work, we provide a new perspective by proposing to optimise for risk-aversion towards both epistemic uncertainty (stemming from distributional shift) and aleatoric uncertainty (stemming from environment stochasticity). Unlike previous approaches, our approach (a) is conceptually simpler as it does not combine different techniques, (b) demonstrates that risk-aversion alone is sufficient to avoid distributional shift, (c) does not require a distributional value function (which might be computationally demanding to train), and (d) is model-based.
An MDP is defined by the tuple  $M = (S, A, T, R, s_0, \gamma)$ .  $S$  and  $A$  are the state and action spaces,  $R(s, a)$  is the reward function,  $T(s'|s, a)$  is the transition function,  $s_0$  is the initial state, and  $\gamma \in (0, 1)$  is the discount factor. In this work we consider Markovian policies,  $\pi \in \Pi$ , which map each state to a distribution over actions. In offline RL we only have access to a fixed dataset of transitions from the MDP:  $\mathcal{D} = \{(s_i, a_i, r_i, s_i')\}_{i=1}^{|\mathcal{D}|}$ . The goal is to find a performant policy using the fixed dataset.
Model-Based Offline RL These approaches utilise a model of the MDP. The learnt dynamics model,  $\widehat{T}$ , is typically trained via maximum likelihood estimation:  $\min_{\widehat{T}} \mathbb{E}_{(s,a,s') \sim \mathcal{D}}\left[-\log \widehat{T}(s'|s,a)\right]$ . A model of the reward function,  $\widehat{R}(s,a)$ , is also learnt if it is unknown. This results in the learnt MDP model:  $\widehat{M} = (S,A,\widehat{T},\widehat{R},s_0,\gamma)$ . Thereafter, any planning or RL algorithm can be used to recover the optimal policy in the learnt model,  $\widehat{\pi} = \arg \max_{\pi \in \Pi} J_{\widehat{M}}^{\pi}$ . Here,  $J$  indicates the optimisation objective, which is typically the expected value of the discounted cumulative reward.
However, directly applying this approach does not perform well in the offline setting due to distributional shift. In particular, if the dataset does not cover the entire state-action space, the model will inevitably be inaccurate for some state-action pairs. Thus, naive policy optimisation on a learnt model in the offline setting can result in model exploitation [32, 41, 55], i.e. a policy that chooses actions that the model erroneously predicts will lead to high reward. We propose to mitigate this issue by learning an ensemble of models and optimising a risk-averse policy.
Following previous works [35, 75, 76], our approach utilises model-based policy optimisation (MBPO) [32]. MBPO performs standard off-policy RL using an augmented dataset  $\mathcal{D} \cup \widehat{\mathcal{D}}$ , where  $\widehat{\mathcal{D}}$  is synthetic data generated by simulating short rollouts in the learnt model. To train the policy, minibatches of data are drawn from  $\mathcal{D} \cup \widehat{\mathcal{D}}$ , where each datapoint is sampled from the real data,  $\mathcal{D}$ , with probability  $f$ , and from  $\widehat{\mathcal{D}}$  with probability  $1 - f$ .
Risk Measures We denote the set of all probability distributions by  $\mathcal{P}$ . Consider a probability space,  $(\Omega, \mathcal{F}, P)$ , where  $\Omega$  is the sample space,  $\mathcal{F}$  is the event space, and  $P \in \mathcal{P}$  is a probability distribution over  $\mathcal{F}$ . Let  $\mathcal{Z}$  be the set of all random variables defined over the probability space  $(\Omega, \mathcal{F}, P)$ . We will interpret a random variable  $Z \in \mathcal{Z}$  as a reward, i.e. the greater the realisation of  $Z$ , the better. We denote a  $\xi$ -weighted expectation of  $Z$  by  $\mathbb{E}_{\xi}[Z] := \int_{\omega \in \Omega} P(\omega) \xi(\omega) Z(\omega) \mathrm{d}\omega$ , where  $\xi: \Omega \to \mathbb{R}$  is a function that re-weights the original distribution of  $Z$ .
A risk measure is a function,  $\rho : \mathcal{Z} \to \mathbb{R}$ , that maps any random variable  $Z$  to a real number. An important class of risk measures are coherent risk measures, which satisfy a set of properties that are consistent with rational risk assessments. A detailed description and motivation for coherent risk measures is given by [48]. An example of a coherent risk measure is the conditional value at risk (CVaR). CVaR at confidence level  $\alpha$  is the mean of the  $\alpha$ -portion of the worst outcomes.
All coherent risk measures can be represented using their dual representation [6, 61]. A risk measure,  $\rho$ , is coherent if and only if there exists a convex bounded and closed set,  $\mathcal{B}_{\rho} \in \mathcal{P}$ , such that
$$
\rho (Z) = \min  _ {\xi \in \mathcal {B} _ {\rho} (P)} \mathbb {E} _ {\xi} [ Z ] \tag {1}
$$
Thus, the value of any coherent risk measure for any  $Z \in \mathcal{Z}$  can be viewed as an  $\xi$ -weighted expectation of  $Z$ , where  $\xi$  is the worst-case weighting function chosen from a suitably defined risk envelope,  $\mathcal{B}_{\rho}(P)$ . For more details on coherent risk measures, we refer the reader to [61].
Static and Dynamic Risk In MDPs, the random variable of interest is usually the total reward received at each episode, i.e.  $Z_{\mathrm{MDP}} = \sum_{t=0}^{\infty} \gamma^{t} R(s_{t}, a_{t})$ . The static perspective on risk applies a risk measure directly to this random variable:  $\rho(Z_{\mathrm{MDP}})$ . Thus, static risk does not take into account the temporal structure of the random variable in sequential problems, such as MDPs. This leads to the issue of time inconsistency [10], meaning that actions which are considered less risky at one point in time may not be considered less risky at other points in time. This leads to non-Markovian optimal policies, and prohibits the straightforward application of dynamic programming due to the coupling of risk preferences over time.
This motivates dynamic risk measures, which take into consideration the sequential nature of the stochastic outcome, and are therefore time-consistent. Markov risk measures [60] are a class of dynamic risk measures which are obtained by recursively applying static coherent risk measures. Throughout this paper we will consider dynamic Markov risk over an infinite horizon, denoted by  $\rho_{\infty}$ . For policy  $\pi$ , MDP  $M$ , and static risk measure  $\rho$ , the Markov coherent risk  $\rho_{\infty}^{\pi}(M)$  is defined as
$$
\rho_ {\infty} ^ {\pi} (M) = R \left(s _ {0}, \pi \left(s _ {0}\right)\right) + \rho \left(\gamma R \left(s _ {1}, \pi \left(s _ {1}\right)\right) + \rho \left(\gamma^ {2} R \left(s _ {2}, \pi \left(s _ {2}\right)\right) + \dots\right)\right), \tag {2}
$$
where  $\rho$  is a static coherent risk measure, and  $(s_0,s_1,s_2,\ldots)$  indicates random trajectories drawn from the Markov chain induced by  $\pi$  in  $M$ . Note that in Equation 2 the static risk measure is evaluated at each step according to the distribution over possible successor states.
We define the risk-sensitive value function under some policy  $\pi$  as  $V^{\pi}(s,M) = \rho_{\infty}^{\pi}(M|s_0 = s)$ . This value function can be found by recursively applying the risk-sensitive Bellman equation [60]:
$$
V ^ {\pi} (s, M) = R (s, \pi (s)) + \gamma \min  _ {\xi \in \mathcal {B} _ {\rho} (T (s, \pi (s), \cdot))} \sum_ {s ^ {\prime}} T (s, \pi (s), s ^ {\prime}) \cdot \xi (s ^ {\prime}) \cdot V ^ {\pi} (s ^ {\prime}, M). \tag {3}
$$
Likewise, the optimal risk-sensitive value,  $V^{*}(s,M) = \max_{\pi \in \Pi}\rho_{\infty}(M|s_0 = s)$ , is the unique solution to the risk-sensitive Bellman optimality equation:
$$
V ^ {*} (s, M) = \max  _ {a \in A} \left\{R (s, a) + \gamma \min  _ {\xi \in \mathcal {B} _ {\rho} \left(T (s, a, \cdot)\right)} \sum_ {s ^ {\prime}} T (s, a, s ^ {\prime}) \cdot \xi \left(s ^ {\prime}\right) \cdot V ^ {*} \left(s ^ {\prime}, M\right) \right\}. \tag {4}
$$
Thus, we can view Markov dynamic risk measures as the expected value in an MDP where the transition probabilities at each step are modified adversarially within the risk-envelope for the chosen one-step static risk measure.
4 One Risk to Rule Them All
In this section, we present One Risk to Rule Them All (1R2R), a new algorithm for risk-averse offline RL (Figure 1). Our approach assumes that we can compute an approximation to the posterior distribution over MDPs, given the offline dataset. This distribution represents the epistemic uncertainty over the real environment, and enables us to reason about risk due to epistemic uncertainty. 1R2R then uses an RL algorithm to train a policy offline using synthetic data generated from the distribution over MDP models. To achieve risk-aversion, the transition distribution of the model rollouts is modified adversarially according to the risk envelope of a chosen risk measure. This simple modification to standard model-based RL penalises high uncertainty actions, thus penalising actions that have high epistemic uncertainty (i.e. are out-of-distribution) as well as those that have high aleatoric uncertainty (i.e. are inherently risky). Therefore, by modifying the transition distribution of model rollouts to induce risk-aversion, 1R2R avoids distributional shift and generates risk-averse behaviour.
To simplify the notation required, we assume that the reward function is known and therefore it is only the transition function that must be learnt. Note that in practice, we also learn the reward function
from the dataset. Given the offline dataset,  $\mathcal{D}$ , we assume that we can approximate the posterior distribution over the MDP transition function given the dataset:  $P(T\mid \mathcal{D})$ . We can integrate over all possible transition functions to obtain a single expected transition function:
$$
\bar {T} (s, a, s ^ {\prime}) = \int_ {T} T (s, a, s ^ {\prime}) \cdot P (T \mid \mathcal {D}) \cdot \mathrm {d} T. \tag {5}
$$
This gives rise to a single MDP representing the distribution over plausible MDPs,  $\overline{M} = (S, A, \overline{T}, R, s_0, \gamma)$ . This is similar to the Bayes-Adaptive MDP [27, 17, 28] (BAMDP), except that in the BAMDP the posterior distribution over transition functions is updated after each step of online interaction with the environment. Here, we assume that the posterior distribution is fixed given the offline dataset. To make this distinction, we refer to  $\overline{M}$  as the Bayesian MDP (BMDP) [4].
In this work, we pose offline RL as optimising the dynamic risk in the BMDP induced by the dataset.
Problem 4.1 (Risk-Averse Bayesian MDP for Offline RL). Given offline dataset,  $\mathcal{D}$ , static risk measure  $\rho$ , and belief over the transition dynamics given the dataset,  $P(T \mid \mathcal{D})$ , find the policy with the optimal dynamic risk in the corresponding Bayesian MDP  $\overline{M}$ :
$$
\pi^ {*} = \underset {\pi} {\arg \max } \rho_ {\infty} ^ {\pi} (\bar {M}) \tag {6}
$$
4.2 Motivation
The motivation for Problem 4.1 is that optimising for risk-aversion in the Bayesian MDP results in risk-aversion to both aleatoric and epistemic uncertainty. We first provide intuition for this approach via the following empirical example. Then, we demonstrate this theoretically for the case of Gaussian transition models in Proposition 4.2.
Illustrative Example Figure 2 illustrates a dataset for an MDP where each episode consists of one step. A single reward is received, equal to the value of the successor state,  $s'$ . We approximate  $P(T|\mathcal{D})$  using an ensemble of neural networks, and the shaded region represents the aggregated Bayesian MDP transition function  $\overline{T}(s,a,s')$ .
The dashed red line indicates the standard Q-values computed by drawing transitions from  $\overline{T}(s, a, s')$ . For this value function, the action with the highest Q-value is  $a = 1$ , which is outside of the data distribution. The solid red line indicates the risk-sensitive value function for the risk measure  $\mathrm{CVaR}_{0.1}$  (i.e.  $\alpha = 0.1$ ). This is computed by drawing transitions from a uniform distribution over the worst  $10\%$  of transitions from  $\overline{T}(s, a, s')$ . We observe that the risk-sensitive value function penalises straying far from the dataset, where epistemic uncertainty is high. The action in the dataset with highest expected value is  $a = 0.4$ . However, the action  $a = -0.35$  has a higher risk-sensitive value than  $a = 0.4$ , due to the latter having high-variance transitions (i.e. high aleatoric uncertainty). Thus, we observe that choosing actions according to the risk-sensitive value function penalises actions that are out-of-distribution (i.e. have high epistemic uncertainty) or have high aleatoric uncertainty.
Figure 2: MDP with a single transition to a successor state  $s'$ , with reward  $R(s') = s'$ . Black dots illustrate dataset. Shaded region indicates  $\pm 2$  S.D. of successor states drawn from an approximation of  $\overline{T}(s, a, \cdot)$  (represented by an ensemble [13]). Dashed red line indicates the standard Q-values from running MBPO [32]. Solid red line indicates the risk-sensitive value function computed by reweighting the transitions according to  $\mathrm{CVaR}_{0.1}$ . Vertical arrows indicate optimal actions for each Q-function.
We now consider the specific case where each member of the ensemble is a Gaussian transition function with standard deviation  $\sigma_{A}$ . We assume that in the ensemble, the mean of each Gaussian is also normally distributed with standard deviation  $\sigma_{E}$ . The latter assumption allows us to quantify the epistemic uncertainty in terms of the level of disagreement between members of the ensemble. Proposition 4.2 shows that risk-averse optimisation of transitions sampled from the ensemble results in risk-aversion to both the aleatoric and epistemic uncertainty.
Proposition 4.2. Consider some state  $s$  in a 1D state space, and some action  $a$ . Assume (a) that there is an ensemble of  $N$  Gaussian transition functions, each denoted  $T_{i}$  with mean  $\mu_{i}$  and standard deviation  $\sigma_{A}$ :  $\{T_{i}(s'|s,a) = \mathcal{N}(\mu_{i},\sigma_{A}^{2})\}_{i=1}^{N}$ ; and (b) that the mean of each Gaussian,  $\mu_{i}$ , is normally distributed with mean  $\mu_{0}$  and standard deviation  $\sigma_{E}$ :  $\mu_{i} \sim \mathcal{N}(\mu_{0},\sigma_{E}^{2})$ . The  $N$  transition functions jointly define  $\overline{T}$ , the transition function for Bayesian MDP,  $\overline{M}$ . Assume that for some policy  $\pi$ , the risk-sensitive value is linear around  $\mu_{0}$  with some linearity constant,  $K$ :
$$
V ^ {\pi} \left(s ^ {\prime}, \bar {M}\right) = V ^ {\pi} \left(\mu_ {0}, \bar {M}\right) + K \left(s ^ {\prime} - \mu_ {0}\right)
$$
Then, the value of a randomly drawn successor state is distributed according to  $\mathcal{N}\Big(\mu = V^{\pi}(\mu_0,\overline{M}),\sigma^2 = K^2 (\sigma_A^2 +\sigma_E^2)\Big)$ .
In Proposition 4.2,  $\sigma_{E}$  defines the level of disagreement between the members of the ensemble, and therefore represents the level of epistemic uncertainty.  $\sigma_{A}$  represents the aleatoric uncertainty. The proposition shows that if either  $\sigma_{A}$  or  $\sigma_{E}$  are high, then there is high variability over the value of the successor state when sampling from the ensemble (i.e. sampling from  $\overline{T}$  in Equation 5). Risk measures penalise high variability. Therefore, applying the risk-sensitive Bellman equation (Equation 3) to samples from  $\overline{T}$  penalises executing state-action pairs for which either  $\sigma_{A}$  or  $\sigma_{E}$  is high. For the specific case of CVaR, Corollary 4.3 illustrates how the risk-sensitive value is reduced as  $\sigma_{A}$  or  $\sigma_{E}$  increase.
Corollary 4.3. Under the assumptions in Proposition 1, the CVaR at confidence level  $\alpha$  of the value of a randomly drawn successor state is: