Exploring optimal control of epidemic spread using reinforcement learning

The study addresses how to balance public health mitigation of a fast-spreading pandemic (such as COVID-19) with economic impacts. Traditional interventions (e.g., age-based or fixed n-work-m-lockdown schedules) may not always be applicable or optimal across diseases and contexts. The authors pose key questions: whether long lockdowns are the only solution, when to place lockdowns if cases don’t improve, how to handle resurgence, and how to balance epidemiological control with economic stability. They frame the problem as an optimization task: maximize economic outcomes while minimizing disease spread, solvable via reinforcement learning (RL). Recent advances in deep reinforcement learning (DRL) motivate exploring agent-derived policies that may surpass hand-crafted strategies.

The paper situates its work within compartmental epidemic modeling (e.g., SEIR) and notes widespread use of ODE-based models in epidemiology, which are deterministic and may lack stochasticity and super-spreading phenomena. It references mitigation strategies (age-based lockdowns, cyclic n-work-m-lockdown) and highlights limitations for diseases affecting all ages or during severe waves. The authors survey RL/DRL progress (Q-learning, Double Q-learning, DQN/DDQN, LSTMs) and its successes in complex decision-making (games such as Go, DOTA-2, poker, and dialog systems), motivating DRL for policy discovery in pandemics. They also compare their virtual environment’s reproduction number behavior with empirical COVID-19 estimates in China, supporting realism at certain densities.

The authors implement a stochastic, agent-based virtual environment grounded in an SEIR framework to train a DRL agent to choose movement restrictions. Key elements: - Environment design: 2D grid with 10,000 individuals who move randomly each day (default 15 steps). Individuals transition through SEIR with stochastic durations and outcomes. Exposure (E) lasts 1–2 days before becoming infectious (I). Infectious duration is 21–27 days, after which 80% recover and ~20% die. Recovered are immune and contribute economically. - Movement restrictions (actions): Level-0 (no restriction, full movement), Level-1 (−25% movement; akin to social distancing), Level-2 (−75% movement; akin to lockdown). Restrictions reduce spread but also reduce economic output. - State representation (observations): Seven features provided daily to the agent: Active cases (%), Newly infected (%), Cured (cumulative %), Deaths (cumulative %), Reproduction rate R (estimated in environment), Economy ratio (% of baseline daily contribution), and Current movement restriction level. - Economy model: Each moving individual contributes a random value in [0.8, 1]. Infectious individuals, and those who have died, contribute 0. Increasing deaths and active infections reduce total economic output. - Workflow: Each day, individuals move and potentially contact infectious persons (leading to exposure). Exposed and infectious individuals still move; only infectious can transmit. Deaths remove individuals from the environment. The episode ends when there are zero exposed and zero infectious individuals. - Reward function: R(s) = E_t × exp(−r × A_t) − s × D_t, where E_t is current economy ratio, A_t is active cases (%), D_t is cumulative death (%). Tuning parameters: r = 8 controls decay with active cases; s = 5 penalizes deaths. At low active cases, higher economic stages yield higher reward; beyond a critical active case threshold (~0.82%), rewards across economic levels become nearly equal, shifting priority to disease suppression. - DRL agent and training: A Double Deep Q-Network (DDQN) with memory is used. Network: three bidirectional LSTM layers (128, 64, 64 units per direction) followed by dense layers (128, 64, 32, and 3 outputs for the three actions). The agent uses sequences of past daily reports; memory lengths studied: 7, 15, 30, 45, 60 days. Discount factor γ = 0.9, loss = MSE. Epsilon-greedy exploration decays from 1 with a floor of 0.1; training for 7000 episodes. Additional implementation details: Python/Keras/TensorFlow; environment algorithm has quadratic complexity, constraining population to 10,000. - Model selection: Agents M7, M15, M30, M45, M60 denote memory lengths. M7/M15/M45/M60 initialized from pretrained M30 weights to speed convergence. Performance assessed by reward, infections, deaths, and economic metrics across identical environment scenarios.

- Environment realism and ODE comparison: With densities 0.02–0.03, the virtual environment’s epidemic curves closely match an SEIR ODE model (β=0.12, α=1, γ=1/27, μ=0.009; R0≈3.24). Estimated R0 in the virtual environment: mean ≈3.2±0.30 (density 0.02) and ≈3.4±0.23 (density 0.03), aligning with reported COVID-19 R0 (mean ≈3.28, median ≈2.79). Herd immunity thresholds from the environment are comparable to those implied by ODE. - Baseline interventions in simulation: • No lockdown (Level-0): >80% infected; ~20% die; economy ratio drops to ~0.20 due to population loss and non-contributing infectious individuals. Reproduction rate typically 2–5, with surges after ~160 days. • Social distancing (Level-1): Total infections reduced by ~20% and deaths reduced by ~10% compared to no lockdown; economy ratio declines by ~0.2. • Full lockdown (Level-2): Disease eliminated after ~60–63 days; economy ratio falls below 0.2. Authors note complete elimination via prolonged lockdown is unrealistic in real-world settings. - DRL agent comparison (memory lengths): • Convergence: Pre-initialized models (M7, M15, M45, M60) reached high rewards by ~1750–3000 episodes; M30 required ~6000 episodes from scratch. • Action patterns: Best agents predominantly used Level-0 and Level-2; Level-1 was rarely chosen by the top model (M30). • Outcomes: M30 yielded minimal infections and deaths due to stricter, well-timed lockdowns, and mitigated the disease fastest. It achieved the best average per-day reward and economy, despite using more Level-2 actions, because faster mitigation allowed earlier economic normalization. Aggregate rankings favored M30 across reward/infection/economy. - Learned policy structure (M30): • Early phase: one or more long lockdowns (~30–40 days) to suppress the first surge. • Later phases: cyclic short lockdowns (1–3 day cycles) interspersed with 10–20 day lockdowns to prevent resurgences and push R<1 while maintaining some economic activity. • Decision drivers: Lockdown imposed when both active cases and R are high; removed when R<1. The agent effectively uses the heuristic Increase ≈ ActiveCases × R0 to gauge imminent growth risk. - Comparison to fixed cyclic policies: A 7-work/7-lockdown policy alone leads to 34.5% infected. Adding an initial 40-day full lockdown before 7–7 reduces total infection to 11.5%, yet resurgence persists. The DRL agent’s adaptive policy more effectively flattens the curve than fixed schedules. - Reward threshold insight: At ~0.82% active cases, rewards across economic levels equalize, indicating a critical point where public health suppression should dominate over economy in the reward optimization.

Findings show a DRL agent can learn adaptive, state-dependent intervention policies that outperform simple fixed-cycle strategies in simultaneously mitigating spread and preserving economic activity. The agent learns to combine early decisive lockdowns with later cyclic and shorter lockdowns to manage resurgences, guided by active cases and reproduction rate dynamics. This approach formalizes trade-offs via a reward function that shifts priority from economy to suppression beyond a critical active-case threshold, aligning with capacity constraints (e.g., healthcare limits). The virtual environment’s alignment with ODE models and empirical R0 estimates supports plausibility. The results suggest that data-driven, RL-derived policies can complement public health decision-making by optimizing timing and intensity of non-pharmaceutical interventions under uncertainty.

This work introduces a stochastic SEIR-based virtual environment with an embedded economy and trains a memory-based DDQN agent to optimize pandemic control policies. Contributions include: (1) a virtual environment capturing contact-driven transmission, stochastic disease progression, and economic impacts; (2) systematic evaluation of movement-restriction strategies (none, social distancing, lockdown) and their epidemiological and economic effects; (3) discovery and analysis of adaptive, RL-derived policies that employ early long lockdowns followed by cyclic and short lockdowns to control resurgences while balancing economy. Future research directions: integrate more realistic mobility/contact networks and heterogeneity (age, comorbidities, occupations), incorporate vaccination, testing, and isolation policies, calibrate to region-specific data and healthcare capacities, explore multi-objective RL or constrained RL formulations, quantify robustness to parameter uncertainty, and assess transferability and interpretability of learned policies for real-world deployment.

- Simulation-only environment: Although validated against ODE dynamics and empirical R0 ranges, the virtual model abstracts real-world complexities (behavioral compliance, demographics, contact networks, policy enforcement, healthcare capacity constraints) and may not capture all heterogeneities. - Computational constraints: Quadratic-time environment limited experiments to 10,000 individuals and fixed daily movement, potentially affecting scalability and realism. - Reward shaping and parameter choices: Results depend on r and s values; different settings or alternative objectives could change learned policies. - Exposure not observed by agent: The agent receives no exposed-state data, mirroring practical observability but potentially limiting foresight. - Policy practicality: Long or repeated lockdowns that fully halt spread (e.g., ~60+ days) may be infeasible in practice; compliance and socio-economic factors were not modeled. - Training procedure: Some agents were initialized from M30 weights, which may influence comparative convergence; generalization to unseen scenarios beyond those simulated remains to be tested.