Source-linked AI summary
EpiWorld: Grounding LLM Policy Agents in Epidemiological World Models
Zeeshan Memon, Yiqi Su, Kai Shu, Naren Ramakrishnan, Liang Zhao
TL;DR
Epidemic policy decisions are expressed through language, while LLMs lack grounded epidemic dynamics for testing intervention consequences. EpiWorld combines a protocol-guided LLM actor with a learned regional world model and counterfactual rollouts; across epidemic settings, it improved out-of-distribution forecasting and closed-loop policy outcomes.
Problem
Epidemic policy reasoning needs a grounded decision environment where LLMs can test textual guidance against learned policy-conditioned epidemic dynamics.
Source: §1 · pp. 1–2Method
EpiWorld couples a protocol-guided LLM policy actor and tiered skill library with an action-conditioned regional world model for counterfactual rollouts and lesson updates.
Source: §1 · pp. 2–3Results
Across all epidemic settings, EpiWorld improved out-of-distribution forecasting and closed-loop policy outcomes.
Source: §6 · p. 9Takeaways & Limitations
EpiWorld suggests LLMs can serve as public-health decision-support tools under human oversight.
Source: §6 · p. 9Evaluations are simulator-grounded and retrospective, so they cannot substitute for evaluation in a live deployment.
Source: Limitations · p. 10
Abstract
from arXiv · showhide
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.
1 Introduction
Epidemic response requires decisions about alternative interventions over evolving, interacting regions, but existing forecasting, causal, simulation, and language-model tools each leave part of this task unsupported. EpiWorld connects protocol-guided language reasoning to a learned regional world model for counterfactual evaluation and lesson updates.
1 Introduction
- Public-health policies are language-mediated artefacts, and epidemic response combines surveillance observations, forecasts, and prior decisions under uncertainty.This makes epidemic response a testbed for LLM policy reasoning grounded in quantitative decision environments.Source: §1 · p. 1
- Epidemic policy optimisation requires evaluating alternative intervention trajectories because decisions recur across interacting regions amid delayed clinical observations and policy-endogenous dynamics (Memon et al., 2026).The objective includes public-health outcomes such as hospital burden and mortality.Source: §1 · p. 1
- Existing forecasters predict outcomes under observed covariates, while causal methods struggle with co-occurring policies, delayed outcomes, and prospective multi-step decisions (Shi et al., 2022; Wan et al., 2025; Du et al., 2025; Flaxman et al., 2020).Mechanistic simulators and reinforcement-learning methods depend on predefined environments, while LLMs lack grounded epidemic dynamics to evaluate recommendation consequences (Williams et al., 2023; Chopra et al., 2024; Song et al., 2020; Mai et al., 2023).Source: §1 · pp. 1–2
- EpiWorld couples an action-conditioned regional world model with a protocol-guided LLM actor that uses protocol constraints, surveillance tools, and lessons from prior experience to propose interventions.The world model evaluates candidate policies through regional counterfactual rollouts; outcomes are analysed to update the lesson library, while recommendations remain subject to expert approval.Source: §1 · p. 2
- EpiWorld frames epidemic AI as a grounded, language-mediated closed loop that unifies forecasting, counterfactual reasoning, and prospective intervention optimisation.Its world model integrates delay-aware temporal encoding, graph-structured regional coupling, residual next-state prediction, and policy-conditioned latent modulation; its actor operationalises protocols, surveillance estimators, and after-action lessons.Source: §1 · pp. 2–3
2 Related Work
Related work spans epidemic forecasting, policy optimisation, LLM-based policy and skill methods, and world models. EpiWorld retains the LLM as policy proposer and uses a learned epidemic world model for rollouts.
Forecasters and LLM proposers
- Epidemic forecasting includes mechanistic Bayesian models, neural latent-ODE and graph-CDE extensions, and LLM trend forecasters; prior LLM approaches also propose policies in fixed mechanistic environments (Flaxman et al., 2020; Gibson et al., 2023; Qian et al., 2020; Shi et al., 2022; Wan et al., 2025; Du et al., 2025; Aoki and Ghaffarzadegan, 2026; Mao et al., 2025; Shi et al., 2026).Other LLM approaches build simulators, generate narratives over projections, or model behaviour (Datta et al., 2026a; Datta et al., 2026b; Williams et al., 2023; Hou et al., 2025; Park et al., 2023).Source: §2 · p. 3
- EpiWorld differs from approaches that fold interventions into covariate weightings or evaluate policies only inside fixed mechanistic simulators by closing the loop through a learned world model.The LLM remains the policy proposer.Source: §2 · p. 3
Policy optimisation and skill learning
- Optimal control derives policies under assumed compartmental or metapopulation dynamics, while deep RL trains in simulated epidemic environments but transfers poorly across diseases or periods (Song et al., 2020).LLM-agent research explores skill accumulation and prompt optimisation; EpiWorld applies epidemic-specific skill accumulation against a learned world model (Wang et al., 2023; Shinn et al., 2023; Xia et al., 2026; Agrawal et al., 2025; Khattab et al., 2023; Yang et al., 2023; Chopra et al., 2024; Xi et al., 2025).Source: §2 · p. 3
World models
- World models learn latent dynamics for policy learning through in-model rollouts, with prior work focused largely on robotics and games and some extensions to scientific domains (Hafner et al., 2019, 2020; Hafner et al., 2021, 2023; Micheli et al., 2022; Hansen et al., 2022; Qiao et al., 2026).EpiWorld instantiates epidemic decision-making as an action-conditioned recurrent state-space world model and uses it as the policy rollout substrate (Memon et al., 2026).Source: §2 · p. 3
3 Method
EpiWorld formulates epidemic policy optimisation as decisions across interacting regions and couples a learned world model for counterfactual rollouts with a protocol-guided LLM actor. Its model incorporates epidemic-specific temporal, intervention, and regional structure, while the actor uses a tiered skill library that can adapt through rollout lessons without changing protocol constraints or surveillance structures.
3.1 Problem Formulation
- The formulation represents interacting regions as a graph and seeks a goal-conditioned policy that maximises expected utility over a horizon under unknown policy-conditioned epidemic dynamics.Objectives may include minimising cumulative hospitalisations, reducing peak burden, intervention budgets, or equity constraints; existing approaches trade off counterfactual richness, scalability, adaptability, and controllability.Source: §3.1 · pp. 3–4
3.2 EpiWorld Framework
- EpiWorld couples a model of unknown policy-conditioned epidemic dynamics that enables counterfactual rollouts with an actor that generates interventions from predicted trajectories.The learned model replaces unknown dynamics for forecasting, counterfactual evaluation, and policy optimisation.Source: §3.2 · p. 4
3.3 Graph-Structured Epidemiological World Model
- The world model combines delay-aware temporal encoding, policy-conditioned latent modulation, and graph-structured regional coupling to model epidemic dynamics for counterfactual policy optimisation.It uses latent and observation-space supervision, with a per-channel Huber loss for robustness to reporting noise and count outliers.Source: §3.3 · pp. 4–5
3.4 Protocol-Guided LLM Policy Actor
- The actor exposes Protocol Constraints, Surveillance Skills, and Adaptive Lessons as distinct skill tiers grounded in public-health protocols, quantitative surveillance, and operational learning.Protocol Constraints define an admissible action envelope; surveillance tools provide quantitative summaries, and lessons refine surveillance estimates without expanding that envelope.Source: §3.4 · pp. 5–6
3.5 Closed-Loop Skill Adaptation
- After-action analysis of world-model rollouts adds or supersedes Adaptive Lessons while leaving Protocol Constraints and Surveillance Skill structures untouched.Lesson extraction and bank updates occur every K episodes, separately from step-by-step intervention selection, so online policy inference remains unchanged.Source: §3.5 · p. 6
4 Experiments
The experiments evaluate EpiWorld on three epidemic datasets and compare its forecasting and policy-optimisation performance with baseline families. Forecasts are assessed on 12-week rollouts, while policy optimisation is assessed on 8-week cumulative hospitalisation reduction.
4.1 Data
- The three datasets cover real US COVID-19 surveillance, a policy-conditioned COVID-19 SEIR simulator, and an Influenza SEIRV simulator.Real COVID-19 data comprise weekly state-level observations for 51 states from December 2020 to February 2022; the simulators use a 5-tier COVID-19 action and a two-dimensional influenza action, respectively.Source: §4.1 · p. 6
4.2 Baselines and Evaluation
- Forecasting baselines span Bayesian compartmental, latent ODE, and neural CDE models, each evaluated with policy-blind and policy-conditioned variants.The policy-conditioned variants ingest weekly policy intensity through the same control interface as EpiWorld.Source: §4.2 · p. 6
- Forecast quality uses MAE and Peak-MAE on 12-week autoregressive rollouts, while policy quality is cumulative hospitalisation reduction against the historical policy.Table 1 reports Peak-MAE / MAE, with lower values better; Table 2 reports cumulative hospitalisation reduction over 8-week rollouts, also with lower values better.Source: Table 1 · p. 7, Table 2 · p. 7, §4.2 · p. 7
- Policy-optimisation baselines include PPO/actor-critic, grammar-guided genetic programming, and DQN, trained and rolled out against the same world model and evaluation harness.The methods are EpiPolicy-RL, ADIOS, and EpidRLearn.Source: §4.2 · p. 7
- Evaluation includes temporal, state-wise, and intervention-OOD splits.Policy-alignment rates are reported separately as plausibility diagnostics, not as policy-quality metrics.Source: §4.2 · p. 7
5 Results
EpiWorld achieves the best OOD Peak-MAE across three datasets and reduces cumulative hospitalisation through closed-loop policy optimisation. The Ohio Delta case study illustrates one intervention trajectory and its projected effect.
5.1 World Model Performance
- ∼6%, ∼4%, and ∼8% OOD Peak-MAE gains over EARTH + Policy on COVID-19 SEIR, Influenza SEIR, and Real US (Omicron), respectively.EpiWorld outperforms all baselines on OOD Peak-MAE across the three datasets.Source: §5.1 · p. 7
- MechBayes + Policy roughly halves OOD MAE, while latent-ODE models trail uniformly; EpiWorld combines graph coupling, policy-conditioned modulation, and delay-aware decoding for stronger transfer to unseen dynamics.The Peak-MAE advantage over EARTH + Policy on COVID-19 SEIR state-OOD is significant across 50 held-out replicates (paired Wilcoxon p = 0.039).Source: §5.1 · p. 8
- Spearman ρ = 0.83 links ensemble spread to realised forecast error, but nominal 95% intervals cover only ∼68% of outcomes overall, so spread is a relative reliability signal, not a calibrated interval.Coverage is 73% on the hospitalisation channel.Source: §5.1 · p. 8
5.2 Policy Optimisation
- ∼58% further reduction in cumulative hospitalisation versus RL baselines across three datasets; on COVID-19 Real US, reductions are −59.3% for EpiWorld and −37.5% for EpidRLearn.The Base LLM already improves on every RL baseline before after-action signals are added.Source: §5.2 · p. 8
- +73% independent-ground-truth margin over EpidRLearn versus +58% in-model; learned-world-model scoring understates rather than inflates the closed-loop policy advantage.Across six methods, independent-simulator results remain within ×0.97–1.16 of in-model values, preserving their ordering.Source: §5.2 · p. 8, Table 4 · p. 9
- The full Skill Evolution cycle reduces hospitalisation by ∼16% on average across six backbones; static skills leave Base-LLM performance roughly unchanged, while lessons contribute the largest gain.GPT-4o reaches −30% and Qwen2.5-7B reaches −22%.Source: §5.2 · p. 9, Table 4 · p. 9
- At matched intervention intensity, a fixed severity-proportional rule recovers only 59% of the Base LLM reduction; across 2,640 real-US decisions, only 2 percentage points are unsafe departures surfaced for human approval.The fixed rule's weaker reduction indicates that timing and targeting, rather than stronger measures, account for the gain.Source: §5.2 · p. 9
- Adaptive lessons reduce cumulative hospitalisation by a further 9.0% over Base LLM on an independent simulator, but Omicron-derived lessons worsen first-wave outcomes by 12.1%.Lessons generalise across states within an infection regime but are calibrated to that regime; correction can also regress per-state MAE where baseline error is already low.Source: §5.2 · p. 9
6 Conclusion
EpiWorld grounds LLM policy agents in a learned epidemiological world model and protocol-guided skills for interpretable policies evaluated through counterfactuals. Across epidemic settings, it improves out-of-distribution forecasting and closed-loop outcomes, suggesting a role in public-health decision support under human oversight.
6 Conclusion
- EpiWorld grounds LLM policy agents in a learned epidemiological world model and protocol-guided skills to generate interpretable policies through counterfactual evaluation.The framework is presented as closed-loop.Source: §6 · p. 9
- Across epidemic settings, EpiWorld improves out-of-distribution forecasting and closed-loop policy outcomes.Source: §6 · p. 9
- The reported findings suggest LLMs can serve as public-health decision-support tools under human oversight.Source: §6 · p. 9
Limitations
EpiWorld’s evidence remains retrospective and simulator-grounded, and operational use would require further validation and human oversight.
Limitations
- Evaluations use retrospective forecasts and counterfactual rollouts in learned dynamics, not prospective live deployment.Independent mechanistic-simulator re-execution preserves relative policy rankings, but both evaluations remain simulator-grounded.Source: Limitations · p. 10
- The model omits behavioural fatigue, surveillance heterogeneity, and political and economic constraints outside its goal space and training support.Source: Limitations · p. 10
- Counterfactuals involving action combinations absent from historical data rely on extrapolation, and lessons do not transfer across infection regimes.Adaptive Lessons must be mined separately for each setting.Source: Limitations · p. 10
- The five-member ensemble ranks forecast reliability well but is under-dispersed, so its intervals are not calibrated.Operational use would require dedicated validation studies, equity audits, and continuous human oversight.Source: Limitations · p. 10
A Datasets and Splits
The appendix specifies real and simulated epidemic datasets, their out-of-distribution splits, forecasting and policy baselines, and shared feature definitions.
A Datasets and Splits
- Real US COVID-19 covers 51 states from December 2020 to February 2022; training uses Alpha and Delta, with the Omicron wave held out for temporal-OOD evaluation.Surveillance combines hospitalisations from HHS Protect and cases and deaths from JHU CSSE.Source: A Datasets and Splits · pp. 12–13
- COVID-19 SEIR and Influenza SEIR use simulated state-level trajectories, with ten states held out for state-OOD tests and influenza also testing unseen policy configurations.The COVID-19 action grid has five intensity tiers; influenza uses advisory and vaccination tiers.Source: A Datasets and Splits · p. 13
- The legacy ILI panel serves only as an observation-model sanity check and contributes no reported result.Source: A Datasets and Splits · p. 13, Table 5 · p. 14
- Targets are three log(1+x)-transformed observables per state and week, with hospitalisations per 100k as the primary policy target and eight-week rollouts.The appendix also describes rate conversion, clipping, interpolation of missing weeks, and transformations of targets and mobility flows.Source: A Datasets and Splits · p. 13
B.1 World-Model Forecasting Baselines
- Forecasting baselines include compartmental Gaussian processes, Bayesian SEIR and renewal models, latent ODEs, and a neural CDE, with policy-conditioned variants where specified.The appendix documents adaptations of these methods for the EpiWorld evaluation harness.Source: §B.1 · pp. 13–14
- EARTH is also adapted with a shorter lookback and a policy encoder whose output modulates its CDE vector field.Source: §B.1 · p. 14
B.2 Policy-Optimisation Baselines
- Policy baselines comprise PPO/actor-critic, grammar-guided genetic programming, and DQN approaches, each adapted to run against the EpiWorld world model.The baselines are evaluated using the shared rollout harness and held-out splits.Source: §B.2 · pp. 14–15
C Feature Specification
- Static features comprise 17 per-state indicators shared across datasets, while dynamic features total 14 channels, combining observed and simulated latent channels as applicable.Source: C Feature Specification · p. 15, Table 6 · p. 15, Table 7 · p. 15
D Action-Space Specification
The appendix defines dataset-specific action interfaces and points to the tiered skill catalogue and adaptive-lesson schema.
D Action-Space Specification
- The three datasets use different action interfaces: a 13-dimensional OxCGRT panel, a five-tier COVID-19 intensity scalar, and a two-dimensional influenza advisory-by-vaccination interface.OxCGRT actions are grouped into four envelopes for protocol-constraint checks.Source: D Action-Space Specification · p. 15, Table 8 · p. 16
E Skill Catalogue
- The skill catalogue lists each tier’s skills, while the lesson schema and examples are provided in Appendix G.Source: E Skill Catalogue · p. 15
G Lesson Schema, Tag Vocabulary, and Examples
- Adaptive Lessons record observations, the triggering pattern, and an action-level recommendation, with retirement metadata for superseded lessons.Recommendations are for human inspection.Source: G Lesson Schema, Tag Vocabulary, and… · pp. 15–16
- Lessons link applicable tags to a target skill and bounded parameter shift, and record regime, confidence, sign agreement, and at least two source states.Source: G Lesson Schema, Tag Vocabulary, and… · p. 15
- Protocol Constraints (T1) are deterministic binary admissibility predicates computed once per decision step.Source: Table 9 · p. 16
H Experimental Details
EpiWorld combines a policy-conditioned regional world model with a tiered public-health skill library and an LLM actor. Its experiments use structured action proposals, rollout-based lesson updates, and OOD forecasting and policy-evaluation metrics.
H Experimental Details
- A GRU and Neural-CDE encode temporal history, policy actions modulate latent states, and an adjacency-masked Graph Transformer models regional coupling before residual next-state prediction.Training combines latent-prediction and observation-reconstruction losses.Source: H Experimental Details · p. 17
- World models are trained for 5,000 steps with AdamW, batch size 16, and a learning rate of 3×10−4; the best in-distribution validation-loss checkpoint is used downstream.Training uses log(1+x)-transformed, train-split-standardised targets.Source: H Experimental Details · p. 17
- Each LLM backbone returns structured-JSON actions, and top-five tag-matched lessons are injected; lessons update after rollout batches while protocol and surveillance structures remain fixed.Closed-loop rollouts run for eight weeks.Source: H Experimental Details · p. 17
- Evaluation reports OOD Peak-MAE and mean trajectory MAE for forecasting, plus cumulative hospitalisation per 100k over an eight-week post-intervention horizon for policy performance.Policy outcomes are averaged over held-out states.Source: H Experimental Details · p. 17
- The skill library separates non-overridable protocol constraints from quantitative surveillance estimators and adaptive lessons whose content can evolve.Protocol constraints are derived from WHO, CDC, and ECDC guidelines.Source: Table 10 · p. 17, Table 11 · p. 17, H Experimental Details · p. 18
I Additional Experiments
Additional experiments report skill-cycle and world-model ablations, cross-domain lesson transfer, and per-state variation. Policy-alignment diagnostics compare proposed actions with historical actions but are presented as plausibility checks, not quality measures.
I Additional Experiments
- Static skills shift mean hospitalisation by at most 2%, while Skill Evolution yields 21–31% mean reduction on Influenza SEIR despite its different NPI grid.Figure 3 reports mean per-100k hospitalisation reductions across three datasets and three skill configurations.Source: I Additional Experiments · p. 18, Figure 3 · p. 20
- Removing any single world-model component degrades in-distribution and OOD Peak-MAE; residual prediction is load-bearing, while the multi-stream SIR-coupled head contributes most to OOD generalisation.The component ablation is on COVID-19 SEIR.Source: I Additional Experiments · p. 19, Table 12 · p. 19
- On influenza, removing Neural-CDE raises overall MAE from 0.487 to 0.604, while DelayMixer gives the best Peak-MAE at 0.441 versus 0.463 without it.The ablation uses two seeds per configuration on influenza state-OOD.Source: I Additional Experiments · p. 19
- Transferred real-US lessons worsen mean hospitalisation by +12.1% on five held-out synthetic-SEIR OOD states, where their shrinkage correction amplifies under-prediction.Tier-3 lessons are not universally portable and should be mined under the target world model’s training distribution.Source: I Additional Experiments · p. 20, Table 13 · p. 20
- Skill Evolution improves 15 of 20 Omicron OOD states, with largest reductions on high-baseline-MAE Midwest states; on five low-baseline-MAE states, the correction over-applies and MAE regresses.The low-error states include FL, NJ, CA, IL, and WA.Source: Table 14 · p. 20
J Policy-Alignment Diagnostics
- On COVID-19 SEIR, 90.1–98.9% activation alignment coexists with 8.2–22.0% hospitalisation reduction, while exact-match stays at or below 60%, making alignment a plausibility check rather than a quality metric.Activation agreement concerns which measures are switched on; exact-match concerns the full action bundle.Source: J Policy-Alignment Diagnostics · pp. 20–21, Table 15 · p. 21
K No-LLM Rule-Based Selector
A matched rule-based selector isolates the LLM’s contribution, while the lesson module pools rollout outcomes into stable, retrievable adjustments to surveillance-skill parameters. Protocol constraints remain unchanged.
K No-LLM Rule-Based Selector
- On matched independent-simulator rollouts, removing the LLM roughly halves reduction at essentially unchanged mean intensity (0.485 against ≈0.50); the Base LLM supplies most of the gap, and Skill Evolution adds 2.8 points.The comparison uses a matched subset and is not on the same scale as Table 3.Source: Table 16 · p. 21, K No-LLM Rule-Based Selector · p. 21
L Teacher Module and Lesson Retrieval
- Lessons pool episodes by state profile, target skill, and policy regime, and require at least two distinct states to avoid single-state idiosyncrasies; only Tier-2 parameters are adjusted.Protocol Constraints are never modified.Source: L Teacher Module and Lesson Retrieval · p. 21
- Lesson confidence reflects observation count, sign agreement, and contributing states; same-sign refinements blend, while opposite-sign updates require explicit supersession and preserve the audit trail.The confidence merge is an evidence-weighted mean.Source: L Teacher Module and Lesson Retrieval · p. 21
- At each decision, the actor ranks lessons by tag overlap multiplied by confidence, applies the top five with at least two matching tags, and clips deltas to [−0.8, 0.8].The bounded confidence update and supersession policy are described as stabilising continued adaptation.Source: L Teacher Module and Lesson Retrieval · p. 22