Graph World Models for Constrained Epidemic Policy Planning

Yiqi Su; Rashed Shelim; Lingyi Wang; Walid Saad; Naren Ramakrishnan

cs.LGSep 28, 2026

Graph World Models for Constrained Epidemic Policy Planning

Authors: Yiqi Su, Rashed Shelim, Lingyi Wang, Walid Saad, Naren Ramakrishnan

Organizations: Department of Computer Science, Virginia Tech, Alexandria, VA 22305, USA · Department of Electrical and Computer Engineering, Virginia Tech, Alexandria, VA 22305, USA

Abstract

Epidemic policy planning often requires coordination between geographical regions, taking into account mobility-driven spillovers and how to make use of limited resources. Existing methods either lack action-conditioned models of coupled dynamics or cannot guarantee per-period feasibility. We present EpiMind, a graph world model framework for constrained epidemic policy planning across regions. A graph-factored recurrent state-space model generates joint policy-conditioned rollouts from regional latent beliefs, while graph-temporal ADMM optimizes regional interventions, enforces shared-resource feasibility through projection, and evaluates temporal specifications under the learned model. EpiMind reduces admission RMSE by 29% relative to graph-free dynamics modeling, plans within 1-5% of the best feasible constant policy with guaranteed shared-budget feasibility, and outperforms all deployable baselines across three resource budgets in real-context evaluation. These results demonstrate that graph-structured policy imagination with explicit constrained coordination supports effective epidemic interventions from learned dynamics.

Figures & tables

Figure 1: EpiMind framework. A parameter-shared GF-RSSM updates regional beliefs and generates graph-coupled policy rollouts. GT-ADMM coordinates and projects regional actions, rerolls the feasible allocation for model-relative STL evaluation, and executes its first action.

Statistical baselines
Model	Adm. MAE $\downarrow$	Adm. RMSE $\downarrow$	Cum. err. @5 $\downarrow$	Params
Persistence	1.563 $\pm$ 0.053	1.927 $\pm$ 0.056	0.089 $\pm$ 0.003	0
Climatology	1.129 $\pm$ 0.021	1.283 $\pm$ 0.027	0.528 $\pm$ 0.011	0
Ridge $+$ action	0.341 $\pm$ 0.008	0.437 $\pm$ 0.003	0.068 $\pm$ 0.006	84
VARX(1) $+$ action	0.329 $\pm$ 0.008	0.424 $\pm$ 0.001	0.066 $\pm$ 0.007	1,960
Learned dynamics models

Table 1: World-model fidelity in the synthetic environment. Values are mean

\pm

s.d. Admission errors are measured per 100K; cumulative error is over a five-step rollout. Lower is better.

Figure 2: Policy-conditioned admission response. (a) Peak weekly admissions under alternative NPI intensities applied from a common synthetic state. (b) Predicted cumulative admissions under the same sweep at four held-out U.S. decision origins.

Method	$\lambda=3$	$\lambda=10$	$\lambda=30$	mean $\pm$ s.d.
	Objective $J$ $\downarrow$			Regret at $\lambda=10$ $\downarrow$
Best feasible constant †	54.7	58.9	70.9	0.00 $\pm$ 0.00
PPO	60.4	64.6	76.6	5.69 $\pm$ 3.71
MPC-SEIR	108.6	112.8	124.6	53.9 $\pm$ 30.9
Independent MPC	366.3	431.1	394.1	372 $\pm$ 307
Greedy	7,186	7,189	7,197	7,130 $\pm$ 1,221

Table 2: Matched-objective planning performance. Values are mean

\pm

s.d. Lower is better.

Setting	Feasible (%)	Max excess	Projection displacement	STL satisfaction (%)	STL robustness
Synthetic	100.0	0	0.476 $\pm$ 0.201	100.00	0.0021 $\pm$ 0.0004
Real context	100.0	0	0.0039 $\pm$ 0.0400	100.00	0.0016 $\pm$ 0.0003

Table 3: Constraint handling and model-relative verification.

Method	Adm./100K $\downarrow$	NPI burden	$J^{\dagger}$ $\downarrow$	$\Delta J$ (%)	EpiMind wins ( $p$ )
Standard ADMM	30.24	0.423	31.51	$+1.14$	23/36 (0.132)
Centralized MPC	30.65	0.424	31.92	$+2.45$	32/36 ( $<\!0.001$ )
Independent MPC	30.94	0.424	32.21	$+3.40$	26/36 (0.011)
Uniform allocation	31.23	0.479	32.67	$+4.87$	30/36 ( $<\!0.001$ )
EpidRLearn	35.64	0.500	37.14	$+19.23$	36/36 ( $<\!0.001$ )
PPO	35.69	0.499	37.19	$+19.37$	36/36 ( $<\!0.001$ )

Table 4: Realized policy performance in the real-context semi-synthetic evaluation.

Figure 3: Real-context policy comparisons. Each point reports the mean paired admission difference $\Delta\mathrm{Adm}=\mathrm{Adm}_{\mathrm{comparator}}-\mathrm{Adm}_{EpiMind{}}$ ; positive values (blue circles) favor EpiMind; horizontal bars denote 95% confidence intervals. (a) Retrospective Track A reports model-relative projections at held-out U.S. state-level decision origins. (b) Semi-synthetic Track B reports realized outcomes under known simulation dynamics. Diamonds denote the track-specific reference policy.

Figure 4: Texas case study. (a) Admission forecasts under the observed policy; the inset enlarges weeks 50–60. (b) Model-relative counterfactual trajectories under matched resource constraints; the observed trajectory provides context but is not an outcome of the unexecuted policies. (c) EpiMind’s first-step allocation compared with historical actions in Texas and neighboring states. The dashed vertical line marks the training cutoff. MPC and RL baselines appear only in panel (b).

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure A1: Multi-region policy graph under shared resource constraints. Regions are represented as nodes $v_{i}$ in a dynamic policy graph $G_{t}=(V_{t},E_{t})$ , with edges encoding interregional mobility coupling. Each region receives allocations of vaccine, hospitalization-capacity, and fiscal resources. The joint regional allocations must satisfy graph-level resource budgets, $\sum_{i=1}^{N}a_{t,k}^{i}\leq B_{t,k}$ , for $k\in\{v,h,f\}$ .

Challenge	Epidemiological property	Component	Computational response
(A)	Latent disease burden	GF-RSSM	Latent belief inference
	Policy-dependent surveillance	GF-RSSM	Action-conditioned observation model
	Delayed intervention effects	GF-RSSM	Multi-horizon policy-conditioned rollout
	Cross-region spillovers	GF-RSSM	Graph attention over neighbors
(B)	Shared resource limits	GT-ADMM	Capped-simplex projection
	Cross-border coordination	GT-ADMM	Neighbor-consensus signal

Appendix

Table A1: Epidemiological challenges and their computational treatment in EpiMind.

\hat{\mathbf{x}}_{t+1:t+H}=\mathcal{M}_{\theta}\!\left(\mathbf{b}_{t},\mathbf{z}_{t},a^{i},\mathbf{a}^{-i,k};G_{t}\right)

Appendix

Algorithm 1 EpiMind constrained receding-horizon planning

Symbol	Description	Value	Reference / Justification
Epidemiological
$\beta$	Transmission rate	0.25/day	Park et al. (2020) ; $R_{0}=\beta/\gamma=2.5$
$\sigma$	E $\to$ I rate	0.25/day (4 d)	Li et al. (2020) ; Lauer et al. (2020)
$\gamma$	I $\to$ R rate	0.1/day (10 d)	He et al. (2020)
$\delta_{H}$	Hospitalization rate	0.15	Garg et al. (2020) (pre-Omicron)
$\delta_{D}$	Base death rate	0.01	Meyerowitz-Katz and Merone (2020)

Appendix

Table A2: Parameter inventory.

Figure A2: Extended world-model evaluation on the synthetic benchmark. (a) Held-out admission MAE for learned architectures and statistical baselines under the same evaluation protocol. Error bars for learned models show mean $\pm$ s.d. over three training seeds. (b) Relative cumulative admission error over open-loop rollout horizons $H\in\{5,10,20\}$ , shown on a logarithmic scale.

Figure A3: Sensitivity to operational and model perturbations. Changes in cumulative hospitalizations relative to each seed-matched baseline, reported as mean $\pm$ s.d. The shaded band shows the largest within-condition standard deviation among the operational perturbations. Graph noise, masked regional observations, reporting delays, a mid-horizon budget cut, and regional non-compliance remain within this descriptive variability band. Epidemiological model mismatch produces the only substantially larger mean degradation.

$\lambda$	$H$	Regret $\downarrow$		Model effect $\downarrow$	$p^{\dagger}$
		Learned	Oracle
3	4	32.5 $\pm$ 32.8	27.8 $\pm$ 3.2	4.7 $\pm$ 34.6	1.000
3	8	38.2 $\pm$ 50.0	17.6 $\pm$ 0.7	20.5 $\pm$ 50.3	1.000
3	12	32.5 $\pm$ 34.2	20.6 $\pm$ 0.2	11.9 $\pm$ 34.2	1.000
10	4	217.0 $\pm$ 229.3	103.2 $\pm$ 3.8	113.9 $\pm$ 231.3	1.000
10	8	184.9 $\pm$ 201.1	79.3 $\pm$ 0.4	105.6 $\pm$ 200.9	1.000

Appendix

Table A3: Learned- versus oracle-dynamics planning regret. Values are mean

\pm

s.d.

$\lambda$	Calibrated	Shared dec.	Uncalib.	cells	$p$	cells	$p$
	Regret $\downarrow$			vs. decoder		vs. uncalibrated
3	56.5 $\pm$ 50.5	72.8 $\pm$ 85.4	68.8 $\pm$ 42.2	5/9	1.000	7/9	0.180
10	252.8 $\pm$ 242.3	705.3 $\pm$ 586.5	364.9 $\pm$ 309.3	8/9	0.039	9/9	0.004
30	641.4 $\pm$ 474.1	2401.9 $\pm$ 906.1	882.1 $\pm$ 634.7	9/9	0.004	7/9	0.180

Appendix

Table A4: Calibration ablation. Values are mean

\pm

s.d. over

n{=}9

cells.

Variant	$\boldsymbol{\gamma}$	$\boldsymbol{\eta}$	$\boldsymbol{\mu}$	Cum. adm./100K $\downarrow$	$\boldsymbol{\Delta}$ (%)
Graph-free	–	–	–	662.57	$+0.88$
No edge ( $\gamma$ off)	–	✓	✓	659.87	$+0.47$
No spillover ( $\eta$ off)	✓	–	✓	660.35	$+0.54$
No global ( $\mu$ off)	✓	✓	–	662.99	$+0.94$
EpiMind	✓	✓	✓	656.81	—

Appendix

Table A5: Matched-burden coordination ablation. Simulator-realized outcomes on the synthetic benchmark. All variants follow the reference configuration’s stepwise NPI-burden trajectory. Here,

\gamma

denotes neighbor consensus,

\eta

spillover sensitivity, and

\mu

global resource consensus.

\Delta

is the percentage change relative to the full model; lower is better.

Figure A4: Real interstate mobility graph. Row-normalized Advan device-mobility flows among ten selected high-flow U.S. states in January 2021. Rows denote origins, columns denote destinations, and color indicates each destination’s share of an origin’s outgoing travel among the displayed states. The asymmetric, nonuniform matrix provides mobility-edge weights to the graph world model.

Panel of Figure 19 — Figure A5: Texas real-context evaluation. (a) Observed admissions, the held-out forecast under historical actions, and the model-relative rollout under EpiMind’s proposed policy. The inset enlarges weeks 50–62; shading denotes the predicted difference between the historical- and proposed-policy rollouts. (b) Forecast RMSE versus model-relative admissions projected under each model’s optimized policy. Dashed lines mark the observed-policy admission mean and the selected forecast-error reference. Lower values are preferable on both axes.

Figure A6: Multi-state forecasting and planner projections. Observed weekly admissions (gray), forecasts under the recorded policy (blue), and model-relative rollouts under EpiMind’s proposed policy (red) for ten mobility-connected U.S. states. The vertical dashed line marks the training cutoff. Planner projections represent unexecuted counterfactuals, not observed outcomes.

Explore similar work

SL-BiLEM: Structured Learnable Behavior-in-the-Loop Epidemic Modeling for Forecasting and Policy Evaluation
May 26, 2026Haochun Wang, Sendong Zhao, Jingbo Wang +3Epidemiological Models Policy Evaluation
Faculty of Computing, Harbin Institute of Technology, Harbin, China
Understanding Rollout Error in Graph World Models
Jun 26, 2026Xinyuan Song, Zekun CaiWorld Models Prediction Error
Emory University, Atlanta, GA, USA · The University of Tokyo, Tokyo, Japan · LocationMind, Tokyo, Japan
Beyond Static Graph World Models: Learning Stochastic Latent Dynamics over Evolving Topologies
Sep 23, 2026Alex Schutz, Nick Hawes, Victor-Alexandru DarvariuLatent Dynamics Topology
Oxford Robotics Institute University of Oxford

May 26, 2026, cs.LGJ/K move · Enter open · S save

SL-BiLEM: Structured Learnable Behavior-in-the-Loop Epidemic Modeling for Forecasting and Policy Evaluation

Haochun Wang, Sendong Zhao, Jingbo Wang, Yanrui Du+2

Faculty of Computing, Harbin Institute of Technology, Harbin, China

Epidemiological Models Policy Evaluation

Open paper AskPDF arXiv

Epidemic forecasting faces a fundamental challenge: human behavior dynamically responds to disease spread, creating feedback loops that induce distribution shifts at policy intervention points. This renders data-driven models unreliable under distribution shift. We propose \textbf{SL-BiLEM} (Structured Learnable Behavior-in-the-Loop Epidemic Model), leveraging physical constraints as regularization for robust extrapolation. The framework decomposes effective transmission as

β_{\text{eff}}(t,g) = β_0(g) \times m_{\text{policy}}(t) \times m_{\text{media}}(t) \times m_{\text{comp}}(t,g)

, where monotonicity, smoothness, and bounded-jump constraints on the learned compliance function maintain predictive validity under novel policy regimes. Beyond forecasting, SL-BiLEM enables counterfactual analysis for intervention decision support. We validate forecasting on three real-world datasets (cruise ship, school influenza, and school-district COVID-19 surveillance) and evaluate counterfactual recovery on synthetic benchmarks with known ground truth. SL-BiLEM demonstrates: (1) 76% improvement over neural-mechanistic baselines, with only 53% OOD degradation versus 1142% for neural baselines under policy-induced shift; (2) 100% bootstrap CI coverage across 27 synthetic counterfactual experiments; and (3) Treatment Effect Accuracy exceeding 0.85. These results establish SL-BiLEM as an interpretable tool for public health decision-makers seeking accurate prediction and principled intervention planning.

Abstract

Explore similar work

SL-BiLEM: Structured Learnable Behavior-in-the-Loop Epidemic Modeling for Forecasting and Policy Evaluation

Understanding Rollout Error in Graph World Models

Beyond Static Graph World Models: Learning Stochastic Latent Dynamics over Evolving Topologies