[論文レビュー] 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

ひとこと要約

本論文は、変分オートエンコーダー（VAE）と混合密度ネットワーク-RNN（MDN-RNN）を用いて、sepsis治療における患者状態遷移をノイズ低減および不確実性の捉え込みとともにモデル化する、MIMICに基づくOpenAI Gymシミュレータ「Sepsis World Model」を提案する。このモデルにより、深層強化学習（DRL）ポリシーの学習が促進され、乳酸や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.

研究の動機と目的

• 重症ケアにおける強化学習のための、現実的で微分可能なsepsis患者の遷移シミュレータの開発。
• 変分オートエンコーダー（VAE）を用いて、電子的健康記録（EHR）データからのノイズ低減された潜在的表現を学習すること。
• 混合密度ネットワーク-RNN（MDN-RNN）を用いて、次状態の確率分布を予測することで、患者状態遷移の不確実性をモデル化すること。
• 強化学習（DQN）を用いた治療ポリシーの学習において、シミュレータの有効性を、異なる報酬設計のもとで評価すること。
• 乳酸やSOFAスコアといった中間臨床マーカーに基づく時間依存報酬を組み込むことで、ポリシーの多様性と現実性を向上させること。

提案手法

• 本モデルは、46個の正規化済みEHR特徴量を30次元の潜在空間に圧縮するため、3層の全結合エンコーダーおよびデコーダーを持つVAEを用い、平均二乗誤差を最小化する。
• VAEで得られた潜在状態表現を、行動と隠れ状態に条件づけられた混合ガウス分布として次潜在状態の確率分布をモデル化するMDN-RNNに供給する。
• 各時刻において、予測された混合分布からサンプリングすることで、確率的患者状態遷移をシミュレートし、温度パラメータを調整して不確実性を制御する。
• シミュレータは状態、終了、および予後予測モデルを統合し、システム全体を強化学習用に利用可能なOpenAI Gym環境として公開する。
• 3種類の報酬設計を評価：(1) 疎なエピソード終了報酬、(2) ペナルティ付きの遅延報酬、(3) 乳酸およびSOFAスコアの変化に基づく密度の高い中間報酬。
• DQNエージェントをシミュレータ環境で学習させ、その学習済みポリシーを医師の行動と比較して、行動分布、エピソード長、累積報酬の観点から評価する。

実験結果

リサーチクエスチョン

• RQ1 VAEおよびMDN-RNNを用いた世界モデルは、ノイズの多いEHRデータからsepsis患者の遷移の潜在的ダイナミクスを的確にノイズ低減・表現できるか？
• RQ2 確率的MDN-RNNによる不確実性のモデル化は、決定的RNNベースラインと比較して、より現実的かつ多様な治療ポリシー学習を可能にするか？
• RQ3 特に密度の高い中間報酬を用いた報酬設計戦略は、DQNが学習するsepsis治療ポリシーの現実性と多様性にどのように影響を与えるか？
• RQ4 シミュレータが学習したポリシーは、実臨床における医師行動の行動分布、エピソード長、報酬トレースをどの程度再現できるか？
• RQ5 このシミュレータは、極端または単調な介入に過適合するのを避ける、臨床的に妥当で時間依存性のある治療ポリシーの学習を支援できるか？

主な発見

• VAE+MDN-RNNモデルは、ノイズの多いEHRデータを扱う際、決定的RNNベースラインと比較してより安定的かつ現実的な状態遷移を生成した。
• 疎なエピソード終了報酬を用いた場合、エージェントが単一の行動に集中し、不自然に短いエピソードが生成されるポリシーの崩壊が発生した。
• 乳酸およびSOFAスコアの変化に基づく中間報酬を用いることで、医師の行動に近い多様で現実的な行動分布が得られた。
• 中間報酬を用いたモデルは、よりバランスの取れたエピソード長と高い行動多様性を達成し、時間的ポリシー学習の向上が示された。
• 改善が見られたものの、終了モデルは依然として過学習しており、報酬設計の改善にもかかわらず、一貫して短いエピソード長が継続した。
• シミュレータは、報酬信号がリアルタイムの臨床的変化を反映する場合に特に、ベースラインモデルを上回るポリシーの現実性を実現するDQN学習を成功裏に支援した。