Reward Engineering for Spatial Epidemic Simulations: A Reinforcement Learning Platform for Individual Behavioral Learning

AI summary

Overview

Research area: Reinforcement learning, agent-based epidemiological modeling, and reward engineering.

Technical level: Advanced. The paper assumes familiarity with Markov decision processes, deep RL algorithms (PPO, SAC, A2C), compartmental epidemic models, and statistical hypothesis testing.

Scope: The paper introduces ContagionRL, an open-source Gymnasium-compatible simulation platform for systematically studying how reward function design shapes an individual agent's learned strategies for avoiding infection during a spatial epidemic.

What This Paper Is About

Traditional agent-based epidemic models give their simulated individuals fixed, hand-written behavioral rules, so they cannot adapt to changing conditions. Reinforcement learning could replace those rules with learned policies, but in epidemic settings nobody has systematically asked how the reward signal — the thing that tells the agent what "good behavior" means — shapes what the agent actually learns. This paper builds a simulation platform specifically designed to answer that question, and uses it to compare five different reward designs, ablate their components, and benchmark learned policies against fixed-rule baselines.

Key Contributions

1. ContagionRL platform. A modular, Gymnasium-compatible environment that couples a spatial SIRS+D compartmental epidemic model (adding a "Dead" state) on a toroidal grid with a continuous-control RL agent that simultaneously chooses a movement vector and a level of adherence to non-pharmaceutical interventions (NPIs). It supports both fully observable (MDP) and partially observable (POMDP, visibility-radius-limited) formulations.

2. A systematic reward-function taxonomy. Five reward designs are implemented and compared along a sparse-to-dense spectrum: Constant survival bonus, Reduce Infection Probability, Combined (Constant + Reduce Infection), Max Nearest Distance, and a novel Potential Field reward that treats nearby infected individuals as repulsive forces and rewards movement aligned with the resulting net force vector.

3. A component-level ablation and multi-algorithm benchmark. The Potential Field reward is decomposed into six ablated variants, and learned policies from PPO, SAC, and A2C are statistically compared against stationary, random, and greedy-distance-maximizing baselines under a consistent evaluation protocol (3 seeds × 100 episodes per condition, Mann-Whitney U tests with Bonferroni correction).

4. A deliberate information asymmetry design. The directional and magnitude components of the Potential Field reward are computed using the positions and states of all humans, even when the agent's observations are restricted by a visibility radius. This acts as a privileged critic: the training signal is globally informed, but the deployed policy must generalize from local observations alone.

Main Findings

• Learned policies clearly beat naive baselines. PPO, SAC, and A2C all achieved significantly longer episode durations (time spent uninfected) than both the stationary and random baselines, with strongly significant p-values across the board.

• A hand-crafted greedy heuristic is surprisingly competitive. The Greedy Distance Maximizer won statistically in median episode duration against all three RL agents. However, mean performances overlap in their 95% confidence intervals, so the greedy policy does not meaningfully outperform learned policies on average — it just has a tighter, higher distribution.

• The Potential Field reward is the clear winner. Across the five reward designs, the Potential Field approach produced significantly longer mean survival than every alternative. The Constant reward performed worst, which the authors attribute to its sparse, uninformative signal.

• Simple spatial rewards fall into local optima. Max Nearest Distance and Reduce Infection Probability both gave significantly worse results than Potential Field. Agents trained on these rewards learned to maintain distance from currently infected individuals rather than adopting a globally safer spatial strategy — a myopic failure mode.

• NPI adherence is easy to learn; spatial navigation is hard. Every effective reward function led the agent to learn maximal adherence to NPIs. The differentiating factor between reward designs was how well they supported navigation, not adherence.

• Directional and adherence signals are the critical reward components. Ablating the directional guidance, the overall movement incentive, or the explicit adherence reward each caused substantial, statistically significant performance degradation. By contrast, removing the movement magnitude scaling, the basic health reward, or the repulsion from susceptible individuals produced no significant change.

• Explicit incentives beat implicit ones. When the reward function contained many components, the agent did not infer the value of NPI adherence on its own — the adherence term had to be stated explicitly, or performance dropped significantly.

• Learned behavior has population-level effects. Agents trained with the Potential Field reward reduced the population's infection rate (new susceptible-to-infected transitions per timestep among the 40 non-agent humans) by 21.4% relative to the random baseline and 20.6% relative to the stationary baseline. Other reward functions showed smaller, often non-significant reductions. Since the agent represents only one of many individuals, the authors describe this as a conservative lower bound.

• PPO and A2C performed comparably at the top, with no significant difference between them, while SAC trailed both slightly.

Methodology in Plain English

The researchers built a square grid world that wraps around at the edges. Inside it live 40 "ordinary" simulated people, who move either randomly or in fixed patterns, plus one learning agent. Everyone is in one of four states: susceptible, infected, recovered, or dead. Infection probability depends on how close a susceptible person is to infected people, with transmission weakening over distance. Recovered individuals can lose immunity and become susceptible again, and infected individuals can die.

At each timestep, the learning agent outputs three continuous numbers: an x-movement, a y-movement, and an adherence level between 0 and 1. Adherence reduces but never fully eliminates infection risk. The agent starts at the center of the grid and is always susceptible; the episode ends when it becomes infected or when a time limit is reached. To prevent trivially short episodes, an automatic reinfection mechanism keeps the epidemic alive. Reward designs were swapped in and out while everything else stayed fixed, then PPO, SAC, and A2C were each trained on each design three times and evaluated over 100 episodes per seed. Statistical comparisons used one-sided Mann-Whitney U tests with Bonferroni correction. A separate ablation study removed individual terms from the Potential Field reward to see which components actually mattered.

Why This Matters

Impact on research. Reward engineering is a known pain point in reinforcement learning, but its role in epidemiological ABMs has gone largely unexamined. This work provides the tooling and the empirical evidence to treat reward design as a first-class research variable rather than an incidental implementation detail. It also establishes that a well-tuned greedy heuristic can match learned policies in this domain, which sets a useful bar for future RL work and warns against assuming that learned behavior is automatically superior.

Real-world applications.

• Public health policy modeling: Simulating how individual incentives (subsidies, mandates, information campaigns) translate into adherence and, ultimately, transmission rates.

• Digital health and exposure-notification apps: Designing risk-aware navigation or proximity guidance that individuals will actually follow, informed by what reward structures produce realistic compliance.

• Urban and workplace planning: Testing how layout, density, and movement patterns affect transmission under adaptive rather than fixed human behavior.

• Emergency preparedness: Stress-testing intervention strategies under partial observability, mirroring the incomplete information that real people and policymakers operate with.

Industry relevance. Companies building simulation platforms, digital twins of populations, or RL tooling can reuse ContagionRL directly as a testbed. Insurers and health systems modeling risk, game studios building realistic crowd behavior, and logistics firms planning safe routing all benefit from a validated account of how incentive design shapes learned spatial and compliance behavior.

Future Directions

• Multi-agent extensions. The authors deliberately restricted the study to one learning agent to isolate reward effects. The natural next step is populations of learning agents, where non-stationarity, strategic interaction, and emergent social norms become the central questions.

• Closing the observability gap. The current design uses a privileged critic, where the reward has global state information that the policy does not. Whether the learned policies transfer to settings where the reward itself must be locally computed remains open.

• Reward functions for objectives beyond survival. Every reward tested here optimizes the agent's own time-to-infection. Extensions targeting economic cost, social welfare, or population-level outcomes would require fundamentally different reward structures and may not inherit the same component sensitivities.

• Calibration against real behavioral data. Adherence emerged as trivially learnable in simulation. How well learned adherence patterns match real-world NPI compliance dynamics — which are shaped by fatigue, social norms, and misinformation — is untested.

• Richer human models. The non-learning population currently follows stochastic or fixed movement rules. Coupling ContagionRL with LLM-based or empirically calibrated agents would make the environment's unpredictability more realistic.

Target Audience

This paper is most valuable to reinforcement learning researchers working on reward shaping and safety-critical sequential decision-making, to computational epidemiologists and public health modelers interested in adaptive behavior, and to simulation engineers building agent-based platforms. Practitioners who need a ready-made benchmark for reward engineering in a spatial, partially observable, stochastic environment will find the platform and its evaluation protocol directly reusable. Readers without a background in deep RL will find the conceptual argument accessible but the implementation details dense.

Authors’ abstract

We present ContagionRL, a Gymnasium-compatible reinforcement learning platform specifically designed for systematic reward engineering in spatial epidemic simulations. Unlike traditional agent-based models that rely on fixed behavioral rules, our platform enables rigorous evaluation of how reward function design affects learned survival strategies across diverse epidemic scenarios. ContagionRL integrates a spatial SIRS+D epidemiological model with configurable environmental parameters, allowing researchers to stress-test reward functions under varying conditions including limited observability, different movement patterns, and heterogeneous population dynamics. We evaluate five distinct reward designs, ranging from sparse survival bonuses to a novel potential field approach, across multiple RL algorithms (PPO, SAC, A2C). Through systematic ablation studies, we identify that directional guidance and explicit adherence incentives are critical components for robust policy learning. Our comprehensive evaluation across varying infection rates, grid sizes, visibility constraints, and movement patterns reveals that reward function choice dramatically impacts agent behavior and survival outcomes. Agents trained with our potential field reward consistently achieve superior performance, learning maximal adherence to non-pharmaceutical interventions while developing sophisticated spatial avoidance strategies. The platform's modular design enables systematic exploration of reward-behavior relationships, addressing a knowledge gap in models of this type where reward engineering has received limited attention. ContagionRL is an effective platform for studying adaptive behavioral responses in epidemic contexts and highlight the importance of reward design, information structure, and environmental predictability in learning. Our code is publicly available at https://github.com/redradman/ContagionRL