Offline Reinforcement Learning for Out-of-Distribution ICU Sepsis Decision Support

Technical Summary: Offline Reinforcement Learning for Out-of-Distribution ICU Sepsis Decision Support

Problem Statement
Sepsis management in Intensive Care Units (ICU) is a critical sequential decision-making problem where clinicians adapt interventions based on evolving patient physiology. While Reinforcement Learning (RL) offers a framework for learning treatment policies from logged data, prospective exploration is ethically constrained. Consequently, policies must be evaluated offline. A significant challenge arises under distribution shift: policies trained on standard datasets may behave unpredictably for patients with severe health conditions, who are often underrepresented in training data but are clinically the most consequential. This paper investigates whether standard offline RL methods retain stable, action-sensitive decision-support signals when evaluated on increasingly severe out-of-distribution (OOD) patient cohorts.

Methodology
The study utilizes the MIMIC-III benchmark dataset to construct a rigorous evaluation framework centered on severity-induced distribution shift.

• Data and OOD Splits: The authors construct three severity-enriched OOD test mixtures (25%, 50%, and 75% severe-OOD ratios) by increasing the proportion of held-out severe trajectories. Trajectory severity is defined by the mean shock_severity_proxy value. Policies are trained on less-severe cohorts and evaluated on these progressively harder test sets.
• Offline RL Methods: Seven standard offline RL algorithms are compared, all trained strictly on logged trajectories without environment interaction:
  • Fitted Q-Iteration (FQI)
  • Conservative Q-Learning (CQL)
  • FQI with a CQL-style penalty (FQI+CQL)
  • Batch-Constrained Q-Learning (BCQ)
  • Implicit Q-Learning (IQL)
  • Critic-Regularized Regression (CRR)
  • Advantage-Weighted Actor-Critic (AWAC)
• Learned-Dynamics Off-Policy Evaluation (OPE): A critical component of the framework is a shared, separately learned dynamics ensemble (a feed-forward multilayer perceptron) used exclusively for evaluation. This ensemble predicts next-state changes, rewards, termination probabilities, and survival probabilities based on state-action pairs.
  • Rollout Procedure: Policies generate trajectories by interacting with this learned dynamics model. The rollout stops when the cumulative termination probability reaches a threshold ($p_{done} \in \{0.3, 0.5, 0.7, 0.9\}$) or a maximum horizon ($H_{max}=100$) is reached.
  • Metrics:
    1. Model-Predicted Terminal Survival ($\hat{J}_{term}$): The average survival probability assigned by a calibrated survival evaluator to the final state of the model-generated rollout. This is action-sensitive and model-dependent.
    2. Observed Clinical Survival ($J_{obs}$): The actual survival rate from the logged historical data, which is fixed and policy-independent.
    3. Episode-Level Physiological Stabilization Score (EPSS): A heuristic metric scoring whether physiological variables (e.g., lactate, MAP, creatinine) improve, stabilize, or worsen over time. It is computed for both model-generated rollouts and matched logged clinical trajectories.

Key Contributions

1. Severity-OOD Stress Testing: The authors evaluate standard offline RL policies by progressively increasing the fraction of severe held-out ICU sepsis trajectories in the test cohort, moving beyond standard in-distribution evaluation.
2. Distinct Evaluation of Quantities: The study explicitly separates observed clinical survival ($J_{obs}$) from model-predicted terminal survival ($\hat{J}_{term}$). The authors emphasize that the numerical contrast between these two is a model-dependent OPE diagnostic, not evidence of causal treatment benefit.
3. Physiological Stabilization Analysis: A secondary analysis using EPSS provides a complementary diagnostic to terminal survival, assessing whether offline-policy trajectories exhibit favorable movements in physiological variables compared to historical data.

Results

• Survival Trends: As the severe-OOD ratio increased from 25% to 75%, the observed clinical survival rate ($J_{obs}$) declined monotonically from 67.3% to 48.8%. In contrast, the model-predicted terminal survival returns ($\hat{J}_{term}$) for offline RL methods remained relatively stable, ranging between 0.84 and 0.87 across all severity levels and methods.
• Performance Gap: The gap between the best model-predicted return and observed clinical survival widened as cohort severity increased (from +0.194 at 25% to +0.358 at 75%). This gap was robust across different rollout stopping thresholds ($p_{done}$).
• Method Comparison: No single algorithm dominated all settings. Fitted Q-Iteration (FQI) was most consistently strong, while BCQ performed best in the mildest severe-OOD setting at longer rollouts. The spread in performance across methods was small, suggesting a family of stable model-based behaviors rather than a single superior algorithm.
• Physiological Stabilization: Offline-policy model rollouts achieved higher EPSS values than matched logged clinical trajectories. For the CQL policy, the EPSS gap increased with cohort severity (e.g., +0.32 at 75% severe-OOD at horizon 10). The primary drivers of this improvement were reductions in shock severity, creatinine, lactate, and improvements in Mean Arterial Pressure (MAP).
• Rollout Length: Mean rollout length increased with cohort severity, reflecting the model's prediction of longer survival or delayed termination. The authors note that longer rollouts are more susceptible to accumulated model error.

Significance and Claims
The paper claims that under a shared learned-dynamics OPE protocol, offline RL policies can retain stable, action-sensitive decision-support signals even when evaluated on increasingly severe, out-of-distribution patient cohorts. The stability of $\hat{J}_{term}$ and the favorable EPSS signals suggest that these methods do not collapse under severity shift.

However, the authors maintain a modest stance regarding clinical implications. They explicitly state that:

• The observed gap between model-predicted and observed survival is a model-based contrast, not a causal treatment effect.
• The results are not evidence that deploying these policies would improve real-world patient survival or physiology.
• The metrics inherit potential biases and optimism from the learned dynamics ensemble, particularly in severe OOD states.
• Future work is required to calibrate the dynamics evaluator, validate findings on independent clinical cohorts, and perform causal evaluation before any clinical claims or deployment can be considered.

In summary, the paper positions offline RL as a promising framework for generating stable decision-support signals under distribution shift, while rigorously delineating the boundary between model-based diagnostics and clinical reality.