Continuous State-Space Models for Optimal Sepsis Treatment - a Deep Reinforcement Learning Approach

Aniruddh Raghu; Matthieu Komorowski; Leo Anthony Celi; Peter Szolovits; Marzyeh Ghassemi

Source-linked AI summary

Continuous State-Space Models for Optimal Sepsis Treatment - a Deep Reinforcement Learning Approach

Aniruddh Raghu, Matthieu Komorowski, Leo Anthony Celi, Peter Szolovits, Marzyeh Ghassemi

arXiv:1705.08422v1May 23, 2017cs.LG

TL;DR

Sepsis care requires individualized treatment decisions despite limited bedside decision-support tools and unclear universally optimal strategies. The paper uses continuous-state deep reinforcement learning to learn treatment policies from ICU data, reporting potential mortality reductions of 1.8–3.6% from a 13.7% baseline. The policies are described as clinically interpretable and potentially useful for ICU decision support.

Problem
Sepsis treatment is difficult because patients respond differently, treatment strategies vary, and clinicians lack efficient individualized real-time decision support.
Method
The paper uses deep reinforcement learning with continuous physiological state representations and discretized treatment actions to learn sepsis policies from ICU data.
Results
1.8–3.6% potential reduction in hospital mortality is reported from a baseline mortality of 13.7%.
Takeaways & Limitations
The learned policies are clinically interpretable and could support intensive care clinicians in medical decision making.
Takeaways & Limitations
Offline retrospective data limit state-space exploration and make learning the truly optimal medication policy difficult.

Abstract

from arXiv · show

Sepsis is a leading cause of mortality in intensive care units (ICUs) and costs hospitals billions annually. Treating a septic patient is highly challenging, because individual patients respond very differently to medical interventions and there is no universally agreed-upon treatment for sepsis. Understanding more about a patient's physiological state at a given time could hold the key to effective treatment policies. In this work, we propose a new approach to deduce optimal treatment policies for septic patients by using continuous state-space models and deep reinforcement learning. Learning treatment policies over continuous spaces is important, because we retain more of the patient's physiological information. Our model is able to learn clinically interpretable treatment policies, similar in important aspects to the treatment policies of physicians. Evaluating our algorithm on past ICU patient data, we find that our model could reduce patient mortality in the hospital by up to 3.6% over observed clinical policies, from a baseline mortality of 13.7%. The learned treatment policies could be used to aid intensive care clinicians in medical decision making and improve the likelihood of patient survival.

1. Introduction

Sepsis treatment requires difficult patient-specific decisions, while bedside clinicians lack efficient individualized decision-support tools. The paper proposes continuous-state deep reinforcement learning to discover interpretable treatment policies and potentially reduce mortality.

Sepsis treatment varies across fluids, vasopressors, antibiotics, corticosteroids, timing, and dosage.
The approach applies deep reinforcement learning to discover optimal sepsis treatment strategies from clinical data.
Reinforcement learning accommodates stochastic patient responses, delayed treatment efficacy, and suboptimal training examples.
Continuous-state modeling represents physiology as a continuous vector, using raw data or sparse latent representations, while Deep-Q Learning selects actions.
1.8–3.6% potential mortality reduction is reported from a baseline hospital mortality of 13.7%.
The learned policies are investigated for clinical interpretability and potential use as clinical decision-support tools.

2. Background and Related Work

The paper frames sepsis treatment as a sequential reinforcement-learning problem and extends prior discretized clinical RL work toward continuous physiological states. It uses off-policy learning to derive treatment policies from physician-generated data.

2.1 Reinforcement Learning: Reinforcement learning models sequential states, actions, rewards, and transitions while maximizing discounted future reward.
2.1 Reinforcement Learning: Q-learning estimates optimal action values with the Bellman equation and learns off-policy from transition samples.
2.1 Reinforcement Learning: The model uses patient physiology as state, 25 discretized fluid–vasopressor actions, and terminal survival rewards.
2.2 Reinforcement Learning in Health: RL is preferred over supervised learning because the medical literature does not clearly define a ground-truth good treatment strategy.
2.2 Reinforcement Learning in Health: Unlike prior discretized state-and-action work, this study learns policies directly from continuous physiological states and focuses on in-hospital mortality.

3. Data and Preprocessing

The study uses a Sepsis-3 cohort from MIMIC-III and converts ICU observations into regularly aggregated physiological features. Interventions are represented through a discrete fluid–vasopressor action space.

3.1 Cohort: The cohort comes from MIMIC-III v1.4 and includes patients fulfilling Sepsis-3 criteria.
3.1 Cohort: Table 1 compares cohort statistics for subjects fulfilling the Sepsis-3 criteria.
3.2 State Representation: Physiological data include demographics, laboratory values, vital signs, and intake/output events aggregated into 4-hour windows.
3.2 State Representation: Preprocessing removes excessively missing variables and imputes remaining missing values with k-nearest neighbors, producing a 47 × 1 feature vector per timestep.
3.3 Action Representation: The action space covers IV fluid volume and maximum vasopressor dosage in each 4-hour window through a 5 × 5 intervention grid.

4. Methods

The method learns treatment policies from retrospective physician data using continuous states and discretized actions. Its architecture combines latent-state representations with value-based networks designed to address high dimensionality, overestimation, and sparse rewards.

4. Methods: Offline retrospective data limit exploration of the state space and make learning the truly optimal medication policy difficult.
4.1 Reinforcement Learning: The models use continuous state spaces with discretized action spaces to retain more physiological information while learning off-policy from physician trajectories.
4.2 Model Architecture: The baseline uses discretized states and actions with SARSA to estimate the physician-policy action-value function.
4.2 Model Architecture: Deep Q Networks minimize squared error between predicted action values and targets defined by reward plus discounted next-state value.
4.2 Model Architecture: Double-DQN reduces Q-value overestimation, while Dueling Q Networks separate state value from action advantage.
4.2.1 Model Architecture: The final Dueling DDQN combines Double-DQN and Dueling designs in a single architecture.
4.2.1 Model Architecture: After training, the policy selects the action with the highest estimated Q-value for each patient state.
4.3 Autoencoder Latent State Representation: Autoencoders, including sparse variants, transform high-dimensional physiological vectors into latent representations used by the Dueling DDQN.

5. Evaluation

The paper evaluates learned sepsis policies using multiple off-policy approaches because estimating outcomes under an unobserved policy is difficult. These methods relate expected returns to mortality, estimate policy value, and compare learned treatment choices with physician actions.

Evaluation challenges: Off-policy evaluation is difficult because learned-policy rollouts may lead to different eventual mortality than observed physician trajectories.Direct Q-value comparisons on off-policy data can provide incorrect performance estimates.
Return-to-mortality analysis: Q-values are binned and linked to observed patient outcomes to estimate mortality at different expected-return levels.Trajectories are labeled by whether patients died or survived, and average mortality is computed for each return bin.
Policy-value estimation: Doubly Robust Off-policy Value Evaluation provides an unbiased estimate of the learned policy’s value from sampled training trajectories.The estimated value is averaged across observed trajectories and combined with physician-policy returns.
Policy comparison: The evaluation combines policy-value estimates with the learned mortality-versus-return relationship to assess potential mortality reduction.This enables comparison of returns under the learned policy and the physician policy in mortality terms.
Clinical interpretation: Treatment choices from the learned policy are compared with physician actions to examine clinical differences associated with patient mortality.The comparison is intended to provide additional clinical understanding of the proposed policy.

6. Results

The learned policies achieved lower estimated mortality than the physician policy, while retaining clinically interpretable treatment patterns. Their learned returns were associated with mortality, and policy–clinician dosage agreement coincided with lower observed mortality.

Quantitative policy evaluation: The autoencode Q-N policy had the lowest estimated mortality and could reduce patient mortality by up to 4%.The comparison covered the physician, normal Q-N, and autoencode Q-N policies.
Quantitative policy evaluation: The model’s expected return showed an inverse relationship with mortality risk, although the relationship was not monotonic at low returns.The low-return irregularity was attributed to fewer training examples in that regime.
Quantitative policy evaluation: The sparse autoencoder representation showed a mortality gradient across its first two principal components, suggesting a physiologically informative latent state.The association was presented as a possible explanation for the model’s policy quality.
Qualitative policy examination: Both learned policies prescribed vasopressors sparingly, matching a key feature of the physician policy.The autoencode Q-N model reproduced this behavior more strongly than the normal Q-N model.
Quantifying optimality: Observed mortality was consistently low when optimal and clinician dosages coincided, while larger dosage differences produced less reliable results.For IV fluids, the normal Q-N policy showed a clear improvement over the physician policy at zero dosage difference; the autoencode Q-N policy was weaker.

7. Conclusion

The paper applies deep reinforcement learning to deduce sepsis treatment policies and reports that continuous state-space modeling could reduce hospital mortality while producing clinically interpretable policies.

The study applies deep reinforcement learning to deduce optimal medical treatments for patients with sepsis.
Sparse terminal rewards remain a limitation, motivating clinically informed or inverse reinforcement learning reward functions as possible improvements.The dataset contains trajectories recorded from many physicians, which may support inferring a more appropriate reward function.
The authors begin with a discretized model before exploring continuous state-space modeling for treatment-policy learning.The discretized model represents physiological data averaged over four-hour blocks and actions over IV fluids and vasopressors.
1.8–3.6% potential reduction in hospital mortality is reported for policies learned with continuous state-space modeling.The reported reduction is framed as an opportunity to identify better medication strategies for patients with sepsis.
The learned policies suggest that vasopressors may not be favored as a first response and that intermediate fluid dosages fit recent clinical findings.The authors also describe the policies as clinically interpretable and potentially usable for ICU clinical decision support.

8. APPENDICES

The appendix defines the sepsis cohort using Sepsis-3 criteria and analyzes early management with four-hour physiological time series derived from MIMIC-III data.

Sepsis is defined as suspected infection combined with organ dysfunction indicated by a SOFA score greater than or equal to 2.The cohort definition follows the temporal criteria for sepsis diagnosis and assumes a baseline SOFA score of zero.
MIMIC-III data cover up to 24 hours before through 48 hours after sepsis onset to capture early management and initial resuscitation.The outcome of interest is in-hospital mortality.
The extracted data contain 47 features processed into multidimensional time series with a four-hour time resolution.
The model’s physiological inputs include demographic and clinical measures such as age, SOFA, Glasgow Coma Scale, Shock Index, and SIRS.

Lab Values

The paper uses a broad set of laboratory, vital-sign, respiratory, fluid-output, and ventilation variables, then discretizes training states for a SARSA baseline.

Lab Values: Laboratory inputs include albumin, arterial pH, calcium, glucose, haemoglobin, electrolytes, coagulation measures, renal markers, lactate, bilirubin, and blood-cell counts.
Lab Values: Vital and respiratory inputs include blood pressures, PaCO2, PaO2, FiO2, PaO2/FiO2 ratio, respiratory rate, temperature, weight, heart rate, and SpO2.
Lab Values: Fluid-output variables cover four-hour fluid output, total fluid output, and mechanical ventilation.
Discretized Model: The discretized baseline applies k-means clustering to the training set, producing 1250 state clusters.Training and held-out test sets use an 80/20 split with proportionate patient trajectories.
Discretized Model: Offline SARSA updates the action-value function using sampled tuples of current and next state, action, and reward.The Bellman update is Q(s, a) ← Q(s, a) + α * [r + γQ(s′, a′) - Q(s, a)].
Discretized Model: The learned Q-function is tabulated over finite state-action pairs and evaluated as the action-value function associated with the physician policy.

8.5 Continuous Model Architecture and Implementation Details

The continuous-model implementation uses a two-layer neural architecture with stabilized activations, bounded Q-values, batch normalization, and transition-based training.

Architecture: The final network uses two hidden layers of 128 units, batch normalization, Leaky-ReLU activations, and separate advantage and value streams.The streams are combined and projected onto the action space.
Activation and Stability: The activation function is f(z) = max(z, 0.5z), allowing nonzero negative-side responses for Q-values that may be positive or negative.The authors motivate this choice by concerns about saturation and dead neurons with standard activations.
Activation and Stability: Rewards of ±15 at terminal timesteps and feature scaling were used to improve model stability.
Loss and Training: The Q-network loss penalizes outputs outside ±15, and target-network outputs are clipped to the same thresholds.
Loss and Training: Models use an 80/20 train/test split, batch normalization, TensorFlow v1.0, and Adam optimization.
Loss and Training: Training samples transitions <s, a, r, s′> and updates the main-network weights through feed-forward evaluation, loss computation, and backpropagation.

8.6 Autoencoder Implementation Details

The autoencoder incorporates a desired sparsity level into its training objective by combining reconstruction loss with a KL-divergence penalty. The penalty strength is controlled by a hyperparameter.

The autoencoder minimizes a sparsity-regularized objective combining reconstruction loss with a KL-divergence penalty.The objective is Lsparse(θ) = Lreconstruction(θ) + β Pn.
The desired sparsity ρ is specified before adjusting the autoencoder weights.
The sparsity penalty sums KL(ρ||ρj) across the network’s n hidden neurons.Here, ρj denotes the actual output of neuron j.
The hyperparameter β controls the strength of the sparsity term.