Echoes of Socratic Doubt: Embracing Uncertainty in
Calibrated Evidential Reinforcement Learning
(Supplementary Material)

Artificial intelligence (AI) has seen incredible growth in recent years in terms of interest, research, applications, and performance. Anyone with even basic knowledge of the technology has been able to utilize its tools for processing unprecedented amounts of data from audio, visual, and/or textual sources. However, the insights from the models are only as good as the data they consume regardless of the level of supervision in training. Data, especially from the real world, naturally and often randomly contain noise and inconsistencies, which lead to uncertainty in model predictions and can potentially degrade performance. The two main types of uncertainty that affect the models are referred to as aleatoric and epistemic, where the former is the result of random distortions in the data and is typically unresolvable, and the latter is the result of insufficient knowledge, which can be mitigated with more data and training. Deep reinforcement learning (DRL) is a subset of machine learning (ML) under the umbrella of AI that involves using neural networks to solve sequential decision-making problems that require interaction between an agent and an environment. Recent applications of DRL include drone racing \citepdroneracing2023, healthcare diagnosis and treatment \citephealthcareRL2021, and chip design \citepchipdesignRL2021, which are all distinct tasks. As such, there are several types of DRL algorithms, although in this study we focus on deep Q networks with discrete action spaces \citepMnih2015. In regards to processing data, a key difference between most DRL algorithms and other ML algorithms is their foundation in the Markov Decision Process (MDP) \citepPuterman1994, a mathematical framework for discrete-time decision-making when the desired outcome (e.g., maximum cumulative reward) is stochastic yet partially amenable to a user. Accordingly, DRL is susceptible to uncertainty, which is reflected directly in each interaction and the end result. However, due to the nature of the algorithm’s interactive data processing, one is given explicit parametric control over whether to exploit or explore the environment when seeking rewards, which can be manipulated with uncertainty awareness. Unlike most supervised ML tasks, where uncertainty awareness has seen the most interest lately, reinforcement learning can directly and explicitly incorporate reliable uncertainty measures as feedback in the exploitation/exploration trade-off and choosing actions. There has been a number of studies exploring solutions to important challenges in uncertainty-aware reinforcement learning \citeplockwood2022. One challenge involves ensuring statistical rigor and efficiency in estimating both aleatoric and epistemic uncertainty, whereas most previous works have focused only on epistemic uncertainty due to the intractability of conventional uncertainty estimation methods in DRL (e.g. Bayesian sampling). Another challenge involves accurately separating the two quantities since epistemic uncertainty is often a function of aleatoric uncertainty; an agent unaware of sources of aleatoric uncertainty in the environment (e.g., random traps) may take actions that neglect exploring related sources of epistemic uncertainty (e.g., shortcuts), hindering its ability to gain knowledge of potentially better rewards. Sources of either uncertainty can generally yield positive or negative rewards, but the only way the agent will learn this effectively is through uncertainty awareness, the very essence of socratic doubt. An example of this relationship is given in Fig. 1. A third challenge involves properly handling observations that significantly deviate from the norm (i.e., out-of-distribution or OOD) \citepyang2021oodsurvey, which tend to negatively affect both algorithm performance and traditional uncertainty measures that rely on restrictive assumptions such as Gaussianity. Lastly, unlike traditional ML, the target in DRL is not fixed (whether it is a single value or a distribution), and samples are not considered traditionally independent and identically distributed (i.i.d) so certain theoretical bases become invalid. Therefore, the uncertainty quantification in DRL needs to be highly adaptive and flexible, yet stable and capable of empirically accommodating violations of typical critical assumptions. To address these challenges of uncertainty estimation in DRL, we propose a novel combination of properties of conformal inference \citepvovk2005,shafer2007,angelopoulos2021 and deep evidential learning \citepsensoy2018,amini2020 to further enhance the usage of quantile regression as a basis for distributional reinforcement learning involving deep Q networks (i.e., QR-DQN) \citepbellemare2017distributional,dabney2017distributional. Conformal inference is a simple, widely applicable, and scalable statistical framework for constructing reliable prediction intervals using a finite set of data with guaranteed marginal coverage of the true realization of interest, subject to a predetermined confidence rate and a relevant calibration dataset. It excels at adaptation to complex, heteroskedastic data and fortifies estimation of aleatoric and epistemic uncertainty when paired with an inherently uncertainty-aware algorithm such as quantile regression \citepromano2019, without making strong theoretical assumptions. However, conformal inference alone can be overly conservative and does not perform well on new OOD observations without additional training and/or calibration. On the other hand, evidential learning involves modeling a higher-order probability distribution placed over the likelihood function through learned parameters, which bypasses the need for repetitive sampling typically attributed to Bayesian techniques and allows one to explicitly compute aleatoric and epistemic uncertainty separately. In practice, it is equally suitable as an accurate uncertainty estimation method, but is much superior in handling OOD observations. Nonetheless, evidential learning is more computationally complex and therefore requires careful optimizations to maintain stability depending on the data. By combining the two complementary concepts into one algorithm, we can leverage their benefits and mitigate their weaknesses. Thus, the proposed study makes the following key contributions:

• [itemsep=2pt,topsep=2pt,leftmargin=10pt]

• •

  We introduce a novel framework that combines elements of conformal inference and deep evidential learning to improve the usage of quantile regression in deep distributional reinforcement learning. The algorithm incorporates separate, explicit, and statistically robust computations of aleatoric and epistemic uncertainty as feedback for balancing exploration and exploitation in action selection.

• •

  In utilizing properties of conformal inference, the estimated quantiles of the distribution of returns over each action taken in the environment are calibrated with the distributional Bellman target \citepbellemare2017distributional,dabney2017distributional according to a desired marginal coverage of $90\%$. This enhances the accuracy of the Q-value quantiles as training samples are selected from experience replay. Additionally, the estimated quantiles of each learned evidential distribution (which are used to compute uncertainty) are similarly calibrated, allowing the algorithm to tighten action selection and maintain high score performance well into training.

• •

  Combining calibrated quantile regression with deep evidential learning enhances the algorithm’s uncertainty awareness by providing global network uncertainty estimates over the entire known sample space as opposed to local estimates based solely on the variance and expectation of Q-value quantiles, which can be biased. As a result, the proposed CEQR-DQN algorithm achieves meets or exceeds scores on benchmarks such as MinAtar \citepyoung19minatar set by similar previous work with notably faster learning.

In Section 2, we briefly review work related to our proposed solution for uncertainty awareness in DRL. In Section 3, we discuss the algorithm and underlying theory in detail. In Section 4, we present simulation results. Lastly, Section 5 concludes the study.

The study of uncertainty awareness in deep reinforcement learning has remained under active investigation for several years, in part due to its undeniable impact on returns. Different approaches to the topic have attempted to address the many challenges associated with parameterizing, stabilizing, and separating uncertainty estimates subject to the stochastic nature of reinforcement learning. In this study, we build upon a particular thread of influential uncertainty-aware work that involve deep Q learning, a simple, model-free DRL algorithm that focuses on approximating the Q-value function given a discretized action space (i.e., DQN) \citepMnih2015. Bootstrapped DQN \citetosband2016 was one of the first algorithms to address exploration efficiency through incorporating a sense of uncertainty in choosing actions. As an alternative to Thompson sampling \citepthompson1933, their approach involves using multiple neural network heads with random initializations to provide a notion of uncertainty and encourage an agent to dynamically choose actions based on knowledge (or lack thereof) of potential rewards, as opposed to an $\epsilon$-greedy policy. However, redundant approximations require computational overhead and this method of uncertainty estimation is rather rudimentary as random initilizations do not guarantee diversity or generalizability. Double Uncertain Value Network (DUVN) \citetmoerland2017, the first distributional DRL algorithm to separate and propagate both aleatoric and epistemic uncertainty, offers a more nuanced exploration strategy in comparison to bootstrapped DQN, which leads to better overall outcomes. Nonetheless, it is similarly unsophisticated in its uncertainty estimation as it uses dropout as an indicator of epistemic uncertainty and assumes that the return distribution (from which aleatoric uncertainty is extracted) is Gaussian. Furthermore, it is also computationally expensive as it requires careful optimization of two separate networks for the estimation, neglecting to account for their interaction. Uncertainty-Aware DQN (UA-DQN) \citetclements2020 builds upon the previous work by adopting the quantile regression-based distributional DRL solution proposed in \citepdabney2017distributional, establishing two bootstrapped auxiliary networks to assess the uncertainty quantities, and intelligently integrating MAP sampling to enhance uncertainty awareness. It is shown to properly estimate the uncertainties and be superior in performance, but is subject to the same vulnerabilities as other methods that rely on redundancy and random sampling. Acknowledging the advantages and disadvantages of the aforementioned algorithms, the proposed work is inspired to maintain proper disentanglement in estimating global aleatoric and epistemic uncertainty without bootstrapping auxiliary networks. Additionally, we seek to provide a more statistically rigorous algorithm that can empirically handle OOD observations even in highly complex environments.

This section details the proposed framework for deep uncertainty-aware distributional reinforcement learning involving calibrated quantile regression and evidential learning (i.e., calibrated evidential quantile regression in deep Q networks or CEQR-DQN). The scalable architecture is shown in Fig. 6 in the appendix; it adapts the original DQN feature extractor \citepMnih2015, and includes separate linear network heads for outputting quantiles for actions and the parameters of the empirically learned evidential distributions. We designed this model and the underlying methodology to be closely comparable to \citepclements2020 and easy to adapt to other reinforcement learning solutions. In accordance with the original QR-DQN \citepdabney2017distributional, our algorithm seeks to learn the distribution of returns $Z^{\pi}(x,a)$ (i.e., the state-action value distribution) upon taking action $a$ in state $x$ according to a policy $\pi$ using quantile regression. It follows the Markov Decision Process of modeling interactions between the agent and the environment based on the variables $X$, $A$, $R$, $P$, and $\gamma$, where $X$ represents the state space, $A$ represents the action space, $R$ represents the stochastic reward function, $P$ represents the probability of transitioning states $x\rightarrow x^{\prime}$ after taking an action $a$, and $\gamma$ is the discount factor controlling the weight of future $Z^{\pi}(x,a)$ in loss calculations. Regarding loss, the value distribution is approximated via the temporal difference-based (TD) distributional Bellman target $T^{\pi}Z(x,a)\stackrel{{\scriptstyle d}}{{:=}}R(x,a)+\gamma Z(x^{\prime},a^{\prime})$ where $a^{\prime}\sim$ argmax${}_{a^{\prime}}\mathbb{E}[Z(x^{\prime},a^{\prime})]$ and $x^{\prime}\sim P(\cdot\,|\,x,a)$. It is parameterized with a quantile distribution that is learned by minimizing the quantile Huber loss

with real-valued quantile functions of the cumulative distribution function $F$ (i.e., $\theta=F^{-1}_{Z}(\hat{q})$), where

is the Huber loss \citephuber1964, $q\in[0,1]$ represents a target quantile, $\hat{q}$ represents an estimated quantile, $e$ is error, $\mathbb{I}$ is the indicator function, and $\kappa$ is a hyperparameter used to assert squared error within $[-\kappa,\kappa]$ distance of zero, or standard error otherwise. Quantile huber loss is a smoother version of the traditional quantile regression loss that is convex and asymmetric, penalizing overestimation and underestimation of an estimated conditional quantile. Applying this framework to deep Q learning, $e\equiv TZ_{\theta}(x,a)-\theta_{i}(x,a)$ where $Z_{\theta}$ represents the parameterized quantile distribution of returns mapping state-action pairs to a uniform probability distribution of $N$ quantiles with uniform weight and cumulative probabilities $\tau_{i}=\frac{1}{N}$, $i=1,\dots,N$ given $\theta:X\times A\rightarrow\mathbb{R}^{N}$. Accordingly, the network outputs $N$ quantiles for each discretized action and the resultant quantile regression loss is calculated as

Notably, we approximate only one target value distribution $Z_{\theta}$ per update as opposed to multiples contained in a set of all possible distributions $Z_{Q}$ as is done in ensemble methods. Different methods for DRL, such as policy gradient, actor-critic, and model-based methods, can incorporate uncertainty measures in various other aspects of agent-environment interaction. QR-DQN is an off-policy method and, as such, must incorporate uncertainty measures related to the value distribution. Therefore, we focus on uncertainty estimation in action selection as it is the primary component in maximizing returns. As opposed to other works following a similar procedure, the proposed method utilizes evidential learning to evaluate and exploit uncertainty in choosing actions that maximize the target value distribution. Actions are selected via Thompson sampling of a multivariate normal distribution based on the expectation of action quantiles and the covariance between the evidential uncertainty associated with each action. For quick reference, the pseudocode for action selection, which includes evidential learning, is given in Algorithm 3 and explained in Section 3. {algorithm}[!t] Calibrated Evidential QR-DQN Action Choice {algorithmic} \Require$N$ quantiles, action space $A$, state space $X$, network $\mathbb{N}$, evidential distribution parameters $G=\{\gamma,v,\alpha,\beta\}_{evi}$, and uncertainty hyperparameters $\lambda_{al}$ and $\lambda_{ep}$. \Ensure$\{\lambda_{al},\lambda_{ep},v,\beta\}>0$; {$N,\alpha\}>1$ \StateGather quantiles of actions and evidential parameters given state $x$ in $X$ =¿ $A_{N},G_{N}=\mathbb{N}(x)$ \For$a$ in $A$ \StateCalculate action mean over N: $\mu=\frac{1}{N}\sum_{i=1}^{N}a_{i}$ \StateCalculate evidential SD: $\sigma_{evi}=\sqrt{\left(\frac{\beta}{v(\alpha-1)}\right)}$ \StateCalculate uncertainty for evidential $5^{th}$ percentile: \State$u^{5^{th}}_{evi}=[|(\gamma^{5^{th}}_{evi}+\sigma^{5^{th}}_{evi})-(\gamma^{5^{th}}_{evi}-\sigma^{5^{th}}_{evi})|]$ \StateCalculate uncertainty for evidential $95^{th}$ percentile: \State$u^{95^{th}}_{evi}=[|(\gamma^{95^{th}}_{evi}+\sigma^{95^{th}}_{evi})-(\gamma^{95^{th}}_{evi}-\sigma^{95^{th}}_{evi})|]$ \StateCalculate total epistemic uncertainty: \State$\psi_{ep}=\mathbb{E}_{N}[\frac{1}{2}(u^{5^{th}}_{evi}+u^{95^{th}}_{evi})]$ \StateCalculate total aleatoric uncertainty: \State$\psi_{al}=\mathbb{E}_{N}[|\gamma^{95^{th}}_{evi}-\gamma^{5^{th}}_{evi}|]$ \EndFor\StateRisk aversion (optional): $\mathcal{M}\leftarrow\mathcal{M}-\lambda_{al}\Psi_{al}$ \StateDraw Thompson sample: $S_{A}\sim\mathcal{N}(\mathcal{M},\lambda_{ep}\Psi_{ep})$ \State(where $\mathcal{M}=\{\mu_{1},\dots,\mu_{A}\}$; $\Psi_{el,ap}=\{\psi_{1},\dots,\psi_{A}\}$) \StateOutput action: argmax${}_{a}[S_{A}]$

In this section we discuss results of CEQR-DQN on the miniaturized Atari testbed MinAtar \citepyoung19minatar, which provides simplified versions of the games Asterix, Breakout, Freeway, Seaquest, and Space Invaders in comparison to the seminal Arcade Learning Environment \citepbellemare2013. The games are simplified by reducing their observation feature space by over $20\times$, enforcing a smaller action space, and modifying certain environmental parameters so as to encourage behavioral learning over representation learning. Effectively, it makes evaluating discrete DRL algorithms faster and more easily reproducible. Furthermore, it encourages increased stochasticity through randomized spawn points, difficulty ramping, sticky actions, etc. However, we disable sticky actions, which forces the agent to repeat the previous action with $10\%$ probability, as they are counter-intuitive to uncertainty-aware exploration strategies and hinder our quantile calibration efforts due to such random actions being included in gradient descent. All other MinAtar settings were left default. As is shown in Fig. 3, our algorithm demonstrates considerable improvements in terms of learning speed and score over UA-DQN \citepclements2020, which set the previous high scores. The two algorithms share hyperparameters, which are provided in Table 7 in the appendix. These improvements indicate that statistically robust uncertainty quantification, including calibration and modulated confidence upon OOD observations, can significantly alter an agent’s exploration strategy for the better despite the environment being partially observable as is the case in each game except for Breakout. Furthermore, since each evidential distribution is learned by the network through consideration of all observations, their derived uncertainty quantities allow the agent to maintain global uncertainty awareness in action choice as opposed to local awareness based on variance between two continuously changing auxiliary networks. The empirical benefit of this is especially demonstrated in Seaquest, where rewards are complicated to obtain. Since in both algorithms the uncertainty estimates factored into action selection are based on evaluating the distribution of returns $Z_{\theta}$, such a difficult reward structure decreases their efficacy, especially when localized to a single observation. On the other hand, in Freeway, rewards are scarce, so the less complex a model is, the better the performance. Concurrently, disabling sticky actions deepens the reliance of each algorithm’s performance on uncertainty awareness and further emphasizes the need for non-localized estimates. The agents in the algorithms each take between $5\%-30\%$ uncertain actions throughout training, yet the difference in results are appreciable.

We presented a novel framework for quantifying, calibrating, and exploiting aleatoric and epistemic uncertainty in distributional reinforcement learning involving quantile regression-based deep Q networks. The proposed methodology, applied to a suite of simplified miniature Atari games, combines deep evidential learning with quantile calibration based on properties of conformal inference to significantly improve agent exploration through uncertainty-aware action selection. This approach encourages an agent to learn faster and seek greater rewards in each successive training episode. The underlying theory is both easily adaptable and scalable, enabling its application to any task or dataset that necessitates consideration of uncertainty, particularly in scenarios with frequent out-of-distribution (OOD) observations. {acknowledgements} This work was supported in part by COGNISENSE, one of seven centers in JUMP 2.0, a Semiconductor Research Corporation (SRC) program sponsored by DARPA. \appendix

Fig. 6 depicts the neural network model used for the proposed CEQR-DQN algorithm.

Table 7 lists the hyperparameters used by the proposed CEQR-DQN algorithm to achieve the results showcased in Section 4 on MinAtar \citepyoung19minatar, a miniaturized suite of Atari 2600 games. We used code published by \citetclements2020 to produce the results for UA-DQN. Code for CEQR-DQN will be released at: \urlhttps://github.com/acstutts/CEQR-DQN.git