🧬 Learning Targeted Quarantine and Vaccination Policies on Temporal Contact Networks using Graph Neural Network-Driven Reinforcement Learning

A node-scoring GCN-PPO architecture that learns targeted quarantine and vaccination policies on real-world temporal contact networks, achieving up to 47% infection reduction over uncontrolled epidemics.

📄 Report · 📊 Results · 🏗️ Architecture · 🚀 Quick Start

---

📌 Overview

Controlling epidemic outbreaks in structured populations—hospitals, schools, workplaces—requires rapid, targeted interventions. Classical heuristics like degree or betweenness centrality either ignore the evolving nature of contact patterns or are computationally prohibitive at scale.

EpiGraphRL addresses this with a deep reinforcement learning agent that operates directly on temporal contact graphs. A shared Graph Convolutional Network (GCN) encoder produces per-node embeddings, which are scored by independent policy heads to select optimal quarantine and vaccination targets each day — all under a fixed intervention budget.

Key Contributions

• Node-Scoring GCN Architecture — Per-node MLPs applied to GCN embeddings produce intervention scores that naturally preserve node identity, resolving the information-loss bottleneck of flat GCN-RL pipelines
• Hybrid Dual-Head Policy — Separate quarantine and vaccination scoring heads learn distinct targeting strategies (isolate infected vs. immunize hubs), outperforming single-action baselines by 37–41%
• Two-Phase Training Pipeline — Behavioral cloning from a degree-centrality expert provides a warm start, followed by PPO fine-tuning that discovers strategies surpassing the expert itself
• Rigorous Evaluation — Statistical significance testing (Mann-Whitney U, Cohen's d), DQN baselines, four ablation experiments, β-sensitivity analysis, and computational benchmarks across two real-world datasets

---

🏗️ Architecture

System-Level Pipeline

The agent observes the temporal contact network and disease states, encodes them through a 2-layer GCN, scores each node for quarantine and vaccination independently, and feeds the selected actions back to the SEIQVR environment.

Node-Scoring GCN — Detailed Architecture

The core innovation: instead of flattening GCN embeddings into a vector (which destroys node identity), each node's 16-dimensional embedding is independently scored by per-node MLPs. This means logit[i] is computed exclusively from node i's embedding.

---

🔬 SEIQVR Compartmental Model

An 8-state Markov-compliant compartmental model. Sub-states (I₁/I₂, Q₁/Q₂) encode time-in-compartment, eliminating hidden timers and ensuring the environment satisfies the Markov property required by PPO.

Transition	Description	Mechanism
S → E	Susceptible → Exposed	Contact with I₁/I₂ neighbor (prob. β per contact)
E → I₁ → I₂ → R	Disease progression	Automatic, 1 step each
S/E/I → Q₁ → Q₂ → S	Quarantine cycle	Agent action → 2-step isolation → release
S/E/I → V	Vaccination	Agent action → permanent immunity

---

📊 Key Results

Hospital Contact Network (N=75, T=4 days)

Strategy	Avg Infected (Ever)	Avg Reward	Inference (ms/step)
No Intervention	12.3 ± 10.9	-213.4	0.002
Random	9.7 ± 9.0	-207.3	0.038
Degree Centrality	5.3 ± 4.8	-155.8	0.025
Betweenness Centrality	5.5 ± 5.3	-161.8	10.05
DQN Baseline	7.5 ± 6.1	-261.3	0.870
GNN-PPO (Ours)	5.9 ± 6.3	-157.9	3.32

Primary School Contact Network (N=242, T=8 blocks)

Strategy	Avg Infected (Ever)	Avg Reward	Inference (ms/step)
No Intervention	11.8 ± 6.9	-374.0	0.002
Random	11.1 ± 7.2	-393.3	0.049
Degree Centrality	10.5 ± 6.5	-377.1	0.049
Betweenness Centrality	10.7 ± 6.5	-384.5	132.09
DQN Baseline	6.9 ± 5.2	-388.4	0.942
GNN-PPO (Ours)	6.2 ± 3.4	-276.4	3.40

Key Takeaways:

• 📉 47% infection reduction vs. no intervention on the school network (p < 10⁻⁵)
• ⚡ 39× faster inference than betweenness centrality on the school network
• 🏆 Outperforms DQN on both datasets — 21% fewer infections on the hospital network
• 🎯 Matches degree/betweenness heuristics on hospital while being fully learned, not hand-designed

Ablation Studies

Experiment	Hospital (Avg Infected)	School (Avg Infected)	Key Finding
Full GNN-PPO	5.9	6.2	Best overall performance
No Behavioral Cloning	10.1	11.3	BC essential — cold-start exploration fails without it
No Reward Shaping	6.3	6.9	Shaping helps convergence speed, but agent is robust without it
Quarantine-Only (3Q, 0V)	10.0	10.4	Hybrid outperforms by 37–41%
Vaccination-Only (0Q, 3V)	9.4	10.4	Hybrid outperforms by 37–40%

---

📁 Repository Structure

EpiGraphRL/
│
├── README.md
├── LICENSE
├── requirements.txt
│
├── docs/
│   ├── Architecture Diagram.png           # System-level pipeline
│   ├── Node Scoring GCN.png               # Detailed GCN architecture
│   └── SEIQVR_MODEL.png                   # Compartmental model diagram
│
├── report/
│   └── Report.pdf                         # Full project report
│
└── notebooks/
    ├── hospital/
    │   ├── Hosptial Dataset_Implementation.ipynb          # Main GNN-PPO experiment
    │   ├── hospital_dataset_dqn-baseline.ipynb            # DQN baseline
    │   ├── hospital_dataset-without-behavioural-cloning.ipynb  # No-BC ablation
    │   └── hospital_dataset_without-reward-shaping.ipynb  # No-Shaping ablation
    │
    └── school/
        ├── school-dataset_Implementation.ipynb            # Main GNN-PPO experiment
        ├── school-net-dqn-baseline.ipynb                  # DQN baseline
        ├── school-dataset-no-bc.ipynb                     # No-BC ablation
        └── school-dataset-no shaping.ipynb                # No-Shaping ablation

---

🚀 Quick Start

Prerequisites

pip install -r requirements.txt

Run on Kaggle

All experiments are designed to run on Kaggle with GPU acceleration:

1. Upload any notebook from notebooks/ to Kaggle
2. Add the required dataset:
   • Hospital: SocioPatterns Hospital Ward
   • School: SocioPatterns Primary School
3. Enable GPU accelerator in notebook settings
4. Run all cells — training takes ~15–30 minutes per notebook

---

🧪 Methodology

Training Pipeline

Phase	Description	Details
1	Data Ingestion	Parse SocioPatterns TSV into temporal adjacency matrices (1 per time block)
2	Network Analysis	Compute degree, betweenness, clustering coefficient per block
3	Environment	Custom SEIQVR Gymnasium env with budget-constrained interventions
4	Behavioral Cloning	Imitate degree-centrality heuristic for 500 episodes (warm start)
5	PPO Fine-Tuning	3000–5000 episodes, clipped surrogate objective, linear LR decay
6	Evaluation	100 seeded episodes per strategy, 5 baselines + DQN
7	Analysis	Mann-Whitney U tests, Cohen's d, β-sensitivity, inference timing

Datasets

Dataset	Nodes	Roles	Duration	Time Blocks	Source
Hospital Ward	75	MED, NUR, PAT, ADM	~4 days	4	SocioPatterns
Primary School	242	TCH, STU	~2 days	8	SocioPatterns

Key Hyperparameters

Parameter	Value	Description
β (transmission)	0.15 / 0.04	Per-contact infection probability (hospital / school)
Budget	2Q + 1V	Interventions per time step
GCN depth	2 layers	Receptive field = 2-hop neighborhood
Embedding dim	16	Per-node representation size
PPO clip ε	0.2	Trust region constraint
GAE λ	0.95	Bias-variance tradeoff in advantage estimation
Learning rate	3×10⁻⁴ → 1.5×10⁻⁵	Linear annealing schedule
Training episodes	3000 / 5000	Hospital / School

---

🛠️ Tech Stack

Component	Technology
Deep Learning	PyTorch
RL Environment	Gymnasium
Graph Analysis	NetworkX
Numerical Computing	NumPy, SciPy
Visualization	Matplotlib, Seaborn
Execution	Kaggle Notebooks (GPU P100)

---

📜 License

This project is licensed under the MIT License — see the LICENSE file for details.

🙏 Acknowledgements

• SocioPatterns for the real-world face-to-face contact datasets
• GCN formulation: Kipf & Welling, ICLR 2017
• PPO algorithm: Schulman et al., 2017