Stable CDE Autoencoders with Acuity Regularization for Offline Reinforcement Learning in Sepsis Treatment

TL;DR

This paper demonstrates that stable Controlled Differential Equation (CDE) autoencoders with acuity regularization improve the quality of state representations for offline reinforcement learning in sepsis treatment, leading to better policy performance and meaningful clinical insights.

Contribution

It introduces a training stabilization approach for CDE autoencoders and enforces acuity-aware regularization, advancing representation learning for clinical RL tasks.

Findings

Stable CDE autoencoders produce representations correlated with clinical scores.

Stable training leads to RL policies with high WIS return (> 0.9).

Unstable CDEs result in poor representations and policy failure.

Abstract

Effective reinforcement learning (RL) for sepsis treatment depends on learning stable, clinically meaningful state representations from irregular ICU time series. While previous works have explored representation learning for this task, the critical challenge of training instability in sequential representations and its detrimental impact on policy performance has been overlooked. This work demonstrates that Controlled Differential Equations (CDE) state representation can achieve strong RL policies when two key factors are met: (1) ensuring training stability through early stopping or stabilization methods, and (2) enforcing acuity-aware representations by correlation regularization with clinical scores (SOFA, SAPS-II, OASIS). Experiments on the MIMIC-III sepsis cohort reveal that stable CDE autoencoder produces representations strongly correlated with acuity scores and enables RL policies with superior performance (WIS return $> 0.9$). In contrast, unstable CDE representation leads to degraded representations and policy failure (WIS return $\sim$ 0). Visualizations of the latent space show that stable CDEs not only separate survivor and non-survivor trajectories but also reveal clear acuity score gradients, whereas unstable training fails to capture either pattern. These findings highlight practical guidelines for using CDEs to encode irregular medical time series in clinical RL, emphasizing the need for training stability in sequential representation learning.