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

arXiv: https://arxiv.org/abs/2506.15019 · HTML: https://mlsanity.com/paper/2506.15019/

## 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.

## Placed in

- Offline RL and imitation learning from data / [Safe offline RL](https://mlsanity.com/offlinerl/safe-offline-rl/)

## Builds on (most cited first, 15 of 15)

- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Recurrent Neural Networks for Multivariate Time Series with Missing Values](https://mlsanity.com/paper/1606.01865/) (2016)
- [GAIN: Missing Data Imputation using Generative Adversarial Nets](https://mlsanity.com/paper/1806.02920/) (2018)
- [Neural SDEs as Infinite-Dimensional GANs](https://mlsanity.com/paper/2102.03657/) (2021)
- [State Representation Learning for Control: An Overview](https://mlsanity.com/paper/1802.04181/) (2018)
- [Stiff Neural Ordinary Differential Equations](https://mlsanity.com/paper/2103.15341/) (2021)
- [Neural Controlled Differential Equations for Online Prediction Tasks](https://mlsanity.com/paper/2106.11028/) (2021)
- [An Empirical Study of Representation Learning for Reinforcement Learning in Healthcare](https://mlsanity.com/paper/2011.11235/) (2020)
- [Opening the Blackbox: Accelerating Neural Differential Equations by Regularizing Internal Solver Heuristics](https://mlsanity.com/paper/2105.03918/) (2021)
- [Training Stiff Neural Ordinary Differential Equations with Implicit Single-Step Methods](https://mlsanity.com/paper/2410.05592/) (2024)
- [Enhancing Continuous Time Series Modelling with a Latent ODE-LSTM Approach](https://mlsanity.com/paper/2307.05126/) (2023)
- [Proximal Implicit ODE Solvers for Accelerating Learning Neural ODEs](https://mlsanity.com/paper/2204.08621/) (2022)
- [Modelling EHR timeseries by restricting feature interaction](https://mlsanity.com/paper/1911.06410/) (2019)
- [Anamnesic Neural Differential Equations with Orthogonal Polynomial Projections](https://mlsanity.com/paper/2303.01841/) (2023)
- [Geometric instability of out of distribution data across autoencoder architecture](https://mlsanity.com/paper/2201.11902/) (2022)

## 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.