Skip to content

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

🕒 Published (v1): 2025-11-22 10:02 UTC · Source: Arxiv · Venue: TMLR · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

ContagionRL is a Gymnasium-compatible RL platform embedding a spatial SIRS+D epidemic model, purpose-built to systematically study how reward function design shapes individual-level behavioral policies. Five reward designs are benchmarked across PPO, SAC, and A2C; a novel potential field reward consistently dominates. Directional guidance and explicit NPI-adherence incentives are identified as the critical reward components.

Problem

Traditional agent-based epidemic models rely on fixed, hand-coded behavioral rules rather than learned policies. Prior RL-for-epidemics work focuses on macro-level intervention optimization and does not systematically investigate how reward formulation affects the resulting individual-level learned behavior—leaving reward engineering in this domain largely unexplored.

Method

A single RL agent operates on a toroidal G×G grid populated by N non-learning humans, all governed by SIRS+D compartmental dynamics (S→I transition probability decays exponentially with distance: \(P_{\text{inf}} = 1 - \exp(-\beta \sum_i e^{-k_d d_i})\)). The agent's continuous action space is \((\Delta x, \Delta y, \alpha)\) where \(\alpha \in [0,1]\) is an NPI adherence level that modulates effective infection rate via \(\beta_{\text{eff}} = \beta(\epsilon_\alpha + (1-\epsilon_\alpha)(1-\alpha))\). Five reward functions span sparse-to-dense: Constant, Reduce Infection Probability, Combined, Max Nearest Distance, and a Potential Field that exerts repulsive force vectors from infected neighbors and adds health/adherence terms. Policies are learned with PPO, SAC, and A2C (stable-baselines3). POMDP variants restrict observations to a visibility radius \(r\) while the potential field reward retains global-state access during training (privileged critic). Ablations remove one Potential Field component at a time. Statistical tests use Mann-Whitney U with Bonferroni correction across 3 seeds × 100 evaluation episodes.

Key Contributions

  • ContagionRL: open-source Gymnasium environment unifying spatial SIRS+D epidemic dynamics with continuous-action RL, supporting MDP and POMDP formulations.
  • Systematic five-way reward comparison demonstrating the Potential Field reward achieves significantly longer survival than all alternatives.
  • Ablation study isolating directional guidance and explicit adherence incentives as the critical components; magnitude and susceptible-repulsion terms are dispensable.
  • Demonstration that partial observability (POMDP, \(r \in \{10,15,20\}\)) paradoxically outperforms full observability, attributed to reduced observation noise and implicit information bottleneck robustness.
  • Population-level analysis showing trained agents reduce non-agent S→I transition rates by 10–21% relative to baselines, with Potential Field achieving the largest reduction (\(p < 0.01\)).

Results

  • All three RL algorithms (PPO, SAC, A2C) significantly outperform Random and Stationary baselines in episode duration (\(p < 0.001\), Mann-Whitney U).
  • Greedy heuristic achieves statistically superior median episode duration versus all RL agents, but mean episode durations overlap (95% CI), indicating RL approaches comparable mean performance without hand-coded rules.
  • Potential Field reward yields significantly longer episode durations than Constant (\(p < 0.001\)), Reduce Infection (\(p < 0.01\)), Max Nearest Distance (\(p < 0.001\)), and Combined (\(p < 0.05\)) rewards (PPO, Bonferroni-corrected).
  • Ablations: removing Direction, Movement, or Adherence components causes statistically significant performance degradation (\(p < 0.001\)); removing Magnitude, Health, or Susceptible Repulsion does not.
  • POMDP agents (\(r=10,15,20\)) outperform the full-observability trained model on average reward, episode length, and infections/step; differences among the three radii are not statistically significant.
  • Population infection rate reductions: Potential Field −21.4% vs Random, −20.6% vs Static (\(p < 0.01\)); Constant −15.0%/−14.2% (n.s.); Reduce Infection −11.1% (n.s.).

Limitations

  • Single-agent design by construction; multi-agent dynamics (non-stationarity, emergent social behavior) are explicitly excluded and findings may not directly transfer.
  • Non-learning human population uses fixed stochastic or deterministic movement, limiting behavioral realism of the surrounding population.
  • Potential field reward uses global state as a privileged critic during training, an asymmetry that may not generalize to settings where global state is unavailable even for training.
  • POMDP superiority is counter-intuitive and the explanation (observation noise, dimensionality) is post-hoc; the feature importance analysis is deferred to an appendix without mechanistic proof.
  • Experiments use a small population (40 non-agent humans) on a bounded grid; scalability to large or heterogeneous populations is untested.

Relevance to Agentic AI / LLM Agents

Reward engineering is a foundational challenge for any agentic system that learns from environmental interaction, and this work provides one of the few rigorous ablation studies decomposing which reward components drive robust policy emergence in a complex, partially observable environment. The finding that explicit behavioral incentives (adherence terms) cannot be omitted even when the behavior is implicitly beneficial echoes broader lessons in RLHF and tool-use agent training: dense, component-wise reward shaping consistently outperforms sparse or myopic signals. The privileged-critic POMDP setup—globally informed reward, locally constrained observation—is directly analogous to reward-model-based training of LLM agents where the critic has access to ground-truth annotations the policy cannot observe at inference time. The information-bottleneck robustness finding (limited observation → stronger generalization) is relevant to designing observation spaces for real-world agentic deployments.