EpiWorld: Grounding LLM Policy Agents in Epidemiological World Models

Zeeshan Memon; Yiqi Su; Kai Shu; Naren Ramakrishnan; Liang Zhao

cs.CLOct 2, 2026

EpiWorld: Grounding LLM Policy Agents in Epidemiological World Models

Authors: Zeeshan Memon, Yiqi Su, Kai Shu, Naren Ramakrishnan, Liang Zhao

Organizations: Department of Computer Science, Emory University, Atlanta, GA, USA · Department of Computer Science, Virginia Tech, Alexandria, VA, USA

Abstract

Epidemic intervention policies are textual artefacts that human decision-makers interpret, justify, and revise through natural language, making large language models a natural candidate for epidemic policy reasoning. A naive LLM, however, lacks the epidemic dynamics needed to project intervention consequences, the quantitative surveillance signals required to assess severity, and the institutional constraints that define admissible actions. We present EpiWorld, a closed-loop framework that grounds an LLM policy actor in a learned action-conditioned epidemiological world model and a tiered skill library of public-health protocols, surveillance tools, and adaptive lessons accumulated through after-action analysis. Given a candidate intervention, the world model predicts regional epidemic evolution and enables fast counterfactual rollouts that provide feedback for policy selection and refinement. Outcomes of simulated futures are distilled into reusable lessons while protocol constraints remain fixed, allowing the decision process to improve without sacrificing interpretability or controllability. We evaluate both the world model and the end-to-end framework on retrospective COVID-19 and Influenza datasets: the world model achieves the best out-of-distribution Peak-MAE among all forecasting baselines, and the closed-loop framework reduces cumulative hospitalisation by up to 59% across datasets and by an average of ~16% across six LLM backbones, outperforming reinforcement-learning and optimal-control policy baselines.

Figures & tables

Figure 1: EpiWorld framework. (a) The policy actor retrieves protocol constraints, surveillance skills, and adaptive lessons. (b) The actor generates interventions and reasoning traces. (c) Counterfactual rollouts evaluate candidate policies and produce updated lessons. Yellow highlights denote retrieved evidence and activated skills.

	COVID-19 SEIR		Influenza SEIR		COVID-19 Real US (Omicron)
Method	In-Dist. Peak / MAE	OOD Peak / MAE	In-Dist. Peak / MAE	OOD Peak / MAE	In-Dist. Peak / MAE	OOD Peak / MAE
Compartmental-GP	0.062 / 0.051	0.558 / 0.533	0.354 / 0.331	1.139 / 1.084	0.332 / 0.178	1.102 / 1.068
MechBayes (naive)	0.730 / 0.700	0.860 / 0.831	4.630 / 4.560	1.750 / 1.681	2.550 / 2.480	0.980 / 0.930
MechBayes + Policy	0.470 / 0.448	0.490 / 0.461	2.970 / 2.905	0.980 / 0.929	1.700 / 1.620	0.630 / 0.590
epidemia Mode A (naive)	0.680 / 0.648	0.750 / 0.709	4.300 / 4.185	1.500 / 1.431	2.400 / 2.310	0.910 / 0.860
epidemia Mode B (pretrained)	0.210 / 0.194	0.250 / 0.228	0.770 / 0.723	0.500 / 0.460	0.430 / 0.390	0.310 / 0.280

Table 1: World-model forecast performance (12-week autoregressive rollout, log1p per-100k). Each cell reports Peak-MAE / MAE ; lower is better in both. † PAN-cODE forecasts hospitalisations and deaths only.

	COVID-19 Real US	COVID-19 SEIR	Influenza SEIR
	(Omicron)
Method	Hosp. $\Delta$ (8w) $\downarrow$	Hosp. $\Delta$ (8w) $\downarrow$	Hosp. $\Delta$ (8w) $\downarrow$
EpiPolicy-RL	$-22.2\%$	$-8.2\%$	$-12.4\%$
ADIOS (GGGP)	$-33.4\%$	$-12.5\%$	$-18.6\%$
EpidRLearn	$\mathbf{-37.5\%}$	$\mathbf{-13.9\%}$	$\mathbf{-20.9\%}$
EpiWorld (Base LLM)	$-46.0\%$	$-17.1\%$	$-25.6\%$

Table 2: Policy optimisation performance over 8-week rollouts, reported as cumulative hospitalisation reduction against the reference historical policy (lower is better). For scale, the reference historical policy yields

406

cumulative hospitalisations per 100k over the 8-week horizon against

1{,}544

under no intervention. Policy-alignment diagnostics are reported separately in Appendix J .

Figure 2: Case study: Ohio during the Delta wave. Highlighted text indicates retrieved key reasoning steps.

Method	In-model	Indep. sim	Ratio
EpiPolicy-RL	$-8.2$	$-8.04$	$0.98$
ADIOS (GGGP)	$-12.5$	$-12.13$	$0.97$
EpidRLearn	$-13.9$	$-14.73$	$1.06$
EpiWorld (Base LLM)	$-17.1$	$-18.30$	$1.07$
EpiWorld (Static Skills)	$-17.8$	$-19.58$	$1.10$
EpiWorld (Skill Evol.)	$\mathbf{-22.0}$	$\mathbf{-25.52}$	$1.16$

Table 3: Independent-simulator validation, COVID-19 SEIR state-OOD. Cumulative hospitalisation reduction (%) against the reference historical policy, scored inside the learned world model

f_{\theta}

and re-executed in the mechanistic simulator that generated the corpus. Ratio is off-model over in-model. Real US is observational data and admits no ground-truth generator.

Model	Base LLM	w/ Static Skills	w/ Skill Evolution	$\Delta$ %
Proprietary
GPT-4o	0.1334	0.1371	0.0936	$\mathbf{-29.8\%}$
GPT-5-mini	0.1252	0.1468	0.1172	$-6.4\%$
GPT-5	0.1413	0.1468	0.1395	$-1.3\%$
Open-weights
Qwen2.5-7B	0.0884	0.0875	0.0692	$-21.6\%$

Table 4: Cross-backbone ablation on COVID-19 SEIR (state-OOD). Cumulative hospitalisation per 100k under the proposed policy over the 8-week post-intervention horizon (lower is better).

\Delta\%

is the relative reduction of the full Skill Evolution configuration over Base LLM.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

	COVID-19	Influenza-like Illness (ILI)
Used for	forecast + policy	observation model only
Period (ISO weeks)	2020-W01 to 2022-W52	2012-W30 to 2013-W40 (PAHO), 2010-W40 to 2015-W20 (CDC)
Panels / records	51 states $\times$ 156 weeks ( $7{,}956$ )	15 country panels $\times\sim\!130$ weeks each
Cases / hosp / deaths	JHU CSSE / HHS Protect / JHU CSSE	ILINet wILI / n/a / n/a
Action grid	13-dim OxCGRT ( Hale et al., 2021 )	none
Special structure	per-week point estimates	per-week revision history (initial $\to$ final)

Appendix

Table 5: Legacy two-panel dataset description preserved for reference. The ILI panel does not enter any reported result.

Feature	Description / source
Population	ACS
under_20	age-bracket fraction
over_65	age-bracket fraction
White	race/ethnicity fraction
Black	race/ethnicity fraction
medicaid_coverage	fraction

Appendix

Table 6: Static per-state covariates (

\mathbf{s}_{i}\in\mathbb{R}^{17}

Feature	Type	COVID	Flu
cases_reported	per 100k	✓	✓
hosp_reported	per 100k	✓	✓
death_reported	per 100k	✓
incidence	per 100k	✓	✓
Rt	Cori estimator	✓	✓
policy	intensity scalar	✓	✓

Appendix

Table 7: Dynamic per-state covariates per dataset (

\mathbf{X}_{i}^{\,t-T:t}\in\mathbb{R}^{14\times T}

). Each dataset sums to 14 channels: 9 observed + 5 latent.

Dataset	Dimension	Levels	Description
Real US COVID-19 : 13-D OxCGRT indicator panel, ordinal NPI levels (increasing stringency)
	C1M_School closing	$\{0,1,2,3\}$	open / recommend / require some / require all
	C2M_Workplace closing	$\{0,1,2,3\}$	open / recommend / require some / require all
	C3M_Cancel public events	$\{0,1,2\}$	none / recommend / require
	C4M_Restrictions on gatherings	$\{0,1,2,3,4\}$	no limit $\to$ $\leq$ 10 persons
	C6M_Stay at home requirements	$\{0,1,2,3\}$	none $\to$ minimal-exception lockdown

Appendix

Table 8: Action-space specification across the three datasets. The agent selects a named tier; the world model receives the corresponding scalar intensity. COVID-19 SEIR tiers map uniformly to

\{0.00,0.25,0.50,0.75,1.00\}

; Influenza SEIR intensities are listed in-line.

Skill	Function	Anchoring authority
tier1_classify_pandemic_interval	Classifies the six WHO pandemic intervals (Investigation, Recognition, Initiation, Acceleration, Deceleration, Preparation) from $(R_{t},R_{t}\text{ CI},\text{case trend},\text{ICU},\text{cluster detected})$ ; returns max_npi_tier_permitted , which gates downstream NPI choice.	WHO pandemic-interval framework
tier1_vaccine_allocation	Enforces the priority ordering 1A $\to$ 1B $\to$ 1C $\to$ 2 $\to$ 3; flags equity_score $<0.8$ for human review; prohibits price-, insurance-, or market-based allocation.	ACIP (US) / WHO SAGE
tier1_check_npi_proportionality	Validates a proposed NPI bundle against the admissible tier envelope; requires legal authority for lockdown and movement restrictions; enforces a 14-day step-down minimum at $R_{t}<0.9$ .	ECDC proportionality test
tier1_who_outbreak_investigation	Enforces the WHO nine-step outbreak investigation ordering; blocks quantitative analysis before step 3 and control measures before step 6; mandates IHR notification at $\geq 5$ confirmed cases.	WHO IHR ( World Health Organization, 2008 )

Appendix

Table 9: Protocol Constraints (

\mathcal{T}_{1}

). Deterministic predicates returning

\{0,1\}

admissibility; computed once per decision step.

Skill	Method (source)	Refinable parameters
tier2_estimate_rt	Bayesian renewal-equation estimator for $R_{t}$ ( Cori et al., 2013 ) .	serial-interval mean ( $4.7$ d), s.d. ( $2.9$ d), window $\tau$ (1 w)
tier2_verity_expected_mortality	Age-stratified IFR with case-ascertainment correction.	ascertainment ratio (0.3)
tier2_compute_stringency_index	Nine-indicator OxCGRT normalisation to $[0,100]$ ( Hale et al., 2021 ) .	—
tier2_lookup_svi	ATSDR Social Vulnerability Index $\to$ quartile and equity directives ( Agency for Toxic Substances and Disease Registry, 2022 ) .	—
tier2_compute_growth_rate	Log-linear OLS for exponential growth rate $r$ , doubling/halving time, and fit $R^{2}$ .	stable band ( $\pm 5\%$ /week)
tier2_r_from_growth_rate	Wallinga–Lipsitch closed-form $R$ from growth rate $r$ and generation interval.	—

Appendix

Table 10: Surveillance Skills (

\mathcal{T}_{2}

). Analytical estimators returning quantitative summaries; refinable parameters may be adjusted by Adaptive Lessons within per-skill bounds.

Skill	Signature	Description
tier3_retrieve_lessons	$(\tau,k)$	Returns the top- $k$ lessons ranked by tag overlap with the current context tags $\tau$ , weighted by lesson confidence $\rho_{\ell}$ .
tier3_apply_lesson	$(\ell,\phi)$	Merges the parameter delta $\boldsymbol{\delta}_{\ell}$ into base parameters $\phi$ , clamping to per-skill bounds $B_{\kappa}$ ; only keys in the refinable-parameter set are modulated.

Appendix

Table 11: Adaptive Lessons (

\mathcal{T}_{3}

). Retrieve-and-apply operators over the lesson bank; structure is fixed, content evolves through after-action analysis.

Figure 3: Hospitalisation reduction (mean per-100k) across three datasets and three skill configurations.

Variant	In-Dist. Peak-MAE	OOD Peak-MAE
EpiWorld (full)	0.339	0.414
$-$ dual encoder (CDE-only)	0.385	0.433
$-$ dual encoder (GRU-only, 8w)	0.419	0.540
$-$ long lookback (GRU 3w)	0.587	0.720
$-$ DelayMixer (1-D causal mixer)	0.399	0.426
$-$ multi-stream SIR-coupled head	0.554	0.734

Appendix

Table 12: World-model component ablation on COVID-19 SEIR (Peak-MAE, log1p per-100k; lower better). Each row removes one component from the full model.

Held-out state	Base	+ Real-US lessons	$\Delta$
KS	12.51	12.64	$+1.0\%$
MN	6.87	6.78	$-1.4\%$
NV	57.74	74.79	$+29.5\%$
SC	11.12	10.29	$-7.5\%$
WA	88.24	93.28	$+5.7\%$
Pooled	35.30	39.56	$\mathbf{+12.1\%}$

Appendix

Table 13: Cross-domain lesson transfer is negative. Lessons mined under the real-US world model are applied to the synthetic SEIR world model on five held-out OOD states. Mean hospitalisation per

100

k; lower better.

State	A (Base)	D (+ Skills)	F (+ Skill Evo.)
Arizona	23.34	25.03	14.24
Indiana	20.01	18.84	10.86
Wisconsin	19.16	17.06	10.44
Pennsylvania	15.49	15.62	9.27
Tennessee	15.24	16.18	10.03
Michigan	14.07	14.63	6.28

Appendix

Table 14: Per-state hospitalisation MAE (per 100k) on the real-US Omicron OOD split. Skill Evolution improves

15

20

states and gives its largest reductions on high-baseline-MAE Midwest states (AZ, IN, WI). On the five lowest-baseline-MAE states (FL, NJ, CA, IL, WA) the multiplicative correction over-applies and MAE regresses; this regime-dependence is discussed in Appendix I .

	Real US	COVID	Flu
Method	Act./Ex.	Act./Ex.	Act./Ex.
EpiPolicy-RL	43.1/22.3	92.3/10.4	38.1/17.2
ADIOS (GGGP)	31.2/33.4	90.1/15.5	27.6/25.8
EpidRLearn	31.9/33.7	90.4/15.8	28.2/26.0
EpiWorld (Base LLM)	66.1/50.3	97.2/23.3	58.5/38.8
EpiWorld (Static Skills)	69.2/53.8	97.7/25.0	61.2/41.5

Appendix

Table 15: Policy-alignment diagnostics (activation / exact-match rate, %) for the methods in Table 2 . Reported as a plausibility check on the proposed action distribution, not as a quality metric.

Selector	Cumulative reduction	Mean intensity
Rule-based (no LLM)	$-21.4\%$	$0.485$
Base LLM	$-36.5\%$	$\approx 0.50$
$+$ Skill Evolution	$-39.3\%$	$\approx 0.50$

Appendix

Table 16: No-LLM ablation, COVID-19 SEIR state-OOD, evaluated on the independent mechanistic simulator. All three selectors are scored on an identical set of rollouts, so the values are directly comparable within this table; they are computed on a matched subset and are not on the same scale as Table 3 .

Action-Conditioned World Models

Action-conditioned world models predict how environments change under actions to support planning, control, and simulation.

See all 326 papers

Toward World Models for Epidemiology
Apr 10, 2026Zeeshan Memon, Yiqi Su, Christo Kurisummoottil Thomas +3World Models Causal World Models
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
The Epi-LLM Framework: probing LLM behavioral priors through epidemiological agent-based models
Jun 1, 2026Petra Ferencz, Ava Keeling, Tobias O'Keefe +4Human Behavior Simulation LLM Agent Evaluation
Big Data Institute, Li Ka Shing Center for Health Information and Discovery, University of Oxford, Oxford, United Kingdom · Leverhulme Centre for Demographic Science, Nuffield Department of Population Health, University of Oxford, Oxford, United Kingdom · Pandemic Sciences Institute, Nuffield Department of Medicine, University of Oxford, Oxford, United Kingdom +3
Graph World Models for Constrained Epidemic Policy Planning
Sep 28, 2026Yiqi Su, Rashed Shelim, Lingyi Wang +2Action-Conditioned World Models Graph-Based Planning
Department of Computer Science, Virginia Tech, Alexandria, VA 22305, USA · Department of Electrical and Computer Engineering, Virginia Tech, Alexandria, VA 22305, USA

Apr 10, 2026J/K move · Enter open · S save

Toward World Models for Epidemiology

Zeeshan Memon, Yiqi Su, Christo Kurisummoottil Thomas, Walid Saad+2

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

World Models Causal World Models

Open paper AskPDF arXiv

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.

Abstract

Action-Conditioned World Models

Explore similar work

Toward World Models for Epidemiology

The Epi-LLM Framework: probing LLM behavioral priors through epidemiological agent-based models

Graph World Models for Constrained Epidemic Policy Planning