PaperlayerView exact revision on arXiv

Source-linked AI summary

Graph World Models for Constrained Epidemic Policy Planning

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

arXiv:2609.35545v2cs.LGcs.AI

TL;DR

Regional epidemic planning must account for mobility, scarce shared resources, and uncertain latent conditions, while existing approaches lack action-conditioned coupled models or per-period feasibility guarantees. EpiMind combines graph-based policy rollouts with constrained planning and projection. It reports 53.4% fewer admissions than no intervention while satisfying shared-resource constraints.

  • Problem

    Regional epidemic planning must coordinate mobility-coupled policies and scarce resources, but existing methods lack action-conditioned models of coupled dynamics or per-period feasibility guarantees.

    Source: §1 · pp. 1–2

  • Method

    EpiMind uses a graph-factored recurrent state-space model for joint policy-conditioned rollouts and graph-temporal ADMM with projection to plan feasible regional interventions.

    Source: §1 · p. 2, §2.1 · pp. 2–3, §2.2.3 · p. 5

  • Results

    53.4% fewer admissions than no intervention, with all shared-resource constraints satisfied in real-context evaluation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    EpiMind supports policy comparison and resource allocation through learned dynamics, with exact enforcement of specified linear resource budgets.

    Source: §4 · p. 9

    Predicted trajectories and spillover effects are model-based sensitivities, not causally identified counterfactuals, because real-world outcomes under alternative policies are unobserved.

    Source: §4 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Epidemic planning must coordinate regions coupled by mobility and scarce shared resources, while acting under uncertainty about latent conditions. EpiMind combines graph-based policy-conditioned prediction with constrained planning, and reports reduced admissions while satisfying shared-resource constraints.

1 INTRODUCTION

  • Planning models must infer latent epidemic conditions from noisy, delayed, policy-dependent observations and predict graph-coupled evolution under alternative actions.Surveillance trends may reflect transmission changes, testing changes, or voluntary behavioral adaptation.Source: §1 · p. 1
  • Existing approaches lack a combination of action-conditioned coupled dynamics and per-period shared-resource feasibility for regional epidemic planning.Forecasting methods lack action spaces, while common multi-agent RL penalties do not guarantee per-timestep feasibility; optimization methods require pre-fit transmission dynamics.Source: §1 · pp. 1–2
  • EpiMind couples a parameter-shared graph world model with a constrained multi-region planner for joint policy rollouts, feasible allocation, and temporal-logic evaluation.The framework formulates regional dynamics on a dynamic policy graph coupled through mobility and finite shared resources.Source: §1 · p. 2
  • 53.4% fewer admissions than no intervention, while satisfying all shared-resource constraints.The reported real-context evaluation also describes performance as near-oracle.Source: §1 · p. 2

2 METHOD

EpiMind models regional epidemics on a time-varying policy graph and plans over joint, policy-conditioned trajectories. Its planner projects actions onto shared-resource constraints and evaluates temporal specifications against the learned model.

2.1 PROBLEM FORMULATION

  • The formulation represents regions and mobility or coordination constraints as a time-varying policy graph, with regional interventions planned over multiple steps under shared-resource budgets.Actions include local NPI intensity and regional allocations of shared vaccines, hospital capacity, and fiscal resources; regional surveillance observations do not directly reveal latent epidemic states.Source: §2.1 · pp. 2–3

2.2 EPIMIND FRAMEWORK

  • EpiMind couples a learned graph world model with constrained receding-horizon planning to predict joint trajectories under regional interventions and coordinate actions under shared-resource constraints.The framework rerolls projected actions and evaluates temporal specifications relative to the learned model.Source: §2.2 · p. 3, Figure 1 · p. 3

2.2.1 GRAPH WORLD-MODEL ROLLOUTS

  • GF-RSSM shares neural parameters while retaining regional beliefs and states, recursively rolls joint actions through mobility-weighted context, and decodes admissions and occupancy.Future observations are unavailable during planning, so imagined trajectories use the learned prior recursively.Source: §2.2.1 · pp. 3–4

2.2.2 GRAPH-TEMPORAL CONSTRAINED PLANNING

  • GT-ADMM alternates regional action proposals, coordination updates, and shared-resource projection; its spillover signal predicts neighboring outcome changes.Regional proposals are optimized through the frozen world model while the rollout remains joint, allowing an action to alter connected regions' predicted outcomes.Source: §2.2.2 · pp. 4–5

2.2.3 PROJECTION, VERIFICATION, AND EXECUTION

  • After optimization, EpiMind rerolls the projected allocation, evaluates exact STL robustness on the learned trajectory, and executes its first action; this verification is model-relative.The executed action satisfies the resource budgets represented in the feasible set, while STL satisfaction is not certified for the unknown true environment.Source: §2.2.3 · p. 5

3 EXPERIMENTS

The experiments assess predictive fidelity, planning utility, constraint handling, coordination, and real-context transfer across simulator and U.S. state-level settings. Results cover policy-conditioned predictions, feasible plans, coordination comparisons, and outcomes in real-context evaluations.

3 EXPERIMENTS

  • Experiments address five questions spanning predictive fidelity, planning effectiveness, feasibility and verification, graph coordination, and real-context transfer.Source: §3 · p. 5

3.1 EXPERIMENTAL PROTOCOL

  • Evaluation spans a 5-region, 26-week mobility simulator, a retrospective U.S. state panel, and a real-initialized semi-simulation, with methods compared under common action bounds, budgets, projection, and objective.Source: §3.1 · pp. 5–6

3.2.1 GF-RSSM SUPPORTS ACCURATE POLICY-CONDITIONED PREDICTION.

  • GF-RSSM reduced held-out synthetic admission RMSE by 29% versus the no-graph ablation and preserved monotonic NPI dose ordering in simulation.It overpredicted admissions at low NPI and underpredicted them at high NPI; retrospective curves show stable model sensitivity, not causal effects.Source: §3.2.1 · p. 6, Table 1 · p. 6, Figure 2 · p. 7

3.2.2 LEARNED ROLLOUTS YIELDS EFFECTIVE RESOURCES-FEASIBLE INTERVENTIONS.

  • EpiMind stayed within 1.1%, 1.5%, and 4.9% of the best constant-NPI policy at λNPI = 3, 10, and 30, respectively, and outperformed the other planning baselines.All executed allocations had zero post-projection budget excess; paired regret was 0.93±1.10 at λNPI = 10.Source: Table 2 · p. 7, §3.2.2 · p. 7

3.2.3 RESOURCE PROJECTION GUARANTEES FEASIBLE EXECUTION.

  • All allocations satisfied encoded linear resource constraints with zero maximum budget excess; projected rollouts satisfied STL specifications with positive exact robustness under the learned model.Resource feasibility is guaranteed for encoded constraints, while STL satisfaction is model relative and does not certify the unknown environment.Source: Table 3 · p. 7, §3.2.3 · p. 7

3.2.4 COORDINATION BENEFITS ARE MODEST UNDER MATCHED INTERVENTION BURDEN.

  • At nearly identical NPI burden (0.423–0.424), EpiMind reduced J by 1.14% versus global-only ADMM, 2.13% versus graph-free planning, 2.21% versus time-shuffled planning, and 3.40% versus independent MPC.With stepwise matched burden, coordination gains fell below 1%, indicating much of the uncontrolled difference arose from intervention burden rather than allocation alone.Source: Table 4 · p. 8, §3.2.4 · p. 8

3.2.5 EPIMIND TRANSFERS TO REAL-CONTEXT MULTI-REGION PLANNING.

  • In retrospective Track A, EpiMind projected fewer admissions than most comparators, though Standard ADMM was marginally better on average; in semi-synthetic Track B, EpiMind outperformed every deployable comparator.Track A comparisons are model relative because candidate policies were not executed.Source: Figure 3 · p. 8, §3.2.5 · p. 8

4 CONCLUSION AND LIMITATIONS

EpiMind combines graph-factored regional beliefs and joint policy-conditioned rollouts with projected shared-resource coordination. Its constraints and temporal checks have distinct guarantees, and planning remains model-dependent.

4 CONCLUSION AND LIMITATIONS

  • EpiMind combines a parameter-shared GF-RSSM for joint policy-conditioned prediction with GT-ADMM for intervention coordination and projection onto feasible shared budgets.Experiments demonstrate accurate policy-conditioned prediction, effective planning through learned dynamics, and exact enforcement of specified linear resource budgets.Source: §4 · p. 9
  • Under matched intervention burden, graph coordination provides a measurable but incremental benefit.Real-world outcomes under alternative policies are unobserved, so predicted trajectories and spillover effects are model-based sensitivities rather than causally identified counterfactuals.Source: §4 · p. 9
  • Projection guarantees only the constraints encoded in the feasible set; STL satisfaction applies to learned trajectories, not the unknown environment.Planning quality depends on world-model calibration, and nonconvex GT-ADMM has no global convergence guarantee.Source: §4 · p. 9

5 SOFTWARE AND DATA

The paper reports that its full implementation is released online.

5 SOFTWARE AND DATA

  • The authors release the full implementation at the stated project README URL.https://anonymous.4open.science/r/ epimind-9706/README.mdSource: §5 · p. 9

AI USE DISCLOSURE

The section describes EpiMind’s epidemic representation and planning pipeline, from latent regional dynamics and shared-resource constraints to projected-action verification. It also includes the authors’ AI-use disclosure.

AI USE DISCLOSURE

  • The authors used generative AI tools for literature retrieval and discovery and to improve manuscript clarity and readability, then reviewed and revised all AI-assisted text.They state that citations and literature-derived statements were verified and that the authors take responsibility for the final content.Source: AI USE DISCLOSURE · p. 10

A EPIDEMIOLOGICAL FOUNDATIONS

  • EpiMind represents regions as nodes in a dynamic policy graph, with mobility-coupling edges and global allocations of vaccine, hospital capacity, and fiscal resources.The epidemiological structure section introduces the latent dynamics, observations, and planning constraints underlying the framework.Source: Figure A1 · p. 13, A EPIDEMIOLOGICAL FOUNDATIONS · p. 13

A.1 EPIDEMICS AS GRAPHS.

  • Latent states represent underlying epidemic conditions inferred behind noisy, delayed, policy-dependent surveillance; interventions can also have behavioral mediation and delayed effects.Reported cases may change with testing rather than true infections, while intervention effects can appear in surveillance only after a delay.Source: §A.1 · pp. 13–14
  • Shared resources impose per-period feasibility across pooled regional allocations, so each region’s feasible actions depend on other regions’ requests.The section also frames policy-conditioned prediction, cross-region coordination, and post-projection verification as distinct computational challenges.Source: §A.1 · p. 14, §A.2 · p. 14

A.2 FROM POLICY-CONDITIONED PREDICTION TO CONSTRAINED COORDINATION

  • The pipeline predicts joint trajectories under candidate actions, optimizes regional sequences with GT-ADMM, projects allocations, then rerolls projected actions for temporal verification.The reroll distinguishes the proposed trajectory from the trajectory under the action that will actually be executed; learning remains predictive rather than enforcing constraints through its training loss.Source: §A.2 · pp. 14–15

B EPIMIND FRAMEWORK

  • Algorithm 1 updates regional beliefs, warm-starts action sequences, performs parallel regional updates and shared-allocation projection, then rerolls and evaluates robustness.The algorithm describes a constrained receding-horizon planning loop.Source: B EPIMIND FRAMEWORK · p. 15

B.1 SIGNAL TEMPORAL LOGIC FOR EPIDEMIOLOGICAL RULES

  • The experiments require predicted hospital occupancy to stay below regional capacity throughout the horizon, using differentiable robustness for optimization and exact robustness for verification.Positive exact robustness certifies satisfaction on the learned rollout only, not in the unknown environment.Source: §B.1 · p. 15

C EXPERIMENTAL SETTINGS

The experimental settings describe a weekly U.S. state-level panel, comparison baselines, and an inventory of experimental parameters.

C.1 REAL-WORLD DATA SOURCES AND PREPROCESSING

  • The weekly U.S. state-level panel combines hospital and epidemic measures, CDC surveillance and vaccination, state policy data, mobility flows, and Census population estimates.Reported infections and deaths are treated as noisy observations, and directed state-to-state mobility flows define the time-varying graph.Source: §C.1 · p. 16

C.2 BASELINE METHODS

  • World-model baselines range from persistence and climatology to action-conditioned linear and recurrent predictors, including a GF-RSSM ablation without graph attention.Forecasting models share the same train/test split, observation units, and rollout protocol.Source: §C.2 · p. 16
  • Planning comparisons include constant, allocation-rule, greedy, MPC, ADMM, graph-free, independent, and shuffled controls, alongside policy-learning methods.MPC-SEIR is an informed simulator-based comparator rather than a deployable real-data method; adapted epidemic-policy baselines are representative adaptations, not exact reproductions.Source: §C.2 · p. 16

C.3 PARAMETER TABLE

  • The parameter inventory summarizes experimental epidemiological, intervention, observation, model, planner, budget, objective, training, and evaluation settings.The passage identifies Table A2 as the inventory but supplies no parameter values.Source: §C.3 · p. 17, Table A2 · p. 17

D.1 WORLD MODEL EVALUATION

The full GF-RSSM performs best on admission MAE and short-horizon cumulative error, but its advantage narrows at longer horizons. Removing graph attention increases rollout error.

D.1 WORLD MODEL EVALUATION

  • The full GF-RSSM achieves the lowest admission MAE and cumulative error at H = 5, while all learned models outperform persistence and climatology.Figure A2 reports admission MAE and relative cumulative admission error over H ∈{5, 10, 20}.Source: §D.1 · p. 18, Figure A2 · p. 18
  • At H = 20, the action-conditioned LSTM performs better than the GF-RSSM, indicating greater long-horizon drift in the GF-RSSM.Source: §D.1 · p. 18, Figure A2 · p. 18
  • Removing graph attention consistently increases rollout error, while the full GF-RSSM’s advantage narrows by H = 10.Source: §D.1 · p. 18, Figure A2 · p. 18

D.2 ROBUSTNESS ANALYSIS

Operational perturbations had small average effects on cumulative hospitalizations in the tested runs, whereas epidemiological model mismatch produced the largest degradation. The results are descriptive given the limited seeds.

D.2 ROBUSTNESS ANALYSIS

  • Operational perturbations changed cumulative hospitalizations by less than 1% on average and remained within seed-level variability across five paired seeds.The tested perturbations included graph noise, missing or delayed observations, a mid-horizon budget cut, and regional non-compliance.Source: §D.2 · p. 18, Figure A3 · p. 18
  • Epidemiological model mismatch produced the largest tested degradation, at +6.47%.Figure A3 identifies it as the only perturbation with substantially larger mean degradation than the operational perturbations.Source: §D.2 · p. 18, Figure A3 · p. 18
  • With only five paired seeds, the perturbation results are descriptive rather than evidence of statistical equivalence.Source: §D.2 · p. 18, Figure A3 · p. 18

D.3 PLANNING AND CALIBRATION ANALYSIS

Planning regret was higher with learned than oracle dynamics in every tested configuration, but the variation prevents establishing its magnitude. Calibration benefits were supported at intermediate intervention costs, not uniformly across costs or cells.

D.3 PLANNING AND CALIBRATION ANALYSIS

  • The learned planner has higher mean regret than the oracle-dynamics planner in every configuration, but no sign test is significant.The model-effect standard deviation exceeds its mean in every row, so the grid suggests a rollout-model contribution but does not establish its magnitude at this sample size.Source: §D.3 · p. 19, Table A3 · p. 19
  • The calibration ablation holds the planner and per-cell best-constant comparator fixed, enabling paired comparisons across variants.Table A4 reports mean ± s.d. over n=9 cells.Source: §D.3 · p. 19, Table A4 · p. 19
  • At λ=10, the calibrated hospitalization head reduces mean regret by 64.2% versus the shared decoder and 30.7% versus the uncalibrated head, with both paired tests significant.At λ=30, the reduction versus the shared decoder is 73.3%, but the additional benefit over the uncalibrated head is not significant; neither contrast is established at λ=3.Source: §D.3 · p. 19, Table A4 · p. 19

D.4 COORDINATION ABLATION

The coordination ablation finds incremental gains under matched intervention burden; the real mobility graph is directed and heterogeneous. In real-context evaluation, EpiMind combines the lowest reported forecast RMSE with lower projected admissions than the historical-policy reference, while proposed-policy trajectories remain unexecuted counterfactuals.

D.4 COORDINATION ABLATION

  • Full EpiMind achieves the lowest cumulative admissions (656.81/100K) under matched NPI burden; removing global resource consensus yields the largest individual degradation (+0.94%).Removing spillover sensitivity yields +0.54%, removing neighbor consensus yields +0.47%, and disabling all three signals increases admissions by 0.88%; gains are incremental once burden is controlled.Source: §D.4 · p. 19, Table A5 · p. 20

D.5 REAL MOBILITY GRAPH

  • The January 2021 row-normalized Advan graph for ten selected high-flow U.S. states has directed, heterogeneous flows used as graph-edge weights.Rows denote origins, columns destinations, and each cell gives the destination’s share of an origin’s outgoing travel among the displayed states.Source: §D.5 · p. 20, Figure A4 · p. 20

D.6 REAL-CONTEXT PREDICTION AND PROJECTION

  • EpiMind attains the lowest forecast RMSE (4.1) and projects 9.9 weekly admissions per 100K, versus 11.8 under the historical-policy reference.Several baselines also project admissions below the reference, but with substantially larger forecast errors; Figure A5 compares forecast RMSE with model-relative projected admissions.Source: §D.6 · p. 20, Figure A5 · p. 21
  • Across ten mobility-connected states, historical-policy forecasts track major admission-wave timing but imperfectly calibrate some peaks; proposed-policy rollouts are unexecuted model-relative comparisons, not causal effects.Figure A6 distinguishes observed admissions, recorded-policy forecasts, and proposed-policy planner projections; the dashed line marks the training cutoff.Source: §D.6 · pp. 20–21, Figure A6 · p. 21
Loading 2609.35545v2…