[Paper Review] Enabling risk-aware Reinforcement Learning for medical interventions through uncertainty decomposition

Paul Festor, Giulia Luise+2 more|arXiv (Cornell University)|Sep 16, 2021

Health Systems, Economic Evaluations, Quality of Life4 citations

TL;DR

This paper proposes a distributional reinforcement learning method, UA-DQN, that decomposes aleatoric and epistemic uncertainty in medical decision-making agents. By separating noise in environment dynamics from model knowledge limitations, the approach enables risk-aware clinical decision support with interpretable uncertainty estimates, validated in grid world and sepsis treatment MDPs showing distinct, behaviorally consistent uncertainty patterns.

ABSTRACT

Reinforcement Learning (RL) is emerging as tool for tackling complex control and decision-making problems. However, in high-risk environments such as healthcare, manufacturing, automotive or aerospace, it is often challenging to bridge the gap between an apparently optimal policy learnt by an agent and its real-world deployment, due to the uncertainties and risk associated with it. Broadly speaking RL agents face two kinds of uncertainty, 1. aleatoric uncertainty, which reflects randomness or noise in the dynamics of the world, and 2. epistemic uncertainty, which reflects the bounded knowledge of the agent due to model limitations and finite amount of information/data the agent has acquired about the world. These two types of uncertainty carry fundamentally different implications for the evaluation of performance and the level of risk or trust. Yet these aleatoric and epistemic uncertainties are generally confounded as standard and even distributional RL is agnostic to this difference. Here we propose how a distributional approach (UA-DQN) can be recast to render uncertainties by decomposing the net effects of each uncertainty. We demonstrate the operation of this method in grid world examples to build intuition and then show a proof of concept application for an RL agent operating as a clinical decision support system in critical care

Motivation & Objective

• To address the lack of uncertainty awareness in reinforcement learning for high-stakes medical interventions.
• To distinguish between aleatoric uncertainty (randomness in environment dynamics) and epistemic uncertainty (model knowledge limitations) in RL agents.
• To enable risk-aware decision-making by exposing uncertainty sources to clinicians in clinical decision support systems (CDSS).
• To validate the method in both synthetic grid world environments and a realistic sepsis treatment MDP derived from MIMIC-III data.
• To improve trust and interpretability of RL-based clinical recommender systems through uncertainty decomposition.

Proposed method

• Adapts distributional reinforcement learning (DQN) to model return distributions and decompose uncertainty into aleatoric and epistemic components.
• Uses a neural network to predict parameters of a distribution over returns, from which aleatoric and epistemic uncertainty are derived via variance decomposition.
• Applies uncertainty decomposition to estimate aleatoric uncertainty as the variance of the return distribution and epistemic uncertainty as the variance across multiple forward passes with dropout.
• Trains the agent on a Markov Decision Process (MDP) derived from the MIMIC-III critical care dataset, focusing on vasopressor and fluid dosing for septic patients.
• Validates the method in a grid world with stochastic wind dynamics to demonstrate correct uncertainty behavior under controlled conditions.
• Visualizes uncertainty estimates in state-action space and correlates epistemic uncertainty with state visitation frequency to validate model behavior.

Experimental results

Research questions

• RQ1 Can a distributional RL framework effectively decompose aleatoric and epistemic uncertainty in medical decision-making?
• RQ2 Do the estimated uncertainty components behave as expected in environments with known sources of randomness and data sparsity?
• RQ3 Can uncertainty decomposition improve the interpretability and trustworthiness of RL-based clinical decision support systems?
• RQ4 How do epistemic and aleatoric uncertainties correlate with data sparsity and environmental stochasticity in realistic critical care MDPs?
• RQ5 Can uncertainty estimates guide clinicians in distinguishing between random outcomes and model uncertainty in treatment recommendations?

Key findings

• In the grid world, aleatoric uncertainty was highest near the cliff edge where wind-induced stochasticity most affected outcomes, confirming correct sensitivity to environmental noise.
• Epistemic uncertainty was highest in states with low visitation frequency, particularly in the top-right corner, confirming sensitivity to data scarcity.
• In the sepsis MDP, epistemic uncertainty was anticorrelated with state visitation frequency, validating the method’s ability to detect data-scarce regions.
• The distribution of uncertainty estimates across the state-action space showed distinct, non-overlapping patterns for aleatoric and epistemic uncertainty, supporting successful decomposition.
• The system produced confidence scores (based on epistemic uncertainty) and outcome variability estimates (based on aleatoric uncertainty), suitable for integration into a clinical interface.
• The method demonstrated that high aleatoric uncertainty signals inherent randomness in patient outcomes, while high epistemic uncertainty indicates model uncertainty due to limited data, enabling better clinical risk assessment.