Treatment planning is intrinsically counterfactual: clinicians must compare how a patient's tumor and survival would evolve under competing regimens, whereas a retrospective cohort records only the regimen received and one follow-up scan. A treatment-conditioned generator trained on such data can therefore reproduce historical treatment allocation rather than simulate the interventional disease dynamics needed to rank regimens. We introduce CaMeWM, a causal medical world model for supported post-TACE tumor evolution and treatment planning. CaMeWM implements one intervention-to-decision loop rather than three independent prediction tasks: it first admits a queried regimen only for clinically comparable histories with empirical treatment support; it then generates a distribution of lesion-local follow-up worlds through an action-conditioned response descriptor, deterministic anchor, and residual 3D diffusion transition; and it finally maps sampled worlds to restricted mean survival time (RMST) with a response-Cox utility model. Under pre-treatment adjustment and positivity, this loop estimates a prespecified, support-restricted g-formula target; all selection and nuisance estimation are fold-local, and the induced policy is evaluated by cross-fitted, censoring-adjusted doubly robust estimation. Across 79 paired out-of-fold episodes, CaMeWM achieves 0.618 SSIM and 0.528 change-map SSIM for follow-up CT transition, while its generated-world survival pathway attains a 0.615 C-index; factual-only AIPW estimates a 9.60-week RMST gain (95% CI 0.61--18.13) for the frozen policy. CaMeWM provides an inspectable basis for retrospective, support-restricted treatment planning.
Figures & tables
Figure 1: CaMeWM replaces observational assignment by a supported do(A=a) query. Under consistency, conditional exchangeability, and positivity, the g-formula identifies the regimen-specific population world law and its RMST value on empirical support (Theorem 1 ; full assumptions and proof in the supplement). The response code, deterministic anchor, residual DDPM, and survival head estimate components of this identified functional; their architectures neither establish exchangeability nor identify a unique patient’s counterfactual CT. The lower rail performs leakage-free IPCW/DR policy evaluation. CT thumbnails are schematic, not patient data or results.
Method
Tumor HU-MAE ↓
Δ SSIM ↑
PSNR ↑
SSIM ↑
Persistence
57.192 [52.069, 63.049]
0.157 [0.135, 0.180]
16.877 [16.148, 17.661]
0.469 [0.443, 0.497]
3D U-Net
26.944 [23.224, 31.104]
0.523 [0.481, 0.566]
19.157 [18.359, 19.974]
0.610 [0.584, 0.637]
CVAE-3D
32.300 [28.065, 37.105]
0.508 [0.464, 0.549]
18.815 [18.027, 19.670]
0.587 [0.562, 0.612]
Mask-DDPM
63.721 [58.270, 69.962]
0.051 [0.026, 0.080]
16.665 [15.966, 17.362]
0.451 [0.425, 0.477]
MeWM ∗
62.910 [57.562, 68.555]
0.074 [0.048, 0.102]
16.684 [15.993, 17.373]
0.438 [0.412, 0.465]
CaMeWM
27.145 [23.139, 31.534]
0.528 [0.484, 0.569]
19.193 [18.353, 20.112]
0.618 [0.591, 0.645]
Table 1: Factual post-TACE CT generation on 79 paired OOF episodes. Brackets denote patient-bootstrap 95% confidence intervals. Tumor HU-MAE is reported after inverse transformation to Hounsfield units. ∗ Task-matched reimplementation of MeWM’s action-conditioned generative-dynamics principle, adapted to the two regimens and public HCC-TACE-Seg outputs available in this study. All methods use the same input state, preprocessing, and patient-level OOF folds.
Table 2: Treatment-planning behavior and factual-only causal evaluation over five patient-level out-of-fold folds. Main estimates are shown above patient-bootstrap 95% confidence intervals. Agreement metrics compare recommendations with observed historical regimens and are behavioral references, not counterfactual treatment truth. † Defined only for CaMeWM, which applies a propensity-support guardrail; “–” denotes not applicable. ∗ Task-matched reimplementation inspired by MeWM’s action-conditioned generative-dynamics principle.
Variant
Normalized tumor MAE ↓
SSIM ↑
Global HU-MAE ↓
Δ SSIM ↑
Full CaMeWM
0.070 [0.059, 0.082]
0.618 [0.591, 0.644]
61.102 [54.384, 68.512]
0.528 [0.483, 0.570]
Causal anchor only
0.070 [0.060, 0.081]
0.612 [0.585, 0.637]
61.139 [54.508, 68.499]
0.526 [0.482, 0.569]
w/o response-code conditioning
0.071 [0.060, 0.082]
0.618 [0.591, 0.643]
61.198 [54.274, 68.508]
0.528 [0.485, 0.571]
w/o direct action path
0.075 [0.064, 0.087]
0.617 [0.591, 0.644]
62.239 [55.086, 69.746]
0.524 [0.480, 0.566]
w/o confidence weighting
0.070 [0.060, 0.082]
0.618 [0.592, 0.644]
61.197 [54.610, 68.239]
0.527 [0.482, 0.570]
Table 3: Controlled generation ablation over the same five patient-level out-of-fold folds. Normalized tumor MAE is computed on the model’s [−1,1] CT scale and is not numerically comparable to the HU-MAE in Table 1 . “Deterministic anchor only” removes residual-DDPM refinement; “uniform pseudo-mask weighting” sets all pseudo-mask confidence weights to one. Brackets denote patient-bootstrap 95% confidence intervals.
Method
Uno C ↑
IBS ↓
RMST MAE ↓
C-index ↑ 95% CI
CoxPH
0.580
0.202
39.28
0.581 [0.500, 0.661]
DeepSurv
0.592
0.200
39.17
0.597 [0.525, 0.670]
IPTW-Cox
0.603
0.200
38.73
0.609 [0.531, 0.689]
MeWM ∗
0.461
0.242
44.53
0.462 [0.390, 0.540]
CaMeWM factual
0.589
0.211
40.80
0.592 [0.521, 0.662]
CaMeWM generated worlds
0.612
0.202
39.23
0.615 [0.546, 0.688]
Table 4: Patient-level out-of-fold survival prediction over five folds. The generated-world row evaluates the complete action-conditioned transition-and-utility path. Brackets denote patient-bootstrap 95% confidence intervals. ∗ Task-matched reimplementation inspired by MeWM’s action-conditioned generative-dynamics principle.
Figure 2: Out-of-fold transition fidelity and paired comparison. CaMeWM synthesizes action-conditioned follow-up CT volumes in registered image space, rather than tumor masks. Each row shows pre-treatment CT, observed post-treatment CT, and predictions from 3D U-Net and the mean of eight CaMeWM worlds. Cyan dashed contours mark the pre-treatment lesion only for orientation; they are neither follow-up annotations nor model outputs. Tumor-region MAE uses a fold-local pseudo post-treatment mask, whereas SSIM is computed over the complete follow-up CT. Rows are deterministically selected at the 20th, 50th, and 80th MAE percentiles.
Main-paper contribution
Formal support
Required implementation and evidence
Support-restricted interventional world distribution
Theorem S1 and its complete proof
Joint-regimen propensity, explicit support set, response-factorized g-formula, overlap diagnostics, and semi-synthetic identification recovery.
Response-Cox RMST, K=8 worlds per action, sampling-convergence analysis, and support abstention. The current score is not claimed to be a calibrated finite-sample confidence bound.
Cross-fitted factual policy evaluation
Theorem S6 , its proof, and the censoring identity
survival, planning, support, and policy-value summaries
OOF aggregation
patient-level pooling and paired bootstrap resampling
reported tables, figures, and confidence intervals
Table S4: Reproducibility materials and audit trail. The table intentionally describes analytical stages rather than release-specific script or archive names.
Stage
Episodes
Public HCC-TACE-Seg release
105
Longitudinal local manifest with required imaging and clinical records
99
Two predeclared supported complete joint regimens
80
Empty released source label-2 tumor mask excluded
1
Final patient-level OOF cohort
79
Table S5: Outcome-blind cohort construction for the fixed five-fold OOF protocol.
log(1+AFP) , then training median-centering and SD scaling
Child–Pugh class
Ordinal release code
0 (0.0)
Training median-centering and SD scaling
BCLC stage
Ordinal release code
0 (0.0)
Training median-centering and SD scaling
Tumor count
Binary release code
0 (0.0)
Training median-centering and SD scaling
Table S7: Pre-treatment adjustment covariates and fold-local missing-data handling in the final 79-episode cohort. “Missing” is assessed before any outer-fold preprocessing. All covariates enter the joint propensity and the response-Cox utility model together with their binary missingness indicators.
Figure S1: Observed-action propensity diagnostics. (a) Distribution of OOF observed-regimen propensities; the dashed line is the predeclared support threshold. (b) Action-wise OOF calibration. (c) Standardized mean differences before weighting, after IPTW, and after overlap weighting for the non-missing clinical covariates. These diagnostics characterize observed support and measured-covariate balance; they do not establish exchangeability.
Figure S2: Factual-to-generated survival transport diagnostic. (a–c) IPCW reliability at 12, 24, and 36 months for factual and generated states. (d) Patient-level agreement between factual-state and generated-world predicted RMST. The diagonal is identity. The paired C-index change is reported with a patient-bootstrap 95% interval and is interpreted only as a diagnostic for generated-state shift.
Figure S3: Stability with the number of sampled worlds. (a) OOF recommendation agreement with the K=32 reference over K∈{2,4,8,16,32} . At the locked K=8 , 78/79 (98.7%) recommendations match K=32 . (b) Corresponding recommended-regimen composition and abstention fraction. This supports the computational choice of K=8 ; it is not a finite-sample confidence-calibration result.
Figure S4: Pseudo follow-up mask confidence audit. (a) Spearman associations, with patient-bootstrap 95% intervals, between three fold-local confidence signals and OOF tumor-region MAE. (b) Direct confidence-versus-error scatter. The pseudo-label confidence association is weak ( ρ=0.20 , 95% CI [−0.03,0.41] ), so these quantities are treated as audits rather than calibrated error predictors.
Figure S5: Known-ground-truth semi-synthetic validation. (a) ATE and policy-value recovery with increasing simulated sample size. (b) Double-robustness identity under correct or misspecified propensity and outcome nuisances. (c) Positivity stress. (d) Factor-wise causal-value bound. (e) Monte Carlo convergence with sampled-world count. (f) Sensitivity to an unmeasured common cause. The experiment tests the behavior of the estimators and theoretical bounds under the simulation SCM, not clinical efficacy.
Quantity
OOF estimate
Patients
79
Factual observed-regimen RMST (weeks)
103.62
Factual-only AIPW policy value (weeks)
113.22 [102.34, 124.84]
Factual-only AIPW gain (weeks)
9.60 [0.61, 18.13]
Policy effective sample size
37.4
Mean recommended-action support
0.597
Table S8: Primary factual-only AIPW evaluation of the frozen CaMeWM policy. The outcome nuisance uses pre-treatment clinical history, factual regimen, and observed survival only.
Figure S6: Support-threshold and IPW-cap sensitivity of a historical generated-state plug-in diagnostic. Each cell uses the fixed five-fold OOF planning and propensity predictions. Rows vary the joint observed/recommended support threshold η and columns vary the IPW cap. Panel (a) reports the descriptive plug-in RMST value gain, while panel (b) reports effective sample size. The locked primary setting is η=0.05 with cap 10. This analysis assesses overlap and weight sensitivity; it is not the primary factual-only AIPW estimand and does not convert an observational interval into evidence of clinical benefit.
Clinical decision-making in oncology requires predicting dynamic disease evolution, a task current static AI predictors cannot perform. While world models (WMs) offer a paradigm for generative prediction, existing medical applications remain limited. Existing methods often rely on stochastic diffusion models, focusing on visual reconstruction rather than causal, physiological transitions. Furthermore, in medical domain, models like MeWM typically ignore patient-specific temporal and clinical contexts and lack a feedback mechanism to link predictions to treatment decisions. To address these gaps, we introduce CLARITY, a medical world model that forecasts disease evolution directly within a structured latent space. It explicitly integrates time intervals (temporal context) and patient-specific data (clinical context) to model treatment-conditioned progression as a smooth, interpretable trajectory, and thus generate physiologically faithful, individualized treatment plans. Finally, CLARITY introduces a novel prediction-to-decision framework, translating latent rollouts into transparent, actionable recommendations. CLARITY demonstrates state-of-the-art performance in treatment planning. On the MU-Glioma-Post dataset, our approach outperforms recent MeWM by 12%, and significantly surpasses all other medical-specific large language models.
Tianxingjian Ding, Yuanhao Zou, Chen Chen +2
Institute of Artificial Intelligence, University of Central Florida
World Models (WM) are increasingly seen as a foundation for intelligent agents that can predict, plan, and act beyond their training distribution. In this paper, we study WMs from a causal perspective across multiple levels of abstraction, ranging from perceptual observations to building a conceptual representation of the structure governing the environment dynamics. We argue that useful WMs must go beyond generative capabilities alone: they should also capture entity properties, entity-to-entity interactions, and entity-to-environment interactions that determine and explain the dynamics of a system. We provide a formal definition of Causal WMs (CWMs) grounded in the tasks they are intended to support, connecting world modelling with existing work in causal representation learning, object-centric learning, causal discovery, structural causal models, and model-based decision-making. Finally, we relate CWMs to the literature on identifiability, clarifying when the components of a WM can be recovered from data and up to which equivalence. With this, we ground WMs in representations and structures that support causal reasoning and informed decision-making.
Avinash Kori, Fabrizio Russo
Department of Computing, Imperial College London, London, UK
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