[论文解读] 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 Healthcare参考文献 5被引用 4

一句话总结

该论文提出了一种基于MIMIC的OpenAI Gym模拟器——败血症世界模型（Sepsis World Model），采用变分自编码器（VAE）和混合密度网络-RNN（MDN-RNN）对败血症治疗中的患者状态转移进行建模，降低噪声并捕捉不确定性。该模型使深度强化学习策略的训练更贴近临床决策过程，尤其在使用乳酸和SOFA评分等生理变化作为中间奖励信号时表现更优。

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.

研究动机与目标

• 开发一个真实、可微分的败血症患者轨迹模拟器，用于危重症护理中的强化学习。
• 通过使用变分自编码器（VAE）学习去噪潜在表示，降低电子健康记录（EHR）数据中的噪声。
• 通过使用混合密度网络-RNN（MDN-RNN）预测下一状态的概率分布，对患者状态转移中的不确定性进行建模。
• 评估模拟器在使用深度Q网络（DQN）和不同奖励设计时，训练临床合理治疗策略的能力。
• 通过引入基于乳酸和SOFA评分等中间临床指标的时间相关奖励，提升策略的多样性与真实性。

提出的方法

• 该模型采用具有三层全连接编码器和解码器的VAE，将46个归一化的EHR特征压缩至30维潜在空间，最小化均方重建误差。
• VAE生成的潜在状态表示输入MDN-RNN，后者以动作和隐藏状态为条件，将下一潜在状态的概率分布建模为高斯混合分布。
• 在每个时间步，模型从预测的混合分布中采样，以模拟随机的患者状态转移，通过温度控制调节不确定性。
• 模拟器集成了状态、终止和结果预测模型，整个系统以OpenAI Gym环境形式暴露，用于强化学习训练。
• 评估了三种奖励设计：(1) 稀疏的回合结束奖励，(2) 带惩罚的延迟奖励，(3) 基于乳酸和SOFA评分变化的密集中间奖励。
• 在模拟环境中训练DQN智能体，并从动作分布、回合长度和累积奖励等方面将其学习策略与医生行为进行比较。

实验结果

研究问题

• RQ1 基于VAE与MDN-RNN的世界模型能否有效对从噪声EHR数据中提取的败血症患者轨迹潜在动态进行去噪与表征？
• RQ2 通过使用随机MDN-RNN建模不确定性，是否能带来比确定性RNN基线更真实、更多样化的治疗策略学习？
• RQ3 不同奖励设计策略——尤其是密集中间奖励——如何影响DQN学习到的败血症治疗策略的真实性和多样性？
• RQ4 该模拟器学习到的策略在多大程度上复现了真实医生行为中的动作分布、回合长度和奖励轨迹？
• RQ5 该模拟器能否支持训练出临床合理、具有时间敏感性的治疗策略，避免对极端或单调干预的过拟合？

主要发现

• 与确定性RNN基线相比，VAE+MDN-RNN模型在处理噪声EHR数据时，产生了更稳定、更真实的轨迹。
• 使用稀疏的回合结束奖励导致策略崩溃，智能体集中于单一动作，产生不切实际的短回合。
• 基于乳酸和SOFA评分变化的中间奖励使动作分布更丰富、更真实，与医生行为高度相似。
• 采用中间奖励的模型实现了更均衡的回合长度和更优的动作多样性，表明其在时间策略学习方面表现更佳。
• 尽管奖励设计有所改进，终止模型仍存在过拟合，导致即使在更优奖励设计下，回合长度仍持续偏短。
• 该模拟器成功支持了DQN训练，其策略在真实性和表现上优于基线模型，尤其在奖励信号反映实时临床变化时更为显著。