Technical Summary: Agentifying Patient Dynamics within LLMs through Interacting with Clinical World Model

Problem Definition

Sepsis management in Intensive Care Units (ICUs) constitutes a high-stakes sequential decision-making problem. Effective treatment requires clinicians to titrate intravenous fluids and vasopressors over time to restore perfusion while avoiding downstream organ injury. This challenge is characterized by rapidly evolving patient physiology, where actions improving short-term hemodynamics (e.g., blood pressure) may worsen long-term outcomes (e.g., pulmonary edema or renal failure).

While Large Language Models (LLMs) possess broad clinical knowledge and can reason over medical guidelines, they lack inherent grounding in action-conditioned patient dynamics. An LLM may know general physiological principles but fail to estimate how a specific patient will respond to a specific dose over the next few hours. Existing methods often learn treatment policies directly from retrospective Electronic Health Record (EHR) trajectories, mapping states to actions. These approaches are vulnerable to historical behavior bias and struggle to reason about alternative treatment paths not well-represented in observational data.

Methodology

The authors introduce SepsisAgent, a world model-augmented LLM agent designed for sepsis treatment recommendation. The core innovation is the integration of a Clinical World Model with an LLM policy through a "propose–simulate–refine" workflow and a three-stage training curriculum.

1. The Clinical World Model

The Clinical World Model ($W_\theta$) is a learned, action-conditioned approximation of patient evolution. Given a current patient state $s_t$ and a candidate treatment action $a_t$, it predicts the distribution of the next physiological state $s_{t+1}$ and downstream outcome signals $o_t$.

• Architecture: Implemented using a 2-layer Gated Recurrent Unit (GRU) encoder to process patient trajectory history, followed by a state transition head (predicting Gaussian parameters for the next state) and an outcome head (predicting long-term clinical outcomes like mortality).
• Function: It acts as a predictive simulator, allowing the agent to compare counterfactual treatment options before committing to a prescription.

2. The SepsisAgent Workflow

At inference time, SepsisAgent operates via a propose–simulate–refine loop:

1. Propose: The LLM generates a set of candidate fluid–vasopressor interventions ($C_t$).
2. Simulate: For each candidate, the Clinical World Model estimates the physiological consequences ($\hat{y}_t$), including future states and outcome signals.
3. Refine: The LLM compares these action–response pairs against clinical priors and guidelines to select the final treatment action.

3. Three-Stage Training Curriculum

The authors argue that merely providing an LLM with access to a world model is insufficient; the agent must be explicitly trained to interpret simulated feedback. They propose a staged training pipeline:

• Stage I: Patient-Dynamics Prediction and Guideline-Aware Reasoning (Supervised Fine-Tuning):

  • The model is trained on tasks involving In-Hospital Mortality (IHM) prediction, Vasopressor Requirement (VR) prediction, and one-step treatment reasoning.
  • Data is synthesized using GPT-4.1 to create chain-of-thought rationales that integrate patient state analysis, dynamics prediction, and guideline constraints.
  • Goal: Teach the model to understand patient evolution and reason within sepsis guidelines.

• Stage II: Propose–Simulate–Refine Behavior Cloning:

  • The model learns to interact with the world model through multi-round reasoning traces.
  • Using expert transitions, the authors synthesize interaction trajectories where the agent proposes actions, queries the world model for simulated responses, and refines its decision before selecting the expert action.
  • Goal: Teach the agent how to use world-model feedback during decision-making (imitation learning).

• Stage III: World-Model-Based Agentic Reinforcement Learning (RL):

  • The agent interacts with the Clinical World Model as a virtual environment using Group Relative Policy Optimization (GRPO).
  • The agent samples multiple reasoning trajectories for each state and optimizes a composite reward function comprising terminal outcomes (survival/mortality), intermediate physiological stabilization (SOFA score, lactate dynamics), and guideline consistency penalties.
  • Goal: Optimize long-horizon treatment strategies rather than greedy short-term stabilization.

Key Contributions

1. A World Model-Augmented LLM Agent: The proposal of SepsisAgent, which collaborates with a Clinical World Model to compare counterfactual interventions. It outperforms traditional RL and vanilla LLM baselines on MIMIC-IV sepsis trajectories.
2. A Multi-Stage Training Paradigm: The demonstration that world-model access alone yields inconsistent performance. The paper establishes a curriculum (SFT $\to$ Behavior Cloning $\to$ RL) that teaches the agent to correctly interpret simulated dynamics and align them with clinical safety constraints.
3. Internalization of Patient Dynamics: Evidence that world-model-based agentic RL enhances the LLM's intrinsic ability to predict patient outcomes (IHM and VR) even when the simulator is removed during evaluation. The agent learns regularities in patient evolution rather than merely fitting reward signals.

Experimental Results

The system was evaluated on 725 held-out episodes from the MIMIC-IV dataset.

• Policy Value: SepsisAgent achieved the highest off-policy value scores (Doubly Robust, Weighted Importance Sampling, and Weighted Per-Decision Importance Sampling) among all baselines, including traditional RL methods (DDQN, AI Clinician) and various LLMs (GPT-4.1, Gemini-3, o3).
• Safety and Adherence: SepsisAgent achieved the highest sepsis guideline adherence (97.95%) and the lowest rates of unsafe actions (0.08% underdosing, 0.14% overdosing). In contrast, vanilla LLMs and world-model-augmented LLMs without the specific training pipeline showed higher rates of unsafe actions or lower policy values.
• Ablation Studies:
  • Stage I (SFT) and Stage II (Behavior Cloning) improved safety metrics and some policy values but did not yield stable gains across all off-policy estimators.
  • Stage III (RL) was identified as the primary driver of policy value improvement, suggesting that the world model serves as an effective environment for optimizing long-horizon strategies.
• Generalization: The model demonstrated strong cross-dataset performance when tested on MIMIC-III episodes, outperforming traditional RL baselines.

Significance and Claims

The paper claims to provide a proof of concept for using world-model-based agentic learning to support sequential clinical decision-making in high-acuity settings.

• Beyond Access: The primary finding is that access to a simulator is not enough; the agent must be trained to interpret the simulator's feedback. Generic LLMs tend to over-trust short-term simulated improvements (e.g., immediate blood pressure recovery) while ignoring long-term risks (e.g., fluid overload), a failure mode corrected by the proposed training pipeline.
• Internalization: The work suggests that repeated interaction with a clinical world model allows LLMs to internalize patient dynamics, improving their ability to predict mortality and vasopressor needs even without simulator access at inference time.
• Safety: The approach successfully balances the optimization of treatment value with strict adherence to clinical safety guidelines, avoiding the "unsafe shortcuts" often seen in purely reward-driven RL agents.

The authors explicitly state that the research is conducted for academic purposes using retrospective data and simulated environments, and the models are not intended for direct clinical deployment without human supervision.