Stable CDE Autoencoders with Acuity Regularization for Offline Reinforcement Learning in Sepsis Treatment
On the map: Offline RL and imitation learning from data → Safe offline RL
Abstract (arXiv)
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 ~ 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.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Safe offline RL, one of 13 directions in Offline RL and imitation learning from data. Builds on 15 of its 42 references that are in the map. First posted on arXiv 17 Jun 2025.
Builds on
15 of 42 references.
References · 15
- 2018Neural Ordinary Differential Equationscited by 3,779
- 2016Recurrent Neural Networks for Multivariate Time Series with Missing Valuescited by 554
- 2018GAIN: Missing Data Imputation using Generative Adversarial Netscited by 352
- 2021Neural SDEs as Infinite-Dimensional GANscited by 161
- 2018State Representation Learning for Control: An Overviewcited by 127
- 2021Stiff Neural Ordinary Differential Equationscited by 78
- 2021Neural Controlled Differential Equations for Online Prediction Taskscited by 46
- 2020An Empirical Study of Representation Learning for Reinforcement Learning in Healthcarecited by 34
- 2021Opening the Blackbox: Accelerating Neural Differential Equations by Regularizing Internal Solver Heuristicscited by 26
- 2024Training Stiff Neural Ordinary Differential Equations with Implicit Single-Step Methodscited by 11
- 2023Enhancing Continuous Time Series Modelling with a Latent ODE-LSTM Approachcited by 9
- 2022Proximal Implicit ODE Solvers for Accelerating Learning Neural ODEscited by 6
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026Attention-Based Offline Reinforcement Learning and Clustering for Interpretable Sepsis Treatment
- 2024How Consistent are Clinicians? Evaluating the Predictability of Sepsis Disease Progression with Dynamics Models
- 2025Exploring a Graph-based Approach to Offline Reinforcement Learning for Sepsis Treatment
- 2025Exploring Time-Step Size in Reinforcement Learning for Sepsis Treatment
- 2021Offline reinforcement learning with uncertainty for treatment strategies in sepsis
- 2025MORE-CLEAR: Multimodal Offline Reinforcement learning for Clinical notes Leveraged Enhanced State Representation
- 2024ICU-Sepsis: A Benchmark MDP Built from Real Medical Data
- 2026Formally Verifying and Explaining Sepsis Treatment Policies with COOL-MC