Agentic Framework for Epidemiological Modeling

The paper introduces EPIAGENT, an agentic framework that automates the synthesis, calibration, and verification of epidemiological models by utilizing an intermediate flow graph representation to ensure mechanistic consistency and accelerate the development of valid, interpretable simulators.

The Gist

Imagine you are trying to predict how a storm will move across a city. In the past, meteorologists had to manually draw the storm's path on a map, adjust the wind speeds by hand, and hope they didn't make a mistake. If the storm changed direction or the wind got stronger, they had to erase everything and start over. This is slow, prone to human error, and hard to scale.

Now, imagine having a team of expert AI assistants who can not only draw the map but also check their own work, fix their mistakes, and learn from every attempt until the prediction is perfect.

That is essentially what EPIAGENT does, but for epidemics (like the flu or COVID-19) instead of storms.

Here is a simple breakdown of how it works, using some creative analogies:

1. The Problem: The "Manual Blueprint" Bottleneck

Traditionally, building a model to predict disease spread is like building a house with a hand-drawn blueprint.

• If the architect (the scientist) wants to add a new room (like a new vaccine strategy) or change the foundation (like a new virus variant), they have to redraw the whole blueprint from scratch.
• If they make a tiny mistake in the drawing (like a wall that doesn't connect to the floor), the whole house might collapse later.
• This process is slow and relies entirely on the human expert's time and energy.

2. The Solution: EPIAGENT (The "AI Construction Crew")

EPIAGENT is an agentic framework. Think of "agentic" as a team of specialized robots that work together, rather than just one robot doing everything.

Here is the step-by-step workflow of this AI crew:

Step A: The "Architect" (Flow Graph Synthesis)

Instead of jumping straight to writing code (the construction), the AI first draws a Flow Graph.

• Analogy: Imagine a subway map. The stations are groups of people (Healthy, Sick, Recovered, Deceased), and the lines are how people move between them.
• The Magic: Before the AI writes a single line of computer code, it draws this subway map. It checks: "Wait, can you go directly from 'Healthy' to 'Recovered' without getting sick first? No! That's impossible."
• If the map is wrong, the AI throws it away and redraws it immediately. This prevents "bad blueprints" from ever becoming "bad buildings."

Step B: The "Translator" (Code Generation)

Once the subway map (Flow Graph) is verified as correct, a different AI agent translates that map into executable code.

• Analogy: This is like a translator who takes the subway map and writes the actual train schedule and control software.
• Because the map was already checked, the code is much less likely to have logical errors.

Step C: The "Safety Inspectors" (Verification & Validation)

This is the most important part. The AI doesn't just run the simulation; it has a team of inspectors that watch the simulation run in real-time.

• Inspector 1 (The Math Police): Checks if the numbers make sense. "Hey, you can't have -500 sick people! That's impossible!"
• Inspector 2 (The Logic Detective): Checks if the story makes sense. "If we introduce a vaccine, the number of sick people should go down, not up. Why did it go up?"
• Inspector 3 (The Scenario Tester): Checks if the model reacts correctly to "What if?" questions. "If we lock down the city, does the model show the virus slowing down?"

If any inspector finds a problem, they send a feedback note back to the "Architect" and "Translator" to fix the specific error. They keep doing this loop until the model is perfect.

3. Why This Matters: The "What If" Machine

The real superpower of EPIAGENT is Counterfactual Reasoning.

• Old Way: To see what happens if we vaccinate 50% of people, scientists had to manually rewrite the model, which took days.
• EPIAGENT Way: You can ask, "What if we vaccinate 50%?" or "What if a new super-virus appears?" The AI instantly adjusts its internal "subway map," runs the simulation, and gives you a projection.
• Because the AI checks its own logic, these "What If" answers are trustworthy. It doesn't just guess; it builds a logical structure that respects the laws of biology.

4. The Result: A Self-Improving Expert

The paper shows that EPIAGENT can:

1. Learn faster: It mimics how human experts work but does it in minutes, not days.
2. Avoid disasters: By checking the "blueprint" before building, it stops silly mistakes (like people dying without getting sick first) from happening.
3. Handle complexity: It can juggle many different scenarios (different ages, different vaccines, different virus strains) at the same time without getting confused.

Summary

Think of EPIAGENT as a self-correcting, AI-powered epidemiology lab.
Instead of a single scientist struggling to draw a complex map by hand, you have a team of AI specialists: one draws the map, one writes the code, and a whole squad of inspectors checks the math and logic. If something is wrong, they fix it instantly. This allows public health officials to get reliable, fast answers to critical questions like "How many hospital beds will we need next month?" or "Will a new vaccine stop the spread?"

It turns the slow, manual process of epidemic modeling into a fast, reliable, and automated engine for saving lives.

Technical Summary

1. Problem Statement

Epidemiological modeling is critical for public health planning, yet traditional approaches suffer from significant limitations:

• Manual Redesign: As pathogens, policies, and assumptions evolve (e.g., new variants, vaccination strategies), epidemiologists must manually redesign model structures, revise transition logic, and re-implement simulators. This is time-consuming and limits scalability.
• Fragility of LLMs: Naively prompting Large Language Models (LLMs) to generate compartmental models (e.g., SEIR) often leads to structural inconsistencies, invalid transitions (e.g., direct recovery without infection), and fragile behavior.
• Lack of Mechanistic Constraints: Existing automated scientific modeling systems often focus on syntactic correctness or empirical performance, lacking domain-specific mechanistic constraints (e.g., conservation of mass, non-negativity of populations) required for epidemiological validity.

The core research question is: Can agentic systems be developed to assist in epidemiological modeling by automatically synthesizing, calibrating, and verifying mechanistic simulators while adhering to strict epidemiological constraints?

2. Methodology: EPIAGENT

The authors introduce EPIAGENT, an agentic framework that treats epidemic modeling as an iterative program synthesis problem. The system operates through a multi-stage pipeline involving specialized agents, retrieval-augmented generation (RAG), and a unique intermediate representation.

Key Components:

1. Scenario Parsing & Knowledge Retrieval (RAG):

   • Natural language scenario descriptions are embedded and augmented with relevant epidemiological knowledge retrieved from a curated corpus (textbooks, prior studies) using FAISS.
   • This grounds the model construction in established theory (e.g., Brauer et al., 2012).

2. Epidemiological Flow Graph Synthesis (The Core Innovation):

   • Instead of generating code directly, the system first synthesizes an Epidemiological Flow Graph ($G = (V, E)$).
   • Nodes ($V$): Represent disease compartments (Susceptible, Exposed, Infectious, Recovered, etc.).
   • Edges ($E$): Represent admissible transitions parameterized by disease rates.
   • Verification Loop: A dedicated agent verifies the graph against hard structural constraints (e.g., no direct $S \to R$ without infection, death only from severe states, conservation of mass). Invalid graphs are rejected and regenerated with structured feedback.

3. Graph-to-Model Compilation:

   • Once a flow graph is verified, an LLM Planner Agent compiles it into executable simulator code (Python/PyTorch).
   • Skeleton Code: The system uses a fixed code skeleton to ensure a consistent execution interface, leaving only the epidemiological logic to be generated.
   • Error Recovery: An automated loop monitors compilation and runtime errors (e.g., shape mismatches, solver instability), feeding structured error traces back to the planner for correction.

4. Calibration & Training:

   • The generated simulator is calibrated against observed epidemiological time-series data using gradient-based optimization (MSE loss).
   • Parameters can be time-invariant, time-varying, or augmented with neural networks (hybrid models) depending on scenario requirements.

5. Multi-Agent Verification and Validation (V&V):

   • Verification Module: Checks mathematical correctness (non-negativity, mass conservation, numerical stability).
   • Model Reasoning Module: Explains the model structure to ensure interpretability.
   • Validation Module: Assesses scenario fidelity (e.g., do counterfactual interventions produce expected shifts in outcomes?).
   • Feedback Module: Provides actionable feedback to refine the model structure if validation fails.

3. Key Contributions

1. Retrieval-Augmented Flow Graph Synthesis: A novel approach that bridges natural language scenarios and mathematical structure by generating intermediate flow graphs that satisfy hard structural constraints, significantly improving robustness over direct code generation.
2. Graph-to-Model Compilation: A mechanism that transforms verified graphs into mechanistic simulators, providing an explicit scaffold for parameter learning.
3. Multi-Agent V&V Architecture: A specialized system where agents enforce correctness properties (mathematical validity, scenario fidelity, mechanistic interpretability), mimicking professional expert workflows.
4. Empirical Validation: Demonstrated on real-world COVID-19 scenario modeling data (MIDAS Network) and behavioral SEIR baselines, showing the framework produces epidemiologically consistent counterfactual projections.

4. Results

The framework was evaluated on two case studies:

• Case Study I (MIDAS Network): 19 scenario rounds involving complex counterfactuals (vaccination policies, immune escape).
  • Accuracy: Generated simulators achieved low Mean Absolute Error (MAE) and Root Mean Square Error (RMSE) compared to observed data.
  • Counterfactual Reasoning: The models correctly inferred qualitative behaviors (e.g., higher immune escape $\to$ higher infections; targeted boosters $\to$ reduced burden) without manual encoding.
  • Verification Impact: Removing flow-graph verification led to 87% of cases producing incorrect projections due to structural errors.
• Case Study II (Behavioral SEIR): Applied to real-world data across four locations using a Data-Driven Behavioral (DDB) model.
  • The system successfully instantiated, calibrated, and produced uncertainty-aware probabilistic projections that captured peak timing and epidemic trends.

Ablation Studies:

• Without Flow-Graph Verification: High rate of invalid structures and incorrect projections.
• Without RAG: Models defaulted to oversimplified SEIRD structures, missing scenario-specific mechanisms (e.g., waning immunity).
• Direct Scenario-to-Code: Resulted in frequent structural errors (missing compartments, invalid transitions) that were hard to detect at the code level.

5. Significance

• Automation of Expert Workflows: EPIAGENT successfully mimics the iterative, verification-heavy workflow of human epidemiologists, reducing the time required to adapt models to new scenarios.
• Ensuring Scientific Validity: By enforcing hard constraints at the graph level (before code generation) and via multi-agent V&V, the system prevents "degenerate" solutions (e.g., negative populations) that might otherwise achieve low loss but lack scientific meaning.
• Interpretability: Unlike black-box neural ODEs, EPIAGENT produces mechanistic models with explicit compartmental structures, making the reasoning behind projections transparent.
• Scalability: The framework offers a path to rapidly deploy and calibrate models for emerging pathogens or changing policy landscapes, a critical need for public health decision-making.

In conclusion, EPIAGENT represents a significant step forward in Agentic AI for Science, demonstrating that combining LLMs with structured intermediate representations and rigorous domain-specific verification can automate complex, high-stakes scientific modeling tasks.