Guardian-regularized Safe Offline Reinforcement Learning for Smart Weaning of Mechanical Circulatory Devices

TL;DR

The paper presents CORMPO, a Clinically-aware OOD-regularized Model-based Policy Optimization algorithm, for safe offline reinforcement learning in weaning mechanical circulatory support devices. It couples a Transformer-based probabilistic digital twin for evaluation with a KDE-driven density penalty to confine policy search to data-supported regions, and employs clinically informed reward shaping to guide safe, effective weaning. The authors provide theoretical guarantees and demonstrate substantial improvements over offline RL baselines on real and synthetic data, including higher physiological reward and improved clinical metrics. This work offers a principled, safe framework for data-driven medical decision making in high-stakes settings with uncertain dynamics and limited data, with potential applicability to other complex biomedical sequential decision problems.

Abstract

We study the sequential decision-making problem for automated weaning of mechanical circulatory support (MCS) devices in cardiogenic shock patients. MCS devices are percutaneous micro-axial flow pumps that provide left ventricular unloading and forward blood flow, but current weaning strategies vary significantly across care teams and lack data-driven approaches. Offline reinforcement learning (RL) has proven to be successful in sequential decision-making tasks, but our setting presents challenges for training and evaluating traditional offline RL methods: prohibition of online patient interaction, highly uncertain circulatory dynamics due to concurrent treatments, and limited data availability. We developed an end-to-end machine learning framework with two key contributions (1) Clinically-aware OOD-regularized Model-based Policy Optimization (CORMPO), a density-regularized offline RL algorithm for out-of-distribution suppression that also incorporates clinically-informed reward shaping and (2) a Transformer-based probabilistic digital twin that models MCS circulatory dynamics for policy evaluation with rich physiological and clinical metrics. We prove that \textsf{CORMPO} achieves theoretical performance guarantees under mild assumptions. CORMPO attains a higher reward than the offline RL baselines by 28% and higher scores in clinical metrics by 82.6% on real and synthetic datasets. Our approach offers a principled framework for safe offline policy learning in high-stakes medical applications where domain expertise and safety constraints are essential.

Figures (15)

• Figure 1: Illustration of data sparsity in low reward, low P-Level region (shaded in red).
• Figure 2: System diagram of the proposed framework. Digital twin module: We use a transformer encoder to learn a latent representation of the patient’s history, concatenate the representation with the P-Level input, and then decode the output using a fully connected neural network. Offline RL module with CORMPO: The replay buffer is created from data with clinical guided reward shaping. We learn a density-based guardian model on the data, whose OOD penalty terms are incorporated during policy training. The learned policies are evaluated in the digital twin-supported medical environment with rich medical metrics.
• Figure 3: An example of TDT prediction vs. baselines. The Transformer model is more accurate in reflecting response to P-level change and more expressive when capturing large changes in patient state, resulting in its higher accuracy as shown in Table \ref{['tab:world_model_results']}.
• Figure 4: Trained on the real-life dataset, we compare our CORMPO against baselines in 6-hour TDT rollouts. Our TDT predicts stable hemodynamics, away from the clinical threshold, for MAP when guided by CORMPO’s optimal policy. CORMPO's WS is $1.0$ which is the maximum possible score. While BC policy yields limited weaning, MBPO, MOPO, and SVR acts opposite to the weaning behavior because patient stability suggests gradual decrease in P-level. CORMPO results in the most successful weaning in this sample roll-out.
• Figure 5: Comparison of physiological reward, ACP, and WS distribution of expert, MOPO, and CORMPO policies. CORMPO suppresses actions with high ACP and results in higher WS compared to baselines, also reducing the high portion of negative rewards in expert policy and achieves higher rewards overall.

Theorems & Definitions (4)

• Theorem 1
• Theorem 2
• Lemma 1: Telescoping lemma - Lemma 4.1 in yu2020mopo
• Lemma 2: Discounted bound for signed shaping