Organizations: Shanghai Artificial Intelligence Laboratory · University of the Chinese Academy of Sciences · Fudan University · Shanghai Jiao Tong University · Zhejiang University · University College Dublin · Shanghai Innovation Institute · Ant Group · University of Science and Technology of China · Co-corresponding authors
Pneumonia changes across respiration, inflammation, and organ function, while the measurements that describe it arrive at different times and are often incomplete. We introduce PULMOWORLD, a world model for longitudinal pneumonia dynamics that connects asynchronous state estimation, recursive simulation, and evidence-constrained planning. History-aware filtering separates state propagation from observation-driven correction, so a measurement updates the shared belief only when it is observed. That belief supports recursive forecasts of clinical values and modality representations, and summaries of future prognosis. Cap-MBRL searches imagined continuations inside an authorized action set and replaces a frozen fallback only when a candidate meets the risk budget and a predicted-cost margin. On the pneumonia cohort MASS, recursive forecasts are competitive with task-trained sequence models, and simulated future states retain prognostic information. Trained separately on public MIMIC-IV-CXR chest radiographs and clinical trajectories, this model improves prognostic discrimination over latest-value fusion in development and is competitive on long-horizon error on the locked pneumonia panel. In five known-mechanism MASS-Syn environments, Cap-MBRL lowers selection regret relative to direct fallback execution.
Figures & tables
Figure 1: A shared patient world for longitudinal simulation and evidence-constrained planning. (a) Asynchronous multimodal observations inform (b) an evolving patient belief and (c) simulation and planning interfaces. The lower strip summarizes cohort forecasting and MASS-Syn selection. Treatment branches show imagined continuations scored in MASS-Syn.
Figure 2: Pneumonia state learning and evidence-constrained planning. (a) Irregular clinical, imaging, and BALF observations update one patient belief, which a shared transition then propagates. (b) In MASS-Syn, a capability screen authorizes actions before frozen-world rollouts; predicted return, uncertainty, and risk are used when ranking candidates against the fallback. Below, reconstruction and rollout losses fit the world model.
State estimation
AUROC ↑
AUPRC ↑
ACC ↑
Bal. ACC ↑
Brier ↓
Without auxiliary reconstruction
Asynchronous filtering
0.7565
0.3737
0.8481
0.6092
0.1064
Latest-value fusion
0.7352
0.3538
0.8454
0.6061
0.1079
With auxiliary reconstruction
Asynchronous filtering
0.7748
0.4326
0.8673
0.6311
0.0965
Latest-value fusion
0.7367
0.3548
0.8427
0.6082
0.1075
Table 1: Asynchronous history filtering improves prognostic prediction. MIMIC development, three-seed means on 3,758 patients (524 deaths). The asynchronous filter retains ordered history; fusion uses latest values. Bold marks the better configuration within each reconstruction setting.
Figure 3: Recursive forecasts of clinical measurements and multimodal representations on MASS. (a,b) Clinical19 locked-test mean H1–H4 NMAE and change error ( n=40 ; patient-bootstrap 95% CIs in b). (c,d) Clinical19 development CT/BALF representation MSE, not the 40-patient locked forecast (three-seed means; non-zero axis origins). Blue denotes PulmoWorld ; lower is better. Complete scores are in Tables 10 and 21 .
Figure 4: Observed imaging evolution and recursive clinical forecasts. (a) Observed MIMIC validation CXRs, not generated images or future rollout inputs. H1–H4 denote six-hour windows, not acquisition times. (b) Representation RMSE for the same sequence. (c–e) Separate, outcome-selected locked-test clinical rollouts. Selection and image timing are in Sections D.6 and 12 .
MIMIC locked test: in-hospital death
MASS frozen validation: 28-day death
Method
ACC ↑
AUROC ↑
AUPRC ↑
Brier ↓
ACC ↑
AUROC ↑
AUPRC ↑
Brier ↓
Controls
Current state
0.8400
0.7106
0.3022
0.0899
0.6162
0.6856
0.4263
0.2051
Persistence
0.5500
0.6244
0.2139
0.1299
0.5960
0.6841
0.4254
0.2100
Train median
0.4200
0.5635
0.1947
0.1220
0.6263
0.6408
0.3776
0.2070
General-purpose Direct VLMs
Table 3: Cross-dataset prognosis metrics from predicted future states. One frozen probe per cohort: MIMIC pneumonia locked test ( n=300 ) and MASS frozen validation ( n=99 ). MIMIC uses H4; MASS uses available endpoints through H4. ACC uses frozen calibration-selected thresholds. Bold and underline mark the best and second-best task-trained result in each column.
Policy control
Unauth. (%)
Sel. regret
Unrestricted WM-RL
32.19
–
Uncertainty-aware
32.26
–
Support-aware
33.71
–
Path-consistency weighting
31.58
–
Same authorized action set
Authorized + soft penalties
0.00
0.3636
Table 4: Policy controls for evidence-constrained planning. Five-environment MASS-Syn means (lower is better). Unauthorized-action rate is the rate of executed actions outside the permitted set. Selection regret is the gap to the best action inside that set.
Figure 5: Conditional trajectories and authorized-set selection. (a–c) Conditional MIMIC validation forecasts under four six-hour care templates; zero fluid and vasopressor inputs denote “No intervention” ( Section C.9 ). (d,e) MASS-Syn selection regret: five-environment mean ± population SD and paired environment results.
Appendix figures & tables24 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: From cohort records to model-specific evaluation populations. (a,b) Private MASS and public MIMIC differ in clinical scope, observation clock, and modality availability. Patient splits precede model fitting; numerical transforms and projections are train-fit. (c) Forecasting and prognosis use distinct eligibility rules and must retain their own denominators. MASS prognosis uses the available endpoint through H4, not uniformly H4. Counts describe frozen data extracts, not recruitment estimates.
Quantity
Value
Development patients
683
Patient split (train/validation/calibration)
485 / 106 / 92
Events / observations / factual transitions
3,252 / 13,039 / 2,569
Transition split (train/validation/calibration)
1,806 / 428 / 335
Patients with clinical / CT / BALF
683 / 197 / 426
Clinical / CT / BALF observations
11,583 / 567 / 889
Appendix
Table 5: Clinical19 development data view. The 83-patient locked test is physically excluded; the 40-patient Clinical19 recursive forecasting evaluation is reported separately. Development splits are patient-disjoint.
Figure 7: Cohort composition and observation availability. (a) Patient counts and within-cohort split proportions; MASS development excludes the 83 locked-test patients. The Clinical19 recursive forecasting panel contains 40 H4-eligible patients and 40 evaluated windows. (b) MASS patient-level modality coverage, with observation counts shown separately. Its 567 CT observations occur on 3,252 irregular events (17.4%). (c) MIMIC linked-radiograph counts divided by all 6-hour events in each split. Patient coverage in (b) and event availability in (c) have different denominators. (d) Exact MIMIC patient, stay, event, and H4-window counts. An H4 window contains one source and four ordered future events; overlapping windows are not independent patients.
Flow step
Remaining
Removed
Evidence status
Recorded clinical source cohort
766
–
Source extract
Authoritative patient split assigned
766
0
Exact set match
Locked-test patients physically excluded
683
83
Test closed
Test-excluded development cohort
683
0
Exact non-test set
Patients with at least one event
683
0
Complete coverage
Appendix
Table 6: Patient flow into the development population. Record-level quarantine rows are not unique-patient exclusions. Raw screening and institutional recruitment denominators remain unavailable.
Figure 8: Architecture of the PulmoWorld observation filter. Modality-specific encoders map masked observations into a Gaussian patient belief; imaging features additionally pass through attention and recurrent memory. Available evidence corrects the state, elapsed-time dynamics propagate it, and a mixture head predicts observations. The dashed loop propagates the predicted belief when no new observation arrives. MASS has three observation modalities and seven internal state blocks ( H=256 , Z=64 ); MIMIC has two modalities and two blocks ( H=64 , Z=32 ). These internal partitions are not additional modalities. The locked MIMIC forecast uses this same history-aware belief.
Module
MASS Clinical19
MIMIC
Numeric encoders
537,600
18,304
History/temporal image encoder
939,008
62,720
State prior and initial projections
232,064
8,448
Dynamics
2,649,827
62,596
Prediction/innovation/correction
921,984
24,960
Observation and auxiliary heads
57,689
15,433
Appendix
Table 7: Checkpoint parameter accounting. Counts are obtained from the stored tensor shapes; frozen image backbones, PCA, and target-network copies are excluded. The history/temporal-image row is the imaging memory above: multi-head attention plus a GRU cell. The MIMIC parent and action adapter have identical total size but different optimization scopes. Counts include allocated modules in the reported checkpoints.
Base prediction: AdamW, batch 16, learning rate 5×10−5 , at most 20 epochs. Constrained refinement: 10−5 , at most 12 epochs, recursive/direct horizon weights 0.05/0.25/1.5/2.2.
Calibration selection; seeds 20260915–20260917; 5,338,172 learned parameters per checkpoint.
MIMIC
13 clinical variables; CXR PCA64; two compartments (clinical/radiograph); hidden 64; compartment 32; 8 Fourier time features; four recursive horizons.
Same parent plus two transformed factual action channels in a 6-D reserved interface; action residual and gate enabled.
AdamW on action adapter only; batch 256; learning rate 3×10−4 ; weight decay 10−4 ; gradient clipping 1.0; at most 8 epochs; patience 2 after epoch 3.
Calibration clinical macro RMSE; selected epoch 7; 20,352 action-related parameters within the 192,461-parameter model.
Appendix
Table 8: Model architecture and training hyperparameters. Each cohort has one observation setting. Parameter counts include learned world-model modules but exclude frozen image encoders and train-fitted PCA. Constrained refinement continues from the Clinical19 checkpoint. The MIMIC action adapter initializes from the history-aware parent and updates only action residual and gate parameters.
Fluoroquinolone, tetracycline, aminoglycoside, polymyxin, other antimicrobial
Antifungal
Antifungal
Antiviral
Antiviral
Immunomodulation
Systemic corticosteroid, other immunomodulator
Appendix
Table 9: Frozen MASS action ontology. Twenty privacy-safe macro classes are collapsed by a within-group maximum into ten mechanism blocks.
MIMIC locked test
MASS Task 1 validation
Requests / masked values
271 / 718
105 / 210
Masked NMAE
0.783 cell; 0.834 macro
0.667
Std. RMSE / Gaussian NLL
1.143
3.228
EHR-S@10/15/25
34.6 / 42.8 / 58.8
–
95% coverage / availability Brier
–
90.0% / 0.278
Appendix
Table 11: Masked source reconstruction under two evaluation protocols. The MIMIC column is the frozen locked test: 271 requests, 718 masked cells, the same train-only IQR and scorer as the Direct VLMs in Table 12 . The MASS column is the Clinical19 development validation: 105 examples and 210 masked values. It is not the 40-patient locked forecast, so the columns are not a paired comparison. Lower NMAE, RMSE, NLL, and Brier score are better. Coverage closer to 95% is better.
Time input
Training / input arm
Clinical RMSE ↓
Elapsed
Rollout only
0.7688±0.0022
Elapsed
Reconstruction + rollout
0.7646±0.0021
Paired component contrast
Gain [95% CI]
Reconstruction curriculum
0.00371[0.00196,0.00535]
Appendix
Table 13: Matched component ablations on MIMIC development. Top: clinical RMSE, mean ± sample SD across three training seeds. Bottom: paired patient-level RMSE reduction; positive favors the named component. Each patient’s window/horizon MSE is averaged before taking the square root, then averaged across seeds. These contrasts therefore have a different aggregation from the top-panel scores. Intervals use 5,000 patient-bootstrap resamples over 4,036 patients, conditional on the three seeds.
Source perturbation
Elapsed-time rollout
With reconstruction
None
0.7688±0.0022
0.7646±0.0021
25% source dropout
0.8264±0.0010
0.8239±0.0015
50% source dropout
0.8891±0.0007
0.8884±0.0036
75% source dropout
0.9504±0.0015
0.9500±0.0055
Appendix
Table 14: Inference with missing source evidence. Clinical RMSE, mean ± sample SD over the same three seeds. Source dropout masks each observed clinical value and drops each whole CXR channel at the stated probability, without retraining. Perturbations are paired across arms. These rows mask observed clinical values and drop whole CXR channels together. They are development robustness diagnostics.
Figure 9: A complete single-window ranking of all 16 MIMIC action templates. (a) The frozen action-conditioned model ranks every train-derived fluid–vasopressor template by imagined H4 mortality probability for one representative validation window. The case is selected deterministically as the median score-range window among those whose imagined optimum is not the no-op fallback template. (b) The same ranks in the 4×4 action grid. F1–V0 has the lowest model score, and the displayed branch retains F0–V0.
Method
Selection regret
Unauth. rate
Diagnostic controls without hard authorization
Unrestricted WM-RL
–
0.3219
Uncertainty-aware
–
0.3226
Support-aware
–
0.3371
Path-consistency weighting
–
0.3158
Policy comparison within the same authorized action set
Appendix
Table 17: Authorization-gated world-model RL on known-mechanism MASS-Syn. Means over five independent simulator-parameter and split seeds in the confirmatory experiment; the aggregate met the prespecified criteria. Lower is better. Selection regret is policy error within the authorized set. Rows with nonzero unauthorized action rates can choose actions outside the permitted set, so their selection regret is not comparable. The authorized soft-penalty policy combines the mask with uncertainty, support, and path-consistency terms, without Cap-MBRL’s budgeted candidate selection. Among rows with zero unauthorized action rate, bold and underline mark the best and second-best selection regret. Rates are fractions, not percentages.
Figure 10: What an action is and how it changes a frozen simulation. (a) The MIMIC action pool is a 4×4 train-supported grid over fluid and vasopressor levels; annotations give train-transition counts. The four corners define the diagnostic candidates replayed in panels (b–d). Repeating these four train-derived templates over H1–H4 separates predicted mean arterial pressure, lactate, and urine-output trajectories. Lines are patient-level means over 2,048 validation windows and bands are patient-cluster 95% CIs. This action sensitivity is predictive input to the gate.
Figure 11: Integrated pre-encoder case records. (a) A MASS validation CT pair is shown beside its matched Clinical19 source and observed next-state record. Numerical references are at +24 h, whereas the observed CT follow-up is at +72 h. Missing target cells remain missing; neither images nor target values are generated. The fixed lung-window slices need not depict identical anatomy and do not support pixelwise comparison. The full Clinical19 field list is in Table 18 . (b) A MIMIC validation radiograph, all 13 source clinical values, the missing MAP mask, and its logged fluid/vasopressor context are grouped together. Images retain their grayscale intensities; fitting and padding preserve aspect ratio. Future observations are evaluation references only.
MASS: Clinical19
MIMIC-IV-CXR: 13 clinical
Variable
Source value
Variable
Source value
FiO 2 (fraction)
0.35
FiO 2 (fraction)
0.467
PaO 2 (mmHg)
91.3
PaO 2 (mmHg)
105
PaCO 2 (mmHg)
44.6
SpO 2 (%)
97.42
pH
7.40
Respiratory rate (/min)
18.72
Respiratory rate (/min)
12
Heart rate (/min)
85.81
Appendix
Table 18: Full pre-encoder clinical records. MASS shows the Clinical19 extraction (19/19 observed), the same interface as the action evaluation record in Figure 11 ; MIMIC shows the 13-field source record in that figure (12/13 observed). Missing MAP remains masked. Units are shown where documented for these records.
Figure 12: From one multimodal source record to recursive future outputs. (a) Observed AP radiographs from the source and four future windows. Acquisition offsets are relative to the source radiograph, whereas horizon labels denote nominal six-hour event-window offsets. One image is displayed per window; the underlying image counts are 1, 1, 2, 2, and 3. The model uses the pooled window representation, not only the displayed view. (b) All 13 clinical input fields and their observed-value masks. (c–e) Three clinical coordinates decoded from the same four-step rollout, with observed references and source-value persistence. All other coordinates are retained in Table 19 . Only the source observations initialize the rollout; future images and measurements serve as evaluation references. Images retain their grayscale intensities with aspect-preserving fitting and padding. This single-patient illustration has no calibrated individual prediction intervals.
Variable
Source
H1: P / O
H2: P / O
H3: P / O
H4: P / O
Respiratory rate (/min)
24.09
23.68 / 20.30
23.36 / 21.38
23.15 / 23.12
23.06 / 23.20
SpO 2 (%)
99.33
99.04 / 99.83
98.74 / 99.33
98.48 / 98.67
98.28 / 95.67
FiO 2 (fraction)
0.40
0.40 / 0.40
0.40 / 0.40
0.40 / 0.40
0.40 / 0.40
PaO 2 (mmHg)
113.00
100.99 / –
96.96 / –
97.35 / 161.00
97.74 / –
Creatinine
–
1.29 / –
1.32 / –
1.35 / 0.60
1.38 / –
Urine output (ml)
31.67
45.09 / 30.83
57.67 / 48.75
68.89 / 143.33
78.46 / 125.00
Appendix
Table 19: Complete clinical outputs for the same multimodal case. Each future entry is prediction / observed reference (P / O); “–” means unobserved, with reference mask zero. Source values use physical units before normalization. Decoder means are shown even when predicted availability is below 0.5; these means are not treated as available inputs at the next step. Unrounded values, normalized inputs, masks, availability probabilities, and observed-cell errors are retained in the accompanying source data.
CXR representation error
H1
H2
H3
H4
PulmoWorld
0.902
0.780
0.939
0.756
Persistence
1.228
1.258
1.318
1.129
Appendix
Table 20: CXR representation outputs from the same rollout. RMSE is over 64 train-standardized PCA coordinates, not image pixels. Persistence repeats the source representation. These are the same image-state predictions already used in the complete-CXR case; no new image selection is performed.
Figure 13: Three MASS temporal CT examples. The frozen best-fit, median-fit, and failure-case roles are all retained. Source and future scans are observed validation images, separated by 3, 7, and 4 days, respectively. Only the original lung-window channel is displayed. Representative slices are not anatomically registered.
Figure 14: Four MIMIC input cases with different logged care contexts. Each row combines the original source radiograph, all 13 clinical values, missingness, and the following six-hour action record from the same validation case. The examples cover no recorded hemodynamic action and different fluid/vasopressor exposures. Dose summaries are factual context; these rows contain no observed counterfactual outcomes or treatment recommendations. They are distinct from the main-text forecasting cases.
Figure 15: MASS CT observations and token-space evolution. Validation CT pairs (best / median / worst set-distance gain after an SSIM screen). Columns are current 2D CT, observed future 2D CT, observed token change, and predicted token change. Future CT is observed, not generated. No locked-test patient is shown. These are the frozen spatial-token diagnostic cases, separate from the raw-input atlas.
Figure 16: MIMIC paired-CXR latent evolution. Validation data, separate from the 300-patient locked test; observed rather than generated films. (a) Frozen Qwen2.5-VL-7B encoder + world model versus source-film persistence. (b) VL-Win is the fraction of source/target-CXR pairs whose predicted representation is closer to the target than persistence. Joint@10 is the fraction of eligible clinical cells with at most 10% source-normalized error and a simultaneous visual win. The error denominator is the larger of the absolute source value and 0.1 times the training IQR (with a numerical floor of 10−8 ); these rates have different denominators. Tick labels give the share of windows with a target CXR ( ≈ 4%). (c) One hashed stay with a radiograph at every horizon. (d) Eight selected low-clinical-MSE stays with CXR at H1–H4. Stay labels are pseudonyms, not MIMIC identifiers.
Figure 17: Observed BALF metabolomic sequences and predicted state changes. Three development-validation patients with five BALF visits on protocol days 1/4/7/14/21 (3 of 63 validation patients). (a,d,g) 24 highest-variance ions after the train-only log1p+IQR map; rows index selected ions and colorbars report transformed intensities. (b,e,h) PCA64 scores for the same visits (PC1–PC6), with a shared legend. (c,f,i) recursive forecasting paired MSE gain versus carry-forward on PC1–PC16; positive means PulmoWorld is closer than persistence. The annotations report mean squared errors for the learned forecast and persistence. Direct VLMs do not emit BALF64. Validation examples only, not the locked-test recursive forecasting ranking.
Clinical decisions depend on anticipating how a patient's condition will evolve as treatment is adjusted and new evidence becomes available. Existing medical world models capture patient-state evolution, but often define the evolving state around a particular modality or task, limiting reuse as inputs or prediction targets change. These models also retain history through observation sequences or recurrent or latent summaries rather than exposing preceding patient states and clinical actions as an ordered transition context. This distinction matters because similar current measurements may follow different recent changes, treatment exposures, and observed responses. To tackle these limitations, we propose PatientWM, a clinical world-model framework built around a unified patient state and sequential latent history. Availability-aware encoders integrate observations available at each decision time into the patient state, while typed decoders produce task-specific clinical readouts without defining the state around any single target. PatientWM further conditions each next-state prediction on the current state and supplied action together with a bounded, ordered history of preceding latent patient states and clinical actions, preserving recent treatment and response context. Experiments across MIMIC-III, MIMIC-IV, and three external clinical datasets show superior performance in clinical prediction, multi-step forecasting, and action-planning evaluation. Transfer experiments further demonstrate the generalizability of PatientWM's patient state. Together, these results establish its ability to forecast clinical evolution and support planning.
Ziwen Kan, Jie Peng, Bo Yan +5
University of Science and Technology of China · University of Central Florida · Princeton University +2
Medical diagnosis and treatment are dynamic processes in which patient states evolve over time and clinical interventions alter future outcomes. Although current medical AI can detect disease, estimate risk and generate reports, many systems still return static labels or scores, offering limited insight into how illness may progress or how alternative interventions may reshape its trajectory. Medical world models adapt the world-model idea from artificial intelligence to healthcare by learning internal simulators of patient-state dynamics. Their long-term goal is to help clinicians anticipate deterioration, compare treatment-conditioned futures and tailor care to individual patients. Yet relevant work remains scattered across foundation models, longitudinal modelling, disease simulation, treatment-effect estimation, reinforcement learning and digital twins. To bridge this gap, this review outlines a roadmap for advancing medical AI from isolated diagnosis and prediction toward medical world models that simulate disease evolution and support intervention decisions. This roadmap is organized around three coupled capabilities: patient-state construction, clinical dynamics modelling and intervention decision support. Across representative systems, the comparison highlights what each capability contributes and how partial components can be integrated into more mature perception--dynamics--planning systems. Finally, we identify the challenges involved in turning plausible rollouts into clinically useful simulators. Related literature is available at https://github.com/1999kevin/awesome_medical_world_models.
Ke Liu, Mengxuan Li, Yanyi Bao +4
College of Computer Science, Zhejiang University, Hangzhou, China. · School of Medicine, Zhejiang University, China. · Department of Biomedical Informatics, Harvard University, USA.
World models have emerged as a unifying paradigm for learning latent dynamics, simulating counterfactual futures, and supporting planning under uncertainty. In this paper, we argue that computational epidemiology is a natural and underdeveloped setting for world models. This is because epidemic decision-making requires reasoning about latent disease burden, imperfect and policy-dependent surveillance signals, and intervention effects are mediated by adaptive human behavior. We introduce a conceptual framework for epidemiological world models, formulating epidemics as controlled, partially observed dynamical systems in which (i) the true epidemic state is latent, (ii) observations are noisy and endogenous to policy, and (iii) interventions act as sequential actions whose effects propagate through behavioral and social feedback. We present three case studies that illustrate why explicit world modeling is necessary for policy-relevant reasoning: strategic misreporting in behavioral surveillance, systematic delays in time-lagged signals such as hospitalizations and deaths, and counterfactual intervention analysis where identical histories diverge under alternative action sequences.
Zeeshan Memon, Yiqi Su, Christo Kurisummoottil Thomas +3
Department of Computer Science, Emory University, Atlanta, GA, USA · Department of Computer Science, Virginia Tech, Alexandria, VA, USA · Department of Electrical and Computer Engineering, Worcester Polytechnic Institute, Worcester, MA, USA +1