[Paper Review] Sepsis World Model: A MIMIC-based OpenAI Gym "World Model" Simulator for Sepsis Treatment

Amirhossein Kiani, Chris Wang+1 more|arXiv (Cornell University)|Dec 15, 2019

Machine Learning in Healthcare5 references4 citations

TL;DR

This paper proposes a MIMIC-based OpenAI Gym simulator called the Sepsis World Model that uses a Variational Autoencoder (VAE) and Mixture Density Network-RNN (MDN-RNN) to model patient state transitions in sepsis treatment, reducing noise and capturing uncertainty. The model enables training of deep reinforcement learning policies that more closely mimic clinical decision-making, especially when using intermediate reward signals based on physiological changes like lactate and SOFA scores.

ABSTRACT

Sepsis is a life-threatening condition caused by the body's response to an infection. In order to treat patients with sepsis, physicians must control varying dosages of various antibiotics, fluids, and vasopressors based on a large number of variables in an emergency setting. In this project we employ a "world model" methodology to create a simulator that aims to predict the next state of a patient given a current state and treatment action. In doing so, we hope our simulator learns from a latent and less noisy representation of the EHR data. Using historical sepsis patient records from the MIMIC dataset, our method creates an OpenAI Gym simulator that leverages a Variational Auto-Encoder and a Mixture Density Network combined with a RNN (MDN-RNN) to model the trajectory of any sepsis patient in the hospital. To reduce the effects of noise, we sample from a generated distribution of next steps during simulation and have the option of introducing uncertainty into our simulator by controlling the "temperature" variable. It is worth noting that we do not have access to the ground truth for the best policy because we can only evaluate learned policies by real-world experimentation or expert feedback. Instead, we aim to study our simulator model's performance by evaluating the similarity between our environment's rollouts with the real EHR data and assessing its viability for learning a realistic policy for sepsis treatment using Deep Q-Learning.

Motivation & Objective

• To develop a realistic, differentiable simulator of sepsis patient trajectories for reinforcement learning in critical care.
• To reduce noise in electronic health record (EHR) data by learning a denoised latent representation using a Variational Autoencoder (VAE).
• To model uncertainty in patient state transitions by using a Mixture Density Network-RNN (MDN-RNN) that predicts distributions over next states.
• To evaluate the simulator’s ability to train clinically plausible treatment policies using Deep Q-Networks (DQN) with different reward formulations.
• To improve policy diversity and realism by incorporating time-dependent rewards based on intermediate clinical markers like lactate and SOFA scores.

Proposed method

• The model uses a VAE with three dense encoder and decoder layers to compress 46 normalized EHR features into a 30-dimensional latent space, minimizing mean-squared reconstruction error.
• The latent state representations from the VAE are fed into an MDN-RNN that models the probability distribution of the next latent state as a mixture of Gaussians, conditioned on action and hidden state.
• At each time step, the model samples from the predicted mixture distribution to simulate stochastic patient state transitions, with temperature control to adjust uncertainty.
• The simulator integrates state, termination, and outcome prediction models, with the entire system exposed as an OpenAI Gym environment for RL training.
• Three reward formulations are evaluated: (1) sparse end-of-episode reward, (2) delayed reward with penalty, and (3) dense intermediate rewards based on changes in lactate and SOFA scores.
• A DQN agent is trained in the simulated environment, and its learned policies are compared to physician behavior in terms of action distribution, episode length, and cumulative reward.

Experimental results

Research questions

• RQ1 Can a VAE- and MDN-RNN-based world model effectively denoise and represent the latent dynamics of sepsis patient trajectories from noisy EHR data?
• RQ2 Does modeling uncertainty through a stochastic MDN-RNN lead to more realistic and diverse treatment policy learning compared to deterministic RNN baselines?
• RQ3 How do different reward shaping strategies—especially dense intermediate rewards—impact the realism and diversity of DQN-learned sepsis treatment policies?
• RQ4 To what extent does the simulator’s learned policy replicate the action distribution, episode length, and reward trajectory observed in real physician behavior?
• RQ5 Can the simulator support the training of clinically plausible, time-sensitive treatment policies that avoid overfitting to extreme or monotonous interventions?

Key findings

• The VAE+MDN-RNN model produced more stable and realistic state trajectories than deterministic RNN baselines, particularly in handling noisy EHR data.
• Using a sparse end-of-episode reward led to policy collapse, with agents clustering on a single action and producing unrealistically short episodes.
• Intermediate rewards based on changes in lactate and SOFA scores resulted in a more diverse and realistic distribution of actions, closely resembling physician behavior.
• The model with intermediate rewards achieved a more balanced episode length and better action diversity, indicating improved temporal policy learning.
• Despite improvements, the termination model remained overfit, resulting in consistently short episode lengths even with better reward shaping.
• The simulator successfully enabled DQN training that outperformed baseline models in policy realism, especially when reward signals reflected real-time clinical changes.