Agentic Framework for Epidemiological Modeling

Abstract: Epidemic modeling is essential for public health planning, yet traditional approaches rely on fixed model classes that require manual redesign as pathogens, policies, and scenario assumptions evolve. We introduce EPIAGENT, an agentic framework that automatically synthesizes, calibrates, verifies, and refines epidemiological simulators by modeling disease progression as an iterative program synthesis problem. A central design choice is an explicit epidemiological flow graph intermediate representation that links scenario specifications to model structure and enables strong, modular correctness checks before code is generated. Verified flow graphs are then compiled into mechanistic models supporting interpretable parameter learning under physical and epidemiological constraints. Evaluation on epidemiological scenario case studies demonstrates that EPIAGENT captures complex growth dynamics and produces epidemiologically consistent counterfactual projections across varying vaccination and immune escape assumptions. Our results show that the agentic feedback loop prevents degeneration and significantly accelerates convergence toward valid models by mimicking professional expert workflows.

• The paper introduces EpiAgent, a framework that automates the construction of epidemiological simulators from natural language scenarios, reducing the need for manual adjustments by leveraging Large Language Models (LLMs) for program synthesis, alongside iterative verification processes.
• EpiAgent uses a retrieval-augmented, flow-graph intermediate representation ensuring epidemiological validity by incorporating constraint checks before code generation, significantly reducing projection errors from structural inconsistencies.
• Empirical results, including case studies on COVID-19 scenarios and behavioral SEIR models, demonstrated the framework's effectiveness in generating stable and interpretable projections

EpiAgent, introduced in (2602.00299), is an agentic framework that automates the construction of mechanistic epidemiological simulators from natural-language scenario specifications. The work addresses a concrete bottleneck in operational epidemic modeling: pipelines such as the COVID-19 Scenario Modeling Hub (SMH) require manual redesign of model structure whenever policies, variants, or immunity assumptions change, and the authors argue that LLM-based program synthesis can absorb this burden if—and only if—generation is constrained by explicit domain verification.

Motivation and problem formulation

The paper frames scenario modeling as follows: given a set of scenarios $\mathcal{S}$ described in natural language (interventions, vaccination policies, immune escape assumptions) and observational time series $\mathcal{D}$ over patches $p$ and weekly time steps, learn compartmental models $f^s_\theta$ with parameters $\theta \in \mathcal{C}$, where $\mathcal{C}$ encodes positivity, boundedness, conservation laws, and normalization constraints. Three desiderata are stated explicitly: empirical accuracy against observed trajectories, satisfaction of epidemiological validity constraints, and internal consistency across scenarios so that cross-scenario differences arise from modeled mechanisms rather than independent re-fitting.

The authors demonstrate empirically that naive prompting of LLMs to produce compartmental models yields frequent structural inconsistencies—invalid transitions such as direct $S \to R$ recovery, deaths arising from susceptible compartments, missing waning flows under immune-escape scenarios—which motivates their central architectural choice.

The flow graph intermediate representation

The key design element is an explicit epidemiological flow graph $G = (V, E)$ serving as an intermediate representation between scenario text and executable code. Nodes are health states; edges are parameterized transitions. Graph synthesis is retrieval-augmented: scenario descriptions are embedded with sentence-transformers, relevant passages retrieved from a curated epidemiology corpus via a FAISS index, and the augmented prompt conditions generation on established theory.

Graphs pass through an iterative verification loop enforcing hard structural constraints before any code is generated: admissible transition semantics (e.g., death only from severe states, $D$ terminal, vaccination eligibility restricted to $\{E,I,J,H\}$), plus agentic checks requiring epidemiological justification for each transition and alignment with scenario assumptions. The paper's ablation quantifies the stakes: incorrect flow graphs lead to incorrect scenario projections in over 87% of cases, and in the remaining cases flexible parameterization can mask structural errors, producing trajectories that look numerically plausible while encoding wrong causal mechanisms. This latter failure mode is arguably the more dangerous one for counterfactual policy analysis, and the paper treats it as the primary justification for pre-code structural verification.

Compilation, calibration, and multi-agent V&V

Verified graphs are compiled into differentiable PyTorch simulators by an LLM planner agent operating within a fixed code skeleton that fixes the execution interface (module structure, time indexing, output contract) while leaving model logic open. Constraint-guided prompting enforces execution constraints—notably, non-negativity is deliberately not enforced via clipping or rectification; instead, boundary violations during calibration are treated as signals that the model structure itself must be revised. This is a substantive methodological position: forced non-negativity would mask misspecification rather than expose it.

Calibration uses gradient-based optimization (Adam with ReduceLROnPlateau) under an MSE loss over levels and first differences of infections and deaths, with parameters shared across scenarios unless explicitly overridden. A multi-agent verification-and-validation layer then checks: (i) code–graph consistency, (ii) state and parameter non-negativity, (iii) population mass conservation excluding deceased flows, (iv) monotonicity of cumulative quantities, and (v) numerical stability. A validation module additionally assesses scenario fidelity—whether cross-scenario orderings match intended mechanisms—and a feedback module converts diagnostics into targeted structural revision directives.

Empirical results

Case Study I uses SMH round 17, with six counterfactual scenarios spanning three vaccination policies (no recommendation; 65+ boosters; universal boosters) crossed with low/high immune escape. The synthesized simulators reproduce the expected qualitative structure without manual encoding: within each policy pair, high-escape scenarios yield higher cumulative infections than low-escape counterparts, and increasing coverage induces monotonic reduction in burden ($\mathcal{D}$0, $\mathcal{D}$1). Calibration accuracy improves monotonically with expressiveness:

Aggregation	Model	MAE	RMSE
State-avg	Time-invariant	$\mathcal{D}$2	$\mathcal{D}$3
State-avg	Time-varying ($\mathcal{D}$4)	$\mathcal{D}$5	$\mathcal{D}$6
State-avg	Neural-incorporated	$\mathcal{D}$7	$\mathcal{D}$8
Nationwide	Time-invariant	$\mathcal{D}$9	$p$0
Nationwide	Time-varying ($p$1)	$p$2	$p$3
Nationwide	Neural-incorporated	$p$4	$p$5

The framework defaults to time-invariant parameters to avoid over-parameterization, escalating to time-varying or neural-augmented dynamics only when fit demands it—a deliberate trade-off favoring interpretability in counterfactual settings over point-forecast accuracy.

Case Study II instantiates the Data-Driven Behavioral SEIR baseline of Gozzi et al. across four heterogeneous locations using real surveillance data, mobility signals, and seasonality covariates. EpiAgent recovers the canonical model structure from a specification containing no equations or code, calibrates to weekly deaths, and produces probabilistic projections with 50%/80% intervals that capture peak timing across locations. This demonstrates generality beyond COVID-specific scenario rounds.

Ablations isolate each component. Removing RAG causes the LLM to default to minimal SEIRD-style graphs omitting scenario mechanisms (waning immunity, age-targeted vaccination), and knowledge guidance reduces average generate–verify iterations by roughly 20%. Removing the flow-graph IR entirely (direct scenario-to-code) produces structural errors—missing compartments, duplicated states, invalid transitions—that are difficult to detect at the code level. Appendix analysis further shows that unconstrained graph generation drifts toward unnecessarily complex dynamics, indicating constraints act bidirectionally against both oversimplification and over-elaboration.

Sensitivity to model capacity and scaffolding

A notable finding is that reliability emerges from the interaction of capacity, skeleton code, and constraints rather than from any single factor. High-capacity models (GPT-4.x class) infer structure with limited guidance but still produce epidemiologically invalid or unstable dynamics when constraints are removed; lightweight models require strong scaffolding to be competitive at all. Skeleton code functions as a structural prior whose importance increases as model capacity decreases. The practical implication is that constraint enforcement, not frontier-model access, is the binding requirement for trustworthy automated simulator generation.

Limitations and open questions

The authors concede two limitations directly. First, calibration uses a fixed optimization strategy; the agent cannot adapt learning rates, and bounded iterative refinement means later generations may occasionally regress due to compounding generation noise—the paper does not quantify how frequently this occurs. Second, automatic error recovery operates under a fixed retry bound to guarantee termination, which caps robustness of fully autonomous code generation. Convergence-aware stopping criteria are identified as unresolved. An additional implicit limitation worth noting: evaluation rests on two case studies (one synthetic-scenario round, one behavioral baseline), and no quantitative benchmark against human-expert-designed SMH ensemble members is reported, leaving the question of how agent-generated simulators compare to professional submissions on standard projection metrics open.

Conclusion

EpiAgent reframes epidemic simulator construction as iterative, verification-gated program synthesis, with a verified flow-graph intermediate representation as the load-bearing component. The evidence supports two claims with reasonable confidence: structural verification before code generation is necessary to prevent both overtly incorrect projections and the subtler failure of plausible-looking trajectories built on wrong causal structure, and agentic feedback loops accelerate convergence to valid structures relative to unguided generation. The framework's reliance on fixed optimization, bounded retries, and a limited evaluation footprint remain the principal open issues.