Boosted Distributional Reinforcement Learning: Analysis and Healthcare Applications

TL;DR

This paper introduces Boosted Distributional Reinforcement Learning (BDRL), a novel algorithm for healthcare decision-making that models outcome distributions, enforces comparability, and improves treatment consistency and quality-adjusted life years.

Contribution

The paper proposes BDRL, which optimizes outcome distributions for individual agents, enforces comparability, and stabilizes learning through a convex optimization-based projection.

Findings

BDRL improves the number of quality-adjusted life years.

BDRL enhances consistency of healthcare outcomes.

BDRL modifies treatments for median and vulnerable patients.

Abstract

Researchers and practitioners are increasingly considering reinforcement learning to optimize decisions in complex domains like robotics and healthcare. To date, these efforts have largely utilized expectation-based learning. However, relying on expectation-focused objectives may be insufficient for making consistent decisions in highly uncertain situations involving multiple heterogeneous groups. While distributional reinforcement learning algorithms have been introduced to model the full distributions of outcomes, they can yield large discrepancies in realized benefits among comparable agents. This challenge is particularly acute in healthcare settings, where physicians (controllers) must manage multiple patients (subordinate agents) with uncertain disease progression and heterogeneous treatment responses. We propose a Boosted Distributional Reinforcement Learning (BDRL) algorithm that optimizes agent-specific outcome distributions while enforcing comparability among similar agents and analyze its convergence. To further stabilize learning, we incorporate a post-update projection step formulated as a constrained convex optimization problem, which efficiently aligns individual outcomes with a high-performing reference within a specified tolerance. We apply our algorithm to manage hypertension in a large subset of the US adult population by categorizing individuals into cardiovascular disease risk groups. Our approach modifies treatment plans for median and vulnerable patients by mimicking the behavior of high-performing references in each risk group. Furthermore, we find that BDRL improves the number and consistency of quality-adjusted life years compared with reinforcement learning baselines.