Offline Bayesian Aleatoric and Epistemic Uncertainty Quantification and Posterior Value Optimisation in Finite-State MDPs

TL;DR

The paper tackles offline Bayesian decision-making under model uncertainty in finite-state MDPs by combining exact moment-based uncertainty quantification with Bayesian posterior value optimization. It leverages tractable DP formulations to compute the first two moments of the return and uses the law of total variance to separate epistemic from aleatoric uncertainty, formulating a posterior-valued objective: $\max_\pi \sum_s \rho(s) \mathbb{E}_{\mathcal{M}\sim p(.|D)}[V_\mathcal{M}^\pi(s)]$. A stochastic gradient-based policy optimiser uses posterior resampling to maximize this objective without strong posterior-distribution assumptions, demonstrated on gridworlds, synthetic MDPs, and the AI Clinician/MIMIC datasets. The results show improved posterior value, especially in low-data regimes, and the approach scales to practical offline settings while highlighting conservatism as a robust strategy when data coverage is limited. The work provides a principled, scalable framework for uncertainty-aware offline control in discrete, finite-state environments, with code available for replication.

Abstract

We address the challenge of quantifying Bayesian uncertainty and incorporating it in offline use cases of finite-state Markov Decision Processes (MDPs) with unknown dynamics. Our approach provides a principled method to disentangle epistemic and aleatoric uncertainty, and a novel technique to find policies that optimise Bayesian posterior expected value without relying on strong assumptions about the MDP's posterior distribution. First, we utilise standard Bayesian reinforcement learning methods to capture the posterior uncertainty in MDP parameters based on available data. We then analytically compute the first two moments of the return distribution across posterior samples and apply the law of total variance to disentangle aleatoric and epistemic uncertainties. To find policies that maximise posterior expected value, we leverage the closed-form expression for value as a function of policy. This allows us to propose a stochastic gradient-based approach for solving the problem. We illustrate the uncertainty quantification and Bayesian posterior value optimisation performance of our agent in simple, interpretable gridworlds and validate it through ground-truth evaluations on synthetic MDPs. Finally, we highlight the real-world impact and computational scalability of our method by applying it to the AI Clinician problem, which recommends treatment for patients in intensive care units and has emerged as a key use case of finite-state MDPs with offline data. We discuss the challenges that arise with Bayesian modelling of larger scale MDPs while demonstrating the potential to apply our methods rooted in Bayesian decision theory into the real world. We make our code available at https://github.com/filippovaldettaro/finite-state-mdps .

Figures (13)

• Figure 1: Fig. \ref{['fig:gridworld']} shows the gridworld used in the experiments. The terminal states are the failure F states (cliff) and the goal G state. The agent can move up, down, left, or right (or remain stationary if it hits the boundary of the grid). The transition dynamics have intrinsic stochasticity controlled by the probability $p_{\text{rand}}$, which is the probability of pushing the agent down regardless of action taken. Offline training datasets were generated by randomly sampling actions at random non-terminal states. State $\bigstar$ is chosen as an exemplar state to plot state-dependent uncertainties. In Fig. \ref{['fig:stochasticity']}, the plot shows the epistemic (blue) and aleatoric (red) standard deviations as a function of training dataset size, with different levels of intrinsic stochasticity indicated by solid, dashed, and dotted lines.
• Figure 2: Fig. \ref{['fig:gridbayesianvalues']} shows the average posterior expected return ('Value') as a function of dataset size for a single set of generated datasets, as in the objective in Eq. \ref{['eq:objective']}. The example gridworld has $p_{\text{rand}}=0.25$. As value will be dataset-dependent, we show the average and standard deviation between the pairwise difference in posterior values between ours and the other methods in Figs. \ref{['fig:gradmle']}, \ref{['fig:gradnominal']} and \ref{['fig:gradsecondorder']}, where values above the red dashed line signify an improvement. These plots report the average and standard deviation across 50 generated datasets for each dataset size.
• Figure 3: Ground truth pairwise difference in average performance (and shaded standard error of the mean) on the policies found by each method and rolled out on the ground-truth synthetic MDP. Regions above the red line correspond to improved performance with our method.
• Figure 4: Average and standard deviation (shaded) of posterior expected value difference achieved by our method. Regions above the red line correspond to improved objective optimisation with our method.
• Figure 5: Posterior of each state (blue dots) under our policy and the MLE-optimal policy in the clinical MDP. Points above the diagonal indicate superior performance of our policy on the posterior expect value. The left plot (a) demonstrates the impact of policy choice on performance when employing Bayesian model selection with an optimal parameter of $\alpha_p=0.072$. The right plot (b) shows the same result when using a prior selected through a conservative sparse dynamics model.