Reward Engineering for Spatial Epidemic Simulations: A Reinforcement Learning Platform for Individual Behavioral Learning
Authors: Radman Rakhshandehroo
Daniel Coombs
Venue: Accepted by TMLR ()

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

Published in Transactions on Machine Learning Research (02/2026)
Reward Engineering for Spatial Epidemic Simulations: A Reinforcement Learning Platform for Individual Behavioral Learning
Radman Rakhshandehroo
Clnnr@student.cee ubc.
Deportment of Computer Science University of British Columbia
Daniel Coombs
coombs@matt ubc.
Department of Mathematics and Institute of Applied Mathematics University of British Columbia
Reviewed on OpenReview: https: // openreview. net/ pormff id= PPEASaa3hk
Abstract
We present ContagionRL, Syymnasium-compattie reinforcement learning platform specif- ically designed for systematic reward engineering spatial epidemic simulations. Unlike traditional agent- based models that rely fixed behavioral rules, platform enables rigor- ous evaluation of how reward function design affeets learned survival strategics across diverse epidemic scenarios. ContagionRL. integrates a spatial SIRS+D epidemiologica 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 novel potential field approach, (PPO, SAC, A2C). Through systematic ablation studies, 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--hammacuutcal interventions while developing sophisticated spatial avoidance strategies. The platform's modular design enables systematic exploration of reward-bchavine relationships, addressing knowledge gap in models of this type where reward engineering has received limited altention. Contagion is an effec- tive platform for studying adaptive behavioral responses in epidemic contexts and highlight the importance of reward design, information structure, and environmental predictability in www.addraaaaannnnnaaaaa Cont hhtp:/Coottaaiionin
multiple RL algorithms
learning. Our code is publicly available at https: htpps////tthbb www/redradmaa/////eeeooooo
1 Introduction
The dynamics of infectious outbreaks, epidemics pandemics, are intricately linked individual behavior (Wang et al., 2021; Funk et al., 2010). Behavioral decisions such movements, risk perceptions and adherence to non-pharmacentical interventions (such epidemic unfolds (Wang et al., 2021; Funk et al., 2010; Manrubia & Zanctte, 2022; Ye ct al., 2020; Saad-Roy & Traulsen, 2023). Traditional epidemic models primarily use ordinary differential equations (ODEs) to keep track of disease population-wide with models such as SIR (Kermack McKendrick, 1927) or SEIR (Kermack McKendrick, 1927; Wilson Worcester, 1945). While these models are effective for a high-level overview of the discase outbreak, due to their assumption of homogencity, they neglect nuanced effects of individual behavior and spatial interactions (Bansal et al., 2007; Zachreson et al., 2022). Furthermore, they
and social distancing) could heavily affect how

---

Published in Transactions on Machine Learning Research (02/2026)
are incapable of capturing dynamic behavioral changes adaptive interventions (Funk et al., 2010; Manrubia & Zanette, 2022; Bansal et al., 2007).
The SARS-CoV-2 pandemic highlighted how individual-level decisions and behaviors, such as adherence to pharmaccutticall interventions (NPIs) like mask wearing, could profoundly change the trajectory of a disease outbreak (Saad-Roy Traulsen, 2023). Agent-Based Models (ABMs) (Grimm & Railsback, 2013), also known as Individual-Based Models (IBMs), offer compelling alternative by simulating systems from the bottom up, representing individual agents with specific attributes, spatial contexts, and behavioral rules. These models accommodate heterogeneity and complex interactions, thereby making ABMs particularly suitable for studying epidemics where individual actions drive epidemic-scale outcomes, while leaving them reliant on prescribed or rule-based behaviors rather than learned policies that exhibit more complexily and are more dynamic (Grimm & Railsback, 2013).
The lack of the learned policics could be addressed with Reinforcement Learning (RI) (Sutton & Barto, 2018) and Deep Reinforcement Learning (DRL) (Li, 2017; Mousavi et al., 2016; François-Lavel al., 2018) that provide powerful frameworks to optimize sequential decision making in dynamic systems (Mnih et al., 2015; Silver et al., 2016; Arulkumaran et al., 2017) by modeling them as Markovian decision processes (MDP). Such algorithms have been applied in epidemiological contexts to learn the optimal policies by focusing on centralized control of macroscopic interventions (Libin et al., 2021; Kwak et al., 2021; Du et al., 2023; Wan et al., 2021). Combining of RL or DRL algorithms with ABMs allows simulation of both adaptive and heterogeneous behaviors.
However, the main challenge in creating and developing informative ABMs lies in defining realistic and computationally tractable agent behaviors. Furthermore, RL.-basdd approaches are highly sensitive to the reward function and the reward desigu can have considerable impact the learned policy and the downstream interpretations of the model upon behavioral convergence (Ng et al., 1999; Ibrahim et al., 2024; Nguyen et al., 2025; Shihab et al., 2025).
Currently, there are limited sinulation plat.orms capable of supporting the integration of RL that permit the agents to learn complex strategies for navigation and adherence strategies by direct interaction within their environment. This creates important gap in the tools for investigating impact of reward design on learned dynamic behaviors at the individual level within biologically sound epidemic simulations.
To address this, we introduce CONTACIONRL, a specialized reinforcement learning platform designed to study how different reward structures influence agent behavior during epidemic simulations. This is a novel and parameterizable, Gymnasium compatible environment that unifies spatial compartmental epidemic modeling and reinforcement. learning. We integrate an extension of a Susceptibbe--Ifeeteed hetiiiee--neeeeeee-eoorrrr--d- (SIRS) framework called SIRS+D, which has an additional component for death, to track population-wice epidemiological dynamics, while also enabling the use of RL to model individual-level decision-making. This environment simulates single learning agent among population of non-learning humans, whose movements affect the discase outbreak. In this setting, we frame the agent's 's decision making as scquential control problem in which the agent learns behavioral policies through direct interaction with the environment. These policies include both spatial navigation and adherence to pharmmaaacutti interventions (NPIs), developed in response to localized epidemic conditions. Table compares our models to other approaches.
To demonstrate the elfectiveness of this novel environment, we conduct a variety of experiments. We investigate the nature of the policies learned by this RL agent under varying epidemiological and behavioral conditions. In addition, we conduct ablation studies to understand the critical components of the reward function that is the most impactful in the learning process. Benchmarking shows that these learned policies significantly outperform random and stationary baselines in terms of survival time. An ablation study shows the importance of both directional cues and adherence incentives within the reward function to learn a robust, risk-averse behavior. Across all effective reward structures, the agent consistently learned maximal adherence to NPIs. This framework provides a modular and highly configurable environment for studying emergent behaviors at the individual level within spatial epidemic models.
2

---

Published in Transactions on Machine Learning Research (02/2026)
2
Related Works
Compartmental Models in Epidemiology. Classical epidemiological models, such SIR (Susceptible- divide a population into discrete groups (compartments) based on disease state (Kermack & McKendrick, 1927; Wilson & Worcester, 1945). The flow of individuals between compartments is described by ordinary differential equations (ODEs). Extensions like SEIR (adding an exposed group to act as a discase incubation period) or SVE(R)IRS (adding Vaccinated, accounting for waning immunity) incorporate more states to increase realism (Brauer et al., 2019; Friedman & Kao, 2014; Hethcote, 2000; Chyba et al., 2024; Balisacan al., 2021). While traditional compartmental models provide valuable population-Ievel insights, assumptions of homogeneity limit their ability to represent individual decision-making (Cangiotti et al., 2024; Siegenfeld et. al., 2022; Dimarco et al., 2021; Jiang et al., 2024; Bostanci & Conrad, 2025; Rose et al., 2021; Kong et al., 2016). These models also struggle to capture dynamic behavioral changes or adaptive interventions (Brauer et al., 2019; Siegenfeld et al., 2022). Ignoring such factors can limit accuracy and lead to misleading conclusions (Sudhakar et al., 2024; Siegenfeld et al., 2022). Stochastic formulations are particularly relevant. for small populations or during the early stages of an outbreak (Braner et al., 2019).
Agent-Based Models (ABMs). ABMs address the limitations of compartmental models like SIR by modeling heterogeneous individuals using specific attributes and interaction rules within an environment network (Tracy et al., 2018; Hunter et al., 2018; 2017). Furthermore, ABMs excel at capturing individual vari- ability, stochasticity, local interactions, and complex contact networks, which are difficult for compartmental models (Siegenfeld et al., 2022; Hunter et al., 2017). Recent advances demonstrate substantial sophistication such as large-scale ABMs that can simulate millions of agents with GPU acceleration (Chopra et al., 2021) or employ LLM-based adaptive agents (Park et al., 2023).
Challenge of Reward Engineering. Sophisticated ABMs show potential for learning-enabled epidemic modeling, However, systematic evaluation of reward function design remains important methodological gap. In reinforcement learning, reward design shapes the fundamentals of the learned behavior with different reward struct ures leading to considerably different learned policies at convergence (Ng et al., 1999; Amodei et al., 2016; Nguyen et al., 2025; Shihab et al., 2025). Despite the importance of this problem in RL, epidemic modeling applications have not systematically investigated how reward structure affects learning under varying conditions.
Reinforcement Learning (RL) for Epidemics. RL has been increasingly leveraged to design policies for epidemic intervention at population and individual level. Individual-basedd Reinforcement Learning Epidemic Control Agent (IDRLECA) proposed an RL agent at an individual level using GNNs to estimate infection probabilities and balance epidemic suppression vs cost of mobility in its reward function (Feng et al., 2023). VEHICLE (Feng et al., 2022) handles challenges of unobservable asymptomatic cases and delayed effects by using a combination ion of hierarchical RL with GNNs. Recent work has also focused on using actor critic methods to minimize a cost that considers both the epidemiological, economic and social cost for multi-intervention planning (Mai et al., 2023) using RL for scheduling lockdown periods (Arango & Pelov, 2020) that incorporates the ICU capacity and minimizing time spent in lockdowns. HRL4EC (Du et al., 2023) also leverages hierarchical RL. and multi-componen reward that considers both health and socioeconomic cost. SiRL (Bushaj et al., 2023) uses a detailed ABM with RL and feeds the compartmental levels to the agent and uses rewards that are tied to intervention effectiveness and discase spread. These approaches excel at macro-level policy optimization such as determining lockdown time, resource allocation, and intervention soquencing (Mai al., 2023; Arango & Pelov, 2020; Khadilkar et al., 2020). While valuable for analyzing epidemic policy, these centralized approaches do not examine the impact of different reward formulations on the learned RL policy.
Addressing the Reward Engincering Gap. This work addresses this gap by providing a systematic framework for evaluating reward formulation in spatial epidemic sinulations. Instead of assuming particular behavioral patterns, we investigate the impact of different reward structure, varying environmental conditions such as partial observability, different infection rates, population density, distance docay factor (strength of discase transmission over distance), effectiveness of adherence non-pharmaceutical interventions and spatial constraints, on learned behavior.
3

---

Published in Transactions on Machine Learning Research (02/2026)
SIRS+D Spatial Environment
Compartmental Dynamics
Toruidal Grid & Agent
Agani Interface
my Ahe Ay). M/
Puf
Actiom: (dr dy.a)
Moscoual visior 1 olloreoo conted
Agent (orange) POMDE: Vialbiltty sndins
Reward Function Design Space
Rawoard Taxonamy (Sparsu (left) Duner (right))
Conatast: Allve reward
Redure In- fection Prob.
resand and rudson
Compiizdd: alive
Max Distance
Potential Pield
Multi- Multi-Diimeniioal Evaluation
Algusthin PPO, SAC, Ase
Olees MDD, FOMDD
Dimensions:
Cyelte
Baseltzon: Gready
Systematic Ablation: vacionte
Acherence: to varlutiore
Metrica: Epleads duration, Serviral time. Indections/step
Redhuttin
Statissiron Marn Wettony U, Rowferment comection Rigan coell x 100 optonils (300 coaluntion Lilee feardittion)
Figure 1: ContagionRL System Architecture. Top: SIRS+D spatial epidemic environment with toroidal grid, configurable observability, and continuous agent control interface. Middle: Reward function design from sparse to potential field-based rewards. Bottom: Multi-dimensional experimental evaluation.
3 Methodology
CONTACIONRL is custom Gymmasiuum-compatibbe (Towers et al., 2024) environment that integrates a SIRS+D compartmental epidemiological model (Hethcote, 2000). Built within an agent- modeling framework, the environment simulates single reinforcement learning (Sutton & Barto, 2018) agent interacting with a population humans whose behaviors may be stochastic or deterministic or mix of both. This design enables the study of learned risk-averse navigation and adherence behaviors response to an evolving epidemic (Saad-Roy & Traulsen, 2023) in different contexts.
Single Agent Design Rationale. Epidemic dynamics inherently multi-agent interactions (Manfredi & D'Onofrio, 2013; Wang et al., 2016). However, our single-agent approach is methodological choice that enables transparent attribution of learned behaviors to reward function design, which is the core scientific question of this work. Multi-agent environments introduce complexities such as non-stationarity due to other learning agents (Foerster et al., 2017), emergent social behavior (Leibo et al., 2017) and strategic interactions (Tampuu al., 2017). These factors obscure whether observed behaviors stem from game-theoretic equilibria or from the reward structure itself. Single-agent formulation permits the identification of essential reward components for effective behavior learning. Understanding individual-level is a necessary prerequisite for designing effective reward systems in multi-agent deployments where each agent still learns from its local reward signal. By isolating a single learning agent in controlled population of non-learning humans, we can conduct detailed assessments of how different reward formulations affect the learned strategies without interference from adaptive social behaviors. The controlled nature of this methodology allows insights into the relationship between reward and behavior, that will be transferable more complex multi-agent deployments where understanding individual-level incentive responses becomes critical for designing effective reward systems across a population.
relationships
4

---

Published in Transactions on Machine Learning Research (02/2026)
3.1 SIRS+D Simulation Environment
Environment Dynamics. The simulation happens over discrete timesteps within two-dimensional grid space of size G x G with periodic boundary conditions (Allen et al., 2009; Keeling & Rohani, 2008; Frenkel & Smit, 2002). The enviromment is populated by N non-learning humans and a single controllable Reinforcement Learning (RL) agent. Each individual, human or agent, can exist in one of four states: Susceptible (S), Infected (I), Recovered (R), or Dead (D). We conduct set of sensitivity analysis in Appendix H and vary the infection rate of disease (3) in Figure 12 and Table 15, vary the grid size and population density in Figure 13 and Table 16, vary the adherence effectiveness (da#) in Figure 14 and Table 17, and vary the distance docay of disease transmission in Figure 15 and Table 18.
Human Movement and State Transitions. Non-learning humans exhibit stochastic or deterministic movement based on a specified movement type with a defined movement scale. Stochastic humans make their decisions by sampling from a Gaussian distribution within the bounds of the action space used to mimic the noisy nature of observation perception by the learning RL agent. Conversely, the deterministic non-learming humans are used to mimic the social and spatial constraints of a pandemic such as commuting. Their state transitions are governed by probabilistic rules. A susceptible human h, becomes infected with a probability (m,( determined by their exposure to nearby infected individuals h4. Exposure is calculated based on the distance d(h,,h) h:) between the susceptible human and each infected human, potentially limited by a maximum distance. The infection probability follows Pinr(h,) = exp(-B>.. e e=Ra-d(h,.hi)) where / is the base infection rate and ky is the distance decay factor. Once infected, a human recovers with a fixed probability p at each timestep. After recovery, a human may lose immunity and revert to the susceptible state with a probability À at each timestep. Infected humans may also die with a probability ó at each timestep. Dead humans are removed from contributing to infection spread and cease movement. Appendix C.2 contains description of environmental parameters and their descriptions.
Agent Dynamics and Control. The agent operates under SIRS+D dynamics identical to other individuals but differs in that it actively selects its movement and adherence levels via a learned policy. These choices affect both its spatial trajectory and susceptibility to infection. For a detailed formalization of the agent's 's action space, including adherence effects and boundary handling, see Section 3.2. Figure 8 shows a render of the environment.
Initialization and Episode Structure. Each simulation run begins with the agent positioned the grid's center. The non-agent humans are distributed randomly, with an parameterization number of individuals set to the infectious (I) state. These initial infectees are placed at least a minimum safe distance from each other and from the agent. All remaining humans commence in the 'S' state, ensuring a minimum initial distance from the agent's starting position. The minimum safe distances are used to prevent trivial early termination of a simulation. The agent itself always starts as Susceptible ('S²) with a predefined initial adherence level. To counteract stochastic extinction of the disease and trivial success episodes, configurable reinfection mechanism is employed. If the number of infected individuals reaches zero, specified nunber of susceptible humans (located beyond the minimum safe distance from the agent) are randomly chosen and transitioned back to the T state, An episode concludes upon reaching the maximum simulation time limit. Termination can occur carly if the agent transitions to infected (TT statc), or the infection is cradicated and the reinfection mechanism cannot be activated due to insufficient eligible susceptible individuals.
3.2 RL Agent Formulation
We implement the agent's learning problem as configurable decision-making process that supports both Markov Decision Process (MDP) and Partially Observable MDP (POMDP) formulations. We use both the MDP and POMDP formulation in our experiments. The POMDP employs threshold-based visibility constraint, where agents can only observe other entities within a specified Euclidean distance radius. The agent's s policy T (a)))o)), mapping observations to actions, is learned using three deep reinforcement learning algorithms: Proximal Policy Optimization (PPO) (Schulman et al., 2017), Soft Actor-Critic (SAC) (Haarnoja et al., 2018), and Advantage Actor Critic (A2C) synchronous variant of A3C algorithm (Mnih et al., 2016), all implemented via the STABLE-BASPLINESS library (Raffin et al., 2021).
5

---

Published in Transactions on Machine Learning Research (02/2026)
Observation Space. The observation space provides the agent with information about its internal state and the surrounding environment at each timestep L. The observation or, constructed via the environment's internal observation method, is comprising the agent's current adherence level a, (a scalar in [0. 1]) and a flattened feature vector describing the non-learning population. This vector aggregates information for each non-learning hurman hy. For every human hy: the features consist of the relative position (Aray, Ayos); representing the normalized displacement from the agent calculated considering toroidal boundaries and scaled by grid size (values in [-0.5, 0.5]); the normalized distance (daj/ /mmx), which is the Euclidean distance scaled by the maximum possible grid diagonal distance dmaxi and a binary infection status indicator (Lz).
Action Space. At each timestep f, the agent selects a contimnous-valued action a, = (Az, Ay, a), where (Ar, Ay) E [-1, 1J defines a movement vector and a E [0, IJ specifies the adherence level to non-pharmaceutical interventions (NPIs). NPIs are methods used to reduce the spread of an epidemic disease without any pharmaceutical drug treatments such as wearing masks or social distancing and are considered extremely effective ways of controlling primary outbreaks that reduce discase transmission and consoquentially mortality (Rizvi et. al., 2021; Ferguson el al., 2020). The agent's 's position is updated using the movement vector, scaled by maximum step size, and wrapped around the grid's periodic boundaries:
T++1 (Iz + Ax) mod G, Vt+1 = (yt + Ag) mod G,
(1)
where G is the grid size. The adherence level a directly influences the agent's probability of infection upon exposure to nearby infected individuals by modulating the effective infection rate:
B B 3 (Ea + (1 (1- (1 (1) a))
where to E (0, 1) represents the residual infection risk at full adherence. Thus, higher adherence values reduce susceptibility, but never eliminate risk entirely.
3.3 Reward Functions
To investigate the impact of reward engineering on policy learning within the CONTAGIONRI environment, we implemented and evaluated several distinct reward functions. These functions vary in complexity and the specific agent behaviors they aim to incentivize, with the overarching goal of prolonging the agent's 's susceptible state. The simplest among these is a Constant Reward, which provides fixed positive signal if the agent remains uninfected. The Reduce Infection Probability reward directly incentivizes the agent to minimize its calculated likelihood of infection at cach step, based exposure to nearby infected individuals and its current adherence level. This can be employed as Combined Reward by supplementing it with the constant bonus. Another approach, Maximize Nearest Distance, focuses on spatial dynamics, rewarding the agent for maintaining a distance from the closest susceptible or infected individuals, particularly beyond a predefined threshold.
The most complex of the explored designs is the Potential Field Reward. This function models interactions forces, where other agents (particularly infected ones) exert repulsive forces the agent. The reward encourages movement that aligns with the net resultant force vector, effectively guiding the agent away from high-risk areas. This directional guidance is complemented by terms rewarding the agent for its health status and NPI adherence. These reward function are tested in different environmental configuration. Details of the reward calculation are provided in Appendix C.3.
An important design consideration the information scope of each reward component. The health and adherence terms depend solely on the agent's 's own state and action, making them purely local signals. In contrast, the directional and magnitude components of the potential field reward are computed using the positions and states of all humans in the simulation, regardless of the agent's visibility radius. This creates a deliberate information asymmetry between the agent's observations and the reward signal: under POMDP conditions, observations respect the visibility constraint (Section 3.2), while the reward provides a globally informed training signal. This design functions as a form of privileged critic (Pinto et al., 2017), where the reward shapes policy gradients using complete state information during training, but the learned policy must generalize from local observations alone at evaluation time. Table summarizes the information scope of each reward component.
6

---

Published in Transactions on Machine Learning Research (02/2026)
4 Baselines
To contextualize the performance of the learned policies with RL, we established several non-learning deterministic baselincs. These policies represent simple, fixed strategics and serve as benchmarks to evaluate the complexity and effectiveness of the learned behavior. We implemented three distinct baseline strategies (see Appendix C.6 for more dctails).
Stationary Baseline. This agent remains immobile with zero adherence to NPIs (Ar = Ay= 0) throughout the episode. Its performance depends entirely initial conditions and the stochastic behavior of the surrounding population.
Random Baseline. At each timestep, this agent samples movements Ar, Ay ~ U(-1,1) and adherence a ~ ((0 1), while being oblivious with regard to the environment and observations entirely. It serves as a reference for performance under uninformed, mobile behavior.
Greedy Distance Maximizer Baseline. This heuristie policy implements specific survival strategy focused on immediate risk avoidance. At each step, it first identifies the nearest infected human. It then evaluates a predefined set of potential movement actions across the eight cardinal and intercardinal directions, plus staying stationary. It systematically simulates the outcome of each potential action, by computing the hypothetical distance to the nearest infected human if that move were taken, and selects the action that yields the greatest distance separation. The policy's evaluation of counterfactual outcomes provides a situational awareness absent in simpler baselines. While myopic and handcrafted, this policy provides strong heuristic benchmark by explicitly evaluating the spatial consequence of each potential action.
5 Results
Performance of different algorithms in ContagionRL environment. To evaluate the utility of the CONTACIONRL for obtaining effective learned policies via reinforcement learning (RL), we trained and assessed three distinct RL algorithms: Soft Actor-Critie (SAC) (Schulman et al., 2017), Proximal Policy Optimization (PPO) (Schulman et al., 2017), and Advantage Actor-Critic (A2C) (Mnih et al., 2016). Their performance was benchmarked against baselines mentioned in section (configuration details are in Appendix C.2 and Appendix C.5).
Figure 2 shows the distribution of episode durations for both baselines and the RL agents. The distributions of episode duration for the stationary and random basclines have substantially less variance compared to those of the RL agents and the greedy heuristic. To determine the statistical significance of these different distributions, we conducted pairwise comparisons using the Mann-Whitney U test. The outcomes of these tests, including both two-sided and onc-sided p-valuos, are presented in Table 1.
The stalistical analyses reveal that the learned policies derived from the RL approaches (PPO, SAC, and A2C) significantly outperform both the random and stationary baselines. The greedy heuristic baseline demonstrates superior performance in pairwise comparisons against all three RL agents, achieving statistically significantly longer episode durations (Table 1). However, as depicted by the mean episode durations and their 95% confidence intervals in Figure 9, the performance of the RI. agents is comparable to that of the greedy baseline, with overlapping confidence intervals.
This tells us that while the distributions episode durations and their medians, differ significantly in favor of the greedy approach, the mean performance of greedy does not outperform learned RL policies. These results collectively demonstrate several key points: (1) The observation formulation within the CONTACIONRL environment provides sufficient and informative signal for the RL agents to develop effective control policies that substantially exceed naive baselines. (2) The environment supports the application and comparative evaluation of various RL algorithms, highlighting their compatibility with its inherent stochasticity and complex dynamics, leading to policies competitive with a strong heuristic. (3) While specialized greedy heuristic can achieve statistically superior median performance, the learned RL agents develop sophisticated strategics that yield comparable mcan performance, showing the environment's utility for investigating adaptive decision- making processes that can approach the efficacy of well-designed heuristics without explicit rule-based programming used in A ABMs
7

---

Published in Transactions on Machine Learning Research (02/2026)
8.
300
Episade Duration (steps)
-0
9-
Stationry
Random
PPO
SAC
A2C
Greedy
Algorithm
Figure 2: Episode duration distributions across different agents, including learning-based (PPO, SAC, A2C) and non-learning baselines (Random, Stationary, Greedy). Each small black dot represents one episode in the totality of episodes across seeds 100 evaluation runs. Per-seed means are shown large black dots with white outlines. This figure complements the summary statistics in Figure and statistical comparisons in Table 1.
Table Pairwise comparisons of episode durations between agents using the Mann Whitney U lest. Two-sided p-values assess distributional differences, and one-sided p-values (Bonferroni test for performance advantage. The Winner is the agent with significantly longer episode duration (after correction): 5_" indicates significant difference. Siguificance codes: < 0.05, p < 0.01,
0.001, n.8. not significant.
Agent A
Agent B
p (2-sided)
p (1-sided)
Sig (2) Corrected p (1)
Sig (1) Winner
Stationary
Random
0.5594
0.7204
Stationary
PPO
3.22e-25
1.61e-25
***
24
PPO
Stationary
SAC
3.83e-13
1.92e-13
***
2.88e-12
SAC
Stationary
A2C
2.46e-22
1.23e-22
1.84e-21
A2C
Stationary
Greedy
3.87e-48
1.94o-48
2.91e-47
288
Greedy
Random
PPO
1.39e-25
6.96e-26
1.04e-21
PPO
Random
SAC
1.27e-12
6.32e-13
***
9.49e-12
***
SAC
Random
A2C
1.29e-22
6.45e-23
9.68e-22
A2C
Random
Greedy
9.96e-51
4.98e-51
***
7.47e-50
28*
Greedy
PPO
SAC
0.00103
0.00052
0.00775
PPO
PPO
A2C
0.7604
0.62
11...
1.S..
PPO
Greedy
0.000639
0.000320
0.00479
Greedy
SAC
A2C
0.00430
0.00215
0.0323
A2C
SAC
Greedy
1.06e-12
5. 5.28e-13
7.93e-12
Greedy
A2C
Greedy
0.000341
0.000171
0.00256
Greedy
Comparison of Different Reward functions in ContagionRL. Our primary algorithm evaluations utilized potential field-based reward. Now, we investigate the differential impact of diverse reward structures on performance. We use the environment configuration in Appendix C.2 and alter the reward functions to gauge difference in performance with PPO. The principal objective for the agent remains the avoidance of infoction. This quantified by the episode duration, which measures the number of timesteps during which the agent remains susceptible before experiencing its first infection. In other words, the agent becomes infected for the first time at. timestep I, the achieved episode duration is I Although alternative reward
8

---

Published in Transactions on Machine Learning Research (02/2020)
Far-seed Mean
000
Epicode Duration (stopo)
000
200
100
Constom:
Reduze Infection
Constant a Reduse Int. Reward Function
Max Nearest Distance
Potential Fieid
Figure Comparison of PPO agent performance under five different reward functions. Each model was trained with three random seeds and evaluated over 100 episodes per seed (300 episodes total per reward function). Violin plots show the distribution of episode durations, overlaid with boxplots and per-episode results (black points). Large black dots with white outlines indicate the per-seed mean. One-sided Mann- Whitney U tests (with Bonferroni correction) compare each reward function to the Polential Field baseline. Statistically significant differences (* p < 0.05, p < 0.01, *** p < 0.001) are annotated. See Table 10 for exact test values and Figure 10 for corresponding mcans with confidence intervals.
schemes could be designed to elicit and study different behaviors, our focus here is optimizing for this primary goal.
Figure 3 and Figure 10 visually summarize the performance distributions and mean episode durations, respectively, for policies trained under five distinet reward functions (see Section 3.3 and Appendix C.3). The statistical significance of these differences is further detailed in Table 10, which presents pairwise Mann-Whitman test comparisons against the Potential Field baseline.
The results indicate that the Potential Field reward function, via its dense and spatially informed feedback, results in significantly superior performance, achieving the longest mean episode durations. Policies trained with the Constant reward structure show poor performance. We attribute this to the sparse and uninformative nature of constant bonus, which provides little guidance for muanced decision-makinn Similarly, the Max Nearest Distance reward function also resulted in significantly shorter episode durations compared to the Potential Field approach.
Both the Max Nearest Distance and the Reduce Infection probability reward functions may guide the agent towards suboptimal local optimas. An agent might learn to maintain specific distance from currently infected individuals, thereby maximizing its immediate reward. However, it fails to adopt globally optimal positioning strategy that minimizes long-term infection risk. This highlights the challenge of designing reward functions that encourage far-sighted behavior rather than myopic optimization in complex spatial settings.
An important observation, across all tested reward functions, is the consistent emergence of high NPI adherence the learned policies. This suggests that the benefit of maximizing adherence is a relatively identifiable signal within the CONTAGIONRL environment, discoverable through lcarning irrespective of the specific reward structure guiding navigation. A detailed analysis of learned adherence behavior across all reward functions and ablations is provided (Appendix E). The primary challenge lies in learning effective spatial navigation strategies, a task for which the choice of reward function is more important. The data presented in Figures and 10, and substantiated by the statistical comparisons in Table 10, unequivocally demonstrate the influence of reward function design on the learned policy to achieve the desired infection avoidance behavior.
Population-Level Epidemic Impact. To assess whether the agent's learned behavior affects population- wide dynamics, we measure the infection rate: mean new S-IT transitions per timestep among non-agent.

---

Published in Transactions on Machine Learning Research (02/2026)
Table 2: Population-Ievel epidemic impact of learned policies. New Inf./step reports the mean number of
new S-II transitions per timestep among the 40 non-agent humans, averagod over 90 episodes (3 seeds 30). Lower values indicate reduced epidemic burden. A% columns show the percentage reduction relative each bascline. One-sided Mann-Whitney U tests with Bonferroni correction (5 comparisons per bascline). Significance: 0.05, p < 0.01, n.s. not significant.
Random
Static
Policy
New Inf./step
A A%
Sig.
A% %%
Sig.
Random (baseline)
1.133
Static (baseline)
1.121
Potential Field
0.890
-2..%
20066
Const.+Red. Inf.
0.925
-18.3%
**
-17.5
**
Constant
0.963
-15.0%
11.22%
Max Near, Dist.
0.972
-14.2%
n.s. -1.3%
11.8.
Reduce Infection
1.007
-11.1%
n.s.
-10.2%
n.8.
humans, which controls for differences in episode length across policies. Table 2 shows that trained policies reduce population-level infection rates by 10-21% relative to basclines, with the Potential Field reward achieving the largest reduction (p < 0.01, Bonferroni-coorecttt). This ranking is broadly consistent with individual survival performance (Figure 3), consistent with the mechanism that agents avoiding infection remove themselves as transmission vectors. Given that the agent constitutes only one of many entities, this represents a conservative lower bound on the impact of learned avoidance behavior.
Ablations of reward function. To understand the contributions of individual components to the overall efficacy of the potential field reward function, conduct a systematic ablation study. Specific elements of the full potential field (FPF) reward were individually removed, and the performance of policies trained with these ablated reward structures was compared against the FPF baseline. We use the same configuration as the previous section but changed the value of the reward ablation field (Appendix C.2). We did the following ablations: removal of the movement magnitude component (No Magnitude), removal of the directional component of movement (No Direction), removal of the entire movement-related reward (No Movement), removal of the adherenco-basee, reward (No Adherence), removal of the health status reward (No Health), and removal of the repulsion force from susceptible individuals (No Susceptible Repulsion). The comparative performance distributions are illustrated in Figure 4, mean performances are shown in Figure 11, and detailed statistical comparisons are presented in Table 11.
The results (Table 11) indicate that several components are important to the FPF reward's success. Specifically, the removal of the directional guidance (No Direction), the overall movement incentive (No Movement), or the reward for NPI adherence (No Adherence) led a statistically significant and substantial degradation in performance of learned policy compared to the FPF as our baseline. These elements provide the primary signal guiding the agent's 's navigat ion to safer regions. Without this explicit directional information, the agent's 's ability to learn effective spatial strategies is severely impaired.
The ablation of the movement magnitude component (No Magnitude) did not result in a statistically significant performance change, suggesting that the direction of the suggested movement is more important than its precise scaling in this context. Similarly, removing the basic hcalth reward (No Health) or the repulsion from susceptible individuals (No Susceptible Repulsion) did not significantly impair performance, implying these components are either less influential or their effects are implicitly captured by other elements of the FPF reward.
Upon removing the adherence reward component (No Adherence), we observed siguificant degradation in agent performance. While we might expect agents to implicitly learn the importance of high NPI adherence due to its direct impact on infoction risk (especially under simpler reward schemes), this was not what we observed with the complex FPF reward. It is evident from our findings that when diverse set of reward signals is provided, explicitly rewarding critical behaviors like NPI adherence is crucial for effective learning,
10

---

Published in Transactions on Machine Learning Research (02/2026)
Porened
400
Eghoodo Daration tatopa
208
No Mogriudd
No Direcion
No Movernntt
No Adheeenos
No Health
No Repulsionn
Fill
Fiee
Ablation Varient
Figure 4: Ablation study of the Potential Field reward function. Each variant was evaluated over 100 episodes across 3 training seeds (300 episodes total). Violin plots show the distribution of episode durations, overlaid with boxplots and individual episode results (small black dots). Large black dots with white outlines represent per-scod means. One-sided Mann-Whitncy tests (Bonforroni-corrccee)) compare each ablation the full model (Full Potential Field), with significance annotations (* 0.05, ** p 0.01, *** p 0.001) shown for statistically significant. differences. See Figure 11 for aggregated means with confidence intervals, and Table 11 for full statistical test results.
as opposed to depending on the policy to deduce their siguificance. This highlights the importance of careful component weighting and inclusion when engineering sophisticated reward functions for complex tasks.
Comparison of Partial Obscrvability (POMDP) and Full Observability. To understand the Icarned behavior under limited visibility, we formulate the epidemic control problem as a partially observable Markov decision process (POMDP). In real epidemics, actors have incomplete information about the environment state and infection status of individuals (Gersovitz Hammer, 2004; Geoffard Philipson, 1996; Farboodi et al., 2021). We implement this POMDP formulation using a visibility radius constraint, where individuals within radius of RL agent's position are passed observations and individuals beyond this radius are unobscrvable. This natural information asymmetry captures realistic scenario where agents have limited sensing information networks that decay with distance. Figure 5 demonstrates trained partial visibility models (r=10, r=15, r=20) paradoxically outperform the baselines and the full visibility trained model, Trained Full, in average reward, opisode length and infections per timestep. This counter-intuitive result suggests that complete information can be detrimental to learning, polentially due to observation noise and the increased dimensionality high-dimensional observation spaces (see Appendix for feature importance analysis). Limited observability forces the agent develop more robust strategies and reduces the complexity of the decision space. The trained (Full) model's performance is par with the greedy heuristic. Notably, performance differences between 1=10, r=15 and r=20 conditions are not statistically significant according to Table 12, indicating a performance plateau beyond a certain visibility threshold. We note that the potential field reward is computed from global state even under POMDP conditions (Section 3.3), functioning as a privileged training signal. The agent's improved performance under limited visibility therefore reflects gemuine policy robustness rather than information leakage through observations.
The POMDP perspective reveals that effective epidemic control strategies can come from local information processing. Agents are able to infer general patterns from partial observations. The improved performance occurring under constrained visibility shows that an information bottleneck can enhance robustness.
Social and Economic Constraint in Human Movement. Mobility is a key consideration in epidemio- logical models (Fakir & Bharati, 2021) and literature review shows that mobility patterns remain unchanged in pandemic for essential or low wage workers (Gallagher al., 2021). Inspired by this, we conducted an experiment evaluate the performance of the agent under these socioeconomic constraints mobility. Figure 6 shows that the trained PPO agent outperforms stationary, random and greedy baselines when trained
11

---

Published in Transactions on Machine Learning Research (02/2026)
4000)
Figure 5: Impact of visibility radius constraints on RI agent performance in epidemic control. The figure compares four agent types across different observation capabilities: Full Visibility (agent observes all infected individuals), and Limited Visibility with radius constraints r=10, r=15, and r=20 (agent only observes infectod individuals within the specified radius). Agent types include: Stationary (no movement); Random (random actions); Greedy (heuristic policy avoiding nearest infected), and Trained (PPO-trained RL agents with respective visibility constraints). Results averaged across 3 random seeds with 100 evaluation episodes per seed (N=300 per condition). Error bars show 95% bootstrap confidence intervals from per-seed means, Left: Average cumulative reward per episode. Center: Episode length (survival time in timesteps). Right: Infections per timestep calculated from total infections divided by cpisode length. Hatching patterns distinguish trainod variants: dots (r=10), diagonal lines (r=15), and crosses (r=20). Limited visibility agents (r=15, r=20) achieve higher performance than full visibility, suggesting that observation constraints can improve learning by reducing obscrvation noise and focusing attention on ncarby thrcats.
workplace /home movement cycle. We notice that the stationary and random baselines declined in this cyclic movement, but the trained and greedy baseline improve considerably. Furthermore, the trained RL agent has higher average reward and higher mean episode length, with an extremely low rate of infections per timestep (Table 3). In the random environment, the agent struggles with long predictions and has a short horizon on its reward, as it performs worse than the greedy baseline. We attribute its inability to beat the local optimization of the greedy algorithm to the noisy observations and movements of the random environment. regularities of workplace/ hooplecee/omeeeee home mobility cycle, a trained agent can
By incorporating the formulate more effective long-term strategies. The resulting expansion of its planning horizon is empirically validated by substantial improvement in cumulative average reward over agent exposed to random mobility pattern. We anticipated that trained agents operaling under cyclic movement patterns would exhibit lower standard deviation across performance metrics compared to those under random movement. However, our results contradicted this expectation. Furthermore, while infections per timestep increasod for all agent types, the trained agent demonstrated no such increase, maintaining stable infection rates throughout the evaluation period.
6 Conclusion
We introduced CONTAGIONRL, a novel reinforcement learning platform that enables systematic evaluation of reward function design and for analysis of learned behavior in a spatial epidemics model. Through comprehensive experiments across multiple RT algorithms (PPO, SAC, A2C) and diverse environmental conditions, we show that reward engincering dramatically influences the learned survival strategics. Systematic reward ablations reveal that directional guidance and adherence incentives are essential components for robust. policy learning while agents consistently learn adherence to non pharmaceutical interventions regardless of the reward formulation used for spatial navigation.
We show that different RL algorithms can successfully learn mitigation stralegies in the environment. We compare different reward function formulations on performance and structural ablation of reward function design. Also, we find that constraining agent visibility yields superior performance relative to full observability, attributable reduced dimensional complexity. Furthermore, we reveal structured movement compared to random movement leads to increased longevity of the trained agent. The CONTAGIONRI. environment.
12

---

Published in Transactions on Machine Learning Research (02/2026)
1.
Intectio per Timeatop
Mas Eptaodo Langh
200
Agont Type
Stationery
Random
30
Greedy
Movement Partern
Movement Pataann
Movement Patheen
Figure 6: Performance comparison of RL agents across different human movement patterns in epidemic control. The figure shows three metrics across two movement patterns: Random Movement (contimuous random walk) and Workplace/ Home Cycle (structured movement between workplace and home locations). Four agent types are compared: Stationary, Random, Greedy, and Trained (PPO agent). Results are averaged across 3 random seeds with 100 evaluation episodes per seed. Error bars show 95% bootstrap confidence intervals from per-scod mcans. Left: Average cumulative reward per episode. Center: Mcan episode length (survival time). Right: Infection rate. Trained agents significantly oulperform all baselines in the Workplace/ Home Cycle environment but show mixed results in Random Movement. Statistical results can be found in Table 3
Table 3: Statistical significance tests comparing RL agent performance across movement patterns shown in Figure 6. Trained agents were evaluated using Mann-Whit U tests with Bonferroni correction for multiple comparisons (a = 0.05). Each condition used 300 episodes (3 seeds 100 episodes per seed). Cross-condition compares trained agents between movement patterns.
Movement Pattern
Comparison
Winner
U-stat.
p-valuc
Sig.
Trained vs Stationary Trained vs Random
Trained
80,724.0
p 0.001
Random Movement
Trained
69,720.0
p 0.001
Trained vs Greedy
Trained
49,117.0
p 0.420
n.S.
Trained vs Stationary
Trained
89,067.0
p < 0.001
Workplace/Hon Home Cycle
Trained Random
Trained
87,819.0
p 0.001
Trained vs Greedy
Trained
85,853.0
p < 0.001
Trained Random vs Trained Workplace
Cross-Condition
Workplace
4,100.0
p < 0.001
**
offers insights for developing behaviorally informed epidemic intervention protocols while enabling systematic evaluation of reward function design. It also provides form of advancing our understanding of how reward structure shapes emergent behaviors in complex spatial and temporal systems such as spatial epidemic models.
Broader Impact Statement. ContagionRL provides platform for researchers to study how reward design influences learned behaviour an epidemic setting. In broad sense, experiments performed using platforms of this type can help understand how potential public health intervention strategies, leading to perceived rewards avoiding penalties, might shape protective behaviours. Nonetheless, emphasize that human behaviour around disease avoidance is complex and heterogeneous, and that the relationship between simulation-oncoder rewards and real-world perceived rewards is also challenging to understand. Simulation findings will complement rather than replace epidemiological espertise and empirical validation for the foreseeable future.
13

---

Published in Transactions on Machine Learning Research (02/2026)
References
Michael P. Allen, Dominic J. Tildesley, and Dominic J. Tildesley. Computer simulation of liquids. Oxford science publications. Clarendon Pr, Oxford, reprinted edition, 2009. ISBN 9780198556459.
Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Concrete problems in ai safety. arXiv preprint arXiv:1606. 065665, 2016.
Mauricio Arango and Lyudmil Pelov. Covid-19 pandemic cyclic lockdown optimization using reinforcement learning. arXiv preprint arXiv:2009.0400 04647, 2020.
Kai Arulkumaran, Mare Peter Deisenroth, Miles Brundage, and Anil Anthony Bharath. A brief survey of deep reinforcement learning. arXiv preprint arXiv: 05866, 2017.
Jonas Balisacan, Monique Chyba, and Corey Shanbrom. Two new compartmental epidemiological models and their equilibria. medRaiv, 2021. doi: 10.1101 200..........22222..... 09. .t..03..22220 21263050.
URL https: //www. www.:wwddxiioooo
panteet//erlly/222//09//////
Shweta Bansal, Bryan T Grenfell, and Lauren Ancel Meyers. When individual behaviour matters: homogeneous and network models in epidemiology. Journal of the Royal Sociely Interface, 4(16):879 891, 2007.
Inan Bostanci and Tim Conrad. Integrating agent-basch and compartmental models for infectious discase modeling: A novel hybrid approach. Journal of Artificial Socielies ond Social Simulation, 28(1):5, 2025. ISSN 1460-7425.
doi: 10.18564/jasss.556) URL attps://wat wasss.org/////////////// html.
Fred Brauer, Carlos Castillo-Chavez, Zhilan Feng, el al. Mathematical models epidemiology, volume 32. Springer, 2019.
Sabah Bushaj, Xuecheng Yin, Arjeta Beqiri, Donald Andrews, and f. Esra Büyüktahtakmn. A simulation-deep reinforcement learning (sirl) approach for epidemic control optimization. Annals of Operations Research, 328(1):245 277, September 2023. ISSN 0254-3330, 1572-9338. doi: 10.10077/110479-022---8888 https://link springer. wttp://0..1000/10000000000000000 022-996-77.
URL
Nicolò Cangiotti, Marco Capolli, Mattia Sensi, and Sara Soutlle. A survey on Yyарunov functions for epidemic compartmental models. Bollettino dell'Unione Matematica Italiana, 17(2):241 257, June 2024. ISSN 1972-6724, 2198-2759. doi: 10.1007/<40574--23-00003 00368--.. URL https://link, https:///ink.sprrnngee springer. com/10.1007/
840574-023-00388-
Ayush Chopra, Esma Gel, Jayakumar Subramanian, Balaji Krishnamurthy, Santiago Romero-Brufau, Kalyan S Pasupathy, Thomas C Kingsley, and Ramesh Raskar. Deepabm: scalable, efficient and differentiable agent-based simulations via graph neural networks. arXiv preprint arXim:2110.044 2021.
Monique Chyba, Taylor Klotz, Yuriy Mileyko, and Corey Shanbrom. A look at endemic equilibria of compartmental epidemiological models and model control via vaccination and mitigation. Mathematics of Control, Signals, and Systems, 36(2):297 327, June 2024. ISSN 0932-4194, 1435-568X. doi: 10.1007 000-996023-003653 2. URL /ttpp:://ink... springers com 1 /ttm///00000//04400000000633
G. Dimarco, B. Perthame, G. Toscani, and M. Zanella. Kinetic models for epidemic dynamics with social heterogeneity. Journal of Mathematical Biology, 83(1):4, July 2021. ISSN 0303-6812, 1432-1416. doi: 10.1007/500255--21-011633
springer URL https://link: /ttp:ss//innk.ssiinneer com/10 http//..100./// 100.1000//00000 20001000/000000-001-0000000
Xinqi Du, Hechang Chen, Bo Yang, Cheng Long, and songwei Zhao. Hrl4ec: Hierarchical reinforcement learning for multi-mode epidemic control. Information Sciences, 640:119065, 2023.
Adnan M. S. Fakir and Tushar Bharati. Pandemic catch-22: The role of mobility restrictions and institutional inequalities in halting the spread of covid-19. PLOS ONE, 16(6))40253348, June 2021. ISSN 1932-6203. doi: /os/a///allll..oee...o.
URL http:://dxxxdo...oo http//.....1//jj/ 1371/journal pone. 0253348.
Maryam Farboodi, Gregor Jarosch, and Robert Shimer. Internal and external effects of social distancing in a pandemic. Journal of Economic Theory, 196:105293, 2021.
14

---

Published in Transactions on Machine Learning Research (02/2026)
Tao Feng, Tong Xia, Xiaochen Fan, Huandong Wang, Zefang Zong, and Yong Li. Precise mobility intervention for epidemic control using unobservable information via deep reinforcement learning. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 2882 2892, 2022.
Tao Feng, Sirui Song, Tong Xia, and Yong Li. Contact tracing and epidemic intervention via deep reinforcement learning. ACM Transactions on Knowledge Discovery from Dala, 17(3):1 24, 2023.
Neil M Ferguson, Daniel Laydon, Gemma Nedjati-Gilani, Natsuko Imai, Kylic Ainslic, Marc Baguelin, Sangeeta Bhatia, Adhiratha Boonyasiri, Zulma Cucunubá, Gina Cuomo-Dannenburgg et al. Report 9: Impact of non-pharmaceutical interventions (NPIs) to reduce COVID19 mortality and healthcare demond, volume 16. Imperial College London London, 2020.
Jakob Foerster, Nantas Nardelli, Gregory Farquhar, Triantafyllos Afouras, Philip HS Torr, Pushmeet Kohli, and Shimon Whiteson. Stabilising experience replay for deep multi-agent reinforcement learning. In PMLR, 2017.
International conference on machine learning, pp.
Vincent Frangois-Lavet, Peter Henderson, Riashat Islam, Marc G Bellemare, Joelle Pineau, et al. An introduction deep reinforcement learning. Foundations and Trends0 in Machine Learning, 11(3-4): 219-354, 2018.
Daan Frenkel and Berend Smit. Understanding molecular simalation: from algoritimss to applications. Computational science series. Academic press, San Diego, 2nd ed edition, 2002. ISBN 9780122673511.
Avner Friedman and Yhiu-Yen Kao. Mathemalical modeling of biological processes, volume 1. Springer, 2014.
Sebastian Funk, Marcel Salathé, and Vincent AA Jansen. Modelling the infiuence of human behaviour on the spread of infectious diseases: a review. Journal of the Royal Society Interface, 1(40))1247-1250, 2010.
Susan Gallagher, Asha Roy, Sandra Domeracki, Todd Mohrmann, Vicki Missar, Janet Jule, Shreela Sharma, and Ryan DeWitt. The low-wage essential worker: Occupational concerns and needs in the covid-19 pandemic round table. Workplace Health & Safely, 69(4):154 160, March 2021. ISSN 2165-0969. doi: 10.1177/2165079990988883 URL http://dx. doi. org/1000100000 10111117111100/ /171117777766777998822
Pierre-Yves Geoffard and Tomas Philipson. Rational epidemics and their public control. International economic review, pp. 603 624, 1996.
Mark Gersovitz and Jeffrey S Hammer. The economical control of infections discases. The Economic Journal, 114(492):1 27, 2004.
Volker Grimm and Steven F Railsback. Individual-based modeling and ccology. In Individual- modeling and ecology. Princeton university press, 2013.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pp.
Pmlr, 2018.
Herbert W Hethcote. The mathematics of infectious discasos. SIAM review, 42(4):599 653, 2000.
Elizabeth Hunter, Brian Mac Namee, and John D. Kelleher. A taxonomy for agent-based models in human infections discase opidemiology. Journal of Artificial Societies and Social Simulation, 20(3):2, 2017. ISSN 1460-7425. doi: 10. 111661//ssss.3411 URL http://jasse: sosasssssc.....................s soc. surrey a wk/.///////3//// / html.
Elizabeth Hunter, Brian Mac Namee, and John Kelleher. An open-data-driven agent-based model to simulate infectious disease outbreaks. PLOS ONF, 13(12))a0208775, December 2018. ISSN 1932-6203. doi: 10.1371 10.1371/journal) wttp:///wma.......//./...... URL https://dx.r plos wot.:o/aa..ooooa...pp.oo. 0208775.
Sinan Ibrahim, Mostafa Mostafa, Ali Jnadi, Hadi Salloum, and Pavel Osinenko. Comprchensive overview of reward engineering and shaping in advancing reinforcement learning applications. IEEE Access, 12:
175473 175500, 2024. doi:
15

---

Published in Transactions on Machine Learning Research (02/2026)
Ning Jiang, Weiqi Chu, and Yao Li. Modeling, inference, and prediction in mobility- compartmental models for epidemiology. ArXiv, abs/2406.12002, 2024. URL. https://app: wwa.p..semaniccccollaeee org/ 2775555:
Matt J. Keeling and Pejman Rohani. Modeling Infectious Diseases in Humans and Animals. Princeton University Press, 2008. ISBN 9780691116174. URL http://www.jeet jstor ott/ ctvcm4gk0.
William Ogilvy Kermack and Anderson G McKendrick. A contribution to the mathematical theory of epidemics. Proceedings of the royal society of london. Series A, Containing papers of mathematical and physicol characler, 115(772):700 721, 1927.
Harshad Khadilkar, Tanuja Gamu, and Deva P Scetharam. Optimising lockdown policies for epidemic control using reinforcement learning: An ai-driven control approach compatible with existing disease and network models. Transactions of the Indian National Academy of Engineering, 5(2):129-18) 2020.
Lingcai Kong, Jinfeng Wang, Weiguo Han, and Zhidong Cao. Modeling heterogeneity in direct infectious disease transmission in compartmental model. International Journal of Enwironmental Research and Public Health, 13(3):253, February 2016. ISSN 1660-4601. doi: //ww:.mdpi.com 0ttppp.ccmm160044411111110 1600-46011///3//255
URL https:
Gloria Hyunjung Kwak, Lowell Ling, and Pan Hui. Deep reinforcement learning approaches for global public health strategies for covid-19 pandemic. PloS one, (((5)e00211500 2021.
Joel Z Leibo, Vinicius Zambaldi, Mare Lanctot, Janusz Marecki, and Thore Graepel. Multi-agent reinforcement learning in sequential social dilemmas. arXiv preprint arXiv: arXiv:1702.0883 700.00337, 2017.
Yuxi Li. Deep reinforcement learning: An overview. arXiv preprint arXiv:1701.07777, 2017.
Picter JK Libin, Arno Moonens, Timothy Verstracten, Fabian Porez-Sanjines, Niel Hens, Philippe Lemey, and Ann Nowé. Deep reinforcement learning for large-scale epidemic control. In Machine Learning and Knowledge Discovery in Databases. Applied Data Science and Demo Track: European Conference, ECML PKDD 2020, Ghent, Belgium, September 14-18, 2020, Procsedings, Part V, pp. 155-170. Springer, 2021.
Anh Mai, Nikunj Gupta, Azza Abouzied, and Dennis Shasha. Planning multiple epidemic interventions with reinforcement learning. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, pp. 6147 6155, Macau, SAR China, August 2023. International Joint Conferences on Artificial Intelligence Organization. ISBN 9781956792034. doi: 10.24963 tijcai20223, /682. URL https: prgppproceedngg/2223//03 //www.1jcal ijcai. hwtptaaa......p...eee.ee
Piero Manfredi and Alberto D'Onofrio. Modeling the interplay between human behavior and the spread of infectious diseases. Springer Science & Business Media, 2013.
Susanna Manrubia and Damián Н Zanette. Individual risk- aversioer responses tune epidemics critical transmissibility (r= 1). Royal Society Open Science, 944)2216667, 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nalure, 518(7540):529 533, 2015.
Volodymyr Muih, Adria Puiglomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukeuoglu. Asynchronous methods for deep reinforcement learning. In Maria Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 1928 1937, New York, New York, USA, 20-21 Jun 2016. PMLR. URL https: //poceee//ccccsssssooooo mww.inng.m.......mm............m wwt..mles.r.....eeemeeem wwwpp///////m.//e.eee hwt.///ninn....mmmmmmmm h.w.ha1..hhe... html.
Seyed Sajad Mousavi, Michael Schukat, and Enda Howley. Deep reinforcement learning: an overview. In Proceedings of SAI intelligent systems conference, pp. 426-440. Springer, 2016.
Andrew Y Ng. Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In Teml, volume 99, pp. 278-287. Citeseer, 1999.
16

---

Published in Transactions on Machine Learning Research (02/2026)
Tai Nguyen, Phong Le, André Biedenkapp, Carola Doerr, and Nguyen Dang. On the Importance of Reward Design in Reinforcement Learning-bused Dynamic Algoritim Configuration: A Case Study on OneMar with (1+(A,)))-GA, pp. 1162 1171. Association for Computing Machinery, New York, NY, USA, 2025. ISBN 9798400714668 URL https://doi.org 10. 1145/3712256 3726395.
Joon Sung Park, Joseph O'Brien, Carrie Jun Cai, Meredith Ringel Morris, Percy Liang, and Michael S Bernstein. Generative agents: Interactive simulacra of human behavior. In Proceedings of the 36th annual acm symposium user interface software and technology, pp. 1-22, 2023.
Lerrel Pinto, Marein Andrychowicz, Peter Welinder, Wojeiech Zaremba, and Picter Abbeel. Asymmetric actor critic for image-based robot learning, 2017. URL htttps:/aarxvv.orr//ab////////// aataaass/11.... 1710. 066422..
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learwing Research, 22(268):1-8, 2021. URL attp://jmml org/ wwwgg/ppers///eeee http///222//2/// 1364.164.h. html.
Rubina F Rizvi, Kelly J Thomas Craig, Rezzan Hekmat, Fredy Reyes, Brett South, Bedda Rosario, William J Kassler, and Gretchen P Jackson. Effectiveness of non-pharmaccutical interventions related social distancing on respiratory viral infectious disease oulcomes: A rapid evidence-based review and meta-analysis. June 2021.
SAGE Open Med., 4:2050312121:1 121222977,,
Christopher Rose, Andrew J. Medford, C. Franklin Goldsmith, Tejs Vegge, Joshua S. Weitz, and Andrew A. Peterson. Heterogencity in susceptibility dictates the order of epidemic models. Journal of Theoretical Biology, 528:110839, November 2021. ISSN 00225193. doi: 10.1016//..itbi.221.1.111
URL https:
//Linkinghub. elsevier. http:///er..eoo//a/////////////0
Chadi M Saad-Roy and Arne Traulsen. Dynamics in behavioral epidemiological model for individual adherence to nonpharmaceutical intervention. Proceedings of the National Academy of Sciences, 120(44): e2311584120, 2023.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alee Radford, and Oleg Klimov. Proximal policy optimization algorithms. Xiv preprint arXio:1707.06344" 2017.
Ibne Farabi Shihab, Sanjeda Akter, and Anuj Sharma. Detecting and mitigating reward hacking in rein- forcement learning systems: A comprehensive empirical study, 2025. URL httpp::/aarti http::///axxvvvv/aaeee wwwaarxiv..ooraaaaaae argg/bss//00/
05619.
Alexander F. Siegenfeld, Pratyush K. Kollepara, and Yaneer Bar- Yam. Modeling complex systems: A case study of compartmental models in epidemiology. Complezity, 2022(1):3007864, January 2022. ISSN 1076-2787, 1099-0526. 1155/2022/3007864
doi: 00/55//2222/3077750 URL wnw.//bbarr www.wwwwwo/oooi//./.... com/doi/10.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Pannecrshelvam, Mare Lanctot, et al. Mastoring the game of go with deep neural networks and tree search. nature, 529(7588):484 489, 2016.
Tarini Sudhakar, Ashna Bhansali, John Walkington, and David Puelz. The disutility of compartmental model forecasts during the covid-19 pandemic. Frontiers in Epidemiology, 4:1389617, June 2024. ISSN 2674-1199. doi: 10.3389 http///pp..22//// 1389177. URL https: h//p::///ww... www.:/rrtatrrrsi.orrrrrriiiiiiie wttp://..............// 2024. 1389617/fu11
Richard S. Sutton and Andrew G. Barto. Reinforcement learning: introduction. Adaptive computation and machine learning series. The MIT Press, Cambridge, Massachusetts, second edition edition, 2018.
ISBN 9780262039246.
Ardi Tampuu, Tambet Matiison, Dorian Kodelja, Tlya Kuzovkin, Kristjan Korjus, Juhan Aru, Jaan Aru, and Raul Vicente. Multiagent cooperation and competition with deep reinforcement learning. PloS one, 12(4): e0172395, 2017.
17

---

Published in Transactions on Machine Learning Research (02/2026)
Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulao, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: A standard interface for reinforcement learning environments. arXiv preprint arXiv24477.17332, 2024.
Melissa Tracy, Magdalena Cerdá, and Katherine M. Keyes. Agent-based modeling in public health: Current. applications and future directions. Annual Review Public Health, 39(1):77-94, April 2018. ISSN 0163-7525, 1545-2093. doi: 10.1146 / aawmuuewwwpuillllell---00000-000 URL
content/journals/LE 10. 1146/ puw.uurrryyplllaaattt0000------0
Runzhe Wan, Xinyu Zhang, and Rui Song. Multi-objective model-based reinforcement learning for infectious discase control. In Proceedings of the 27th ACM SIGKDD Conference Knowledge Discovery and
Data Mining, pp. 1634 1644, Virtual Event Singapore, August 2021. ACM. ISBN 9781450383325, doi:
0441144//44754484467777 URL nttps://dl acm wtt../o////o11.11/5/////200 115534444111534444554554444 3467333
Bing Wang, Zeyang Xie, and Yuexing Han. Impact of individual behavioral changes on epidemie spreading in time-varying networks. Physical Review E, 104(4):044307, 2021.
Zhen Wang, Chris Bauch, Samit Bhattacharyya, Alberto d'Onofrio, Piero Manfredi, Matjaž Perc, Nicola Perra, Marcel Salathé, and Dawei Zhao. Statistical physics of vaccination. Physics Reports, 664:1-113, 2016.
Edwin B Wilson and Jane Worcester. The law of mass action in epidemiology. Proceedings of the National Academy of Sciences, 31(1):24-34 24, 1945.
Yang Ye, Qingpeng Zhang, Zhongyuan Ruan, Zhidong Cao, Qi Xuan, and Daniel Dajun Zeng, Effect of heterogeneons risk perception on information diffusion, behavior change, and disease transmission. Physical review E, 102 (4):044443 2020.
Cameron Zachreson, Sheryl Chang, Nathan Harding, and Mikhail Prokopenko. The effects of local homogemeity assumptions in metapopulation models of infectious disease. Royal Society Open Science, 9(7):211919, 2022.
18

---

Published in Transactions on Machine Learning Research (02/2026)
Appendix
A Compartmental Epidemiology Models Continued
20
A.1 Notation and parameter definitions
20
A.2 Model equations
21
A.3 Comparison of the opidemiological models
21
B Environment Visualization
22
C Methodology Continued
22
C.1 Toroidal Distance Caleulation
22
C.2 Environment Configuration Details
................................
23
C.3 Reward Function Implementation Details
23
C.4 Reward Function Information Scope
25
C.5 Algorithmic Comparison Hyperparameters
26
C.6 Baseline Implementation Details
27
D Results Continued
29
E Learned Adherence Behavior Analysis
32
F Partial Observability Feature Analysis
33
G Compute Resources
33
H Additional Results and Experiments
34
19

---

Published in Transactions on Machine Learning Research (02/2026)
A Compartmental Epidemiology Models Continued
A.1 Notation and parameter definitions
We denote by the numbers of susceptible, infected, recovered, exposed and deceased individuals at time using S(t), I(t), R(t), E(t), D(t), respectively. In the SIR, SIRS and SIRS+D models the total living population is
N(L) 5(1)+((+)))) R(1),
(3)
whereas in the SEIR model
N(1) S(t) E()) + I()) + R(t),
(4)
The model parameters are:
3 transmission rate (per-contact rate at which susceptibles become infected);
7 recovery rate (rate at which infected individuals recover);
σ progression rate (rate at which exposed become infectious, SEIR only):
immmity-wwaing rate (rate at which recovered lose immunity, SIRS and SIRS+D);
/ diseaso-induced mortality rate (rate at which infected die, STRS+D only).
7
SIR:
R
8
7
SIRS:
S
R
3
SEIR:
E
R
SIRS+D:
7
R
D
Figure 7: Compartmental epidemic models visualized: SIR, SIRS, SEIR, and SIRS+D.
20

---

Published in Transactions on Machine Learning Research (02/2026)
A.2 Model equations
The differential equations below provide overview of how compartments are tracked and changed population- wide based on the defined respective parameters in Appendix A.1.
SIR model.
dS SI N 8 ST di N I, dR dt I I.
dt -3
(5)
dI
(6)
(7)
SIRS model.
dS SI N di SI
di 3
+
(8)
dt N 7 1 dR I R. dt
(9)
(10)
SEIR model.
dS SI dt B N dE SI
(11)
E,
(12)
dl dl dt F. I dR
(13)
dt I1
(14)
SIRS+D model.
dS
dt = 3 SI
+ HER,
(15)
N
dt = 3 S1
(y + 1)
(16)
N
dR
dt & R dD
(17)
di pIl
(18)
A.3 Comparison of the epidemiological models
Table 4: Comparison of epidemic models and their capacity for individual-level and adaptive behavior modeling
Comportents
Reinfeetion
Behusioral Modeling
Adoptisn
(Komaadd McKondrick, 1927)
Noue
SIRS (Kemmads & Mekendrisk, 1927)
L.L. R 8. R.
None
No
SE3R [Kimmack & MeKenstrick, 1927; Wileon à Wancester, 1945)
Yos
Naa
No-
50RS +D (Wilka & Wonssler, 1945)
8. R.
Yos
No
Yos
Nata:
No
AUM (Gommm Ralidbadi, 2013)
Custom
Opticnal
Preseribed rules
Purtial
ContagionILL
Modified
Learned policies (RL)
21

---

Published in Transactions on Machine Learning Research (02/2026)
B Environment Visualization
Figure shows single frame of the environment's rendcred vidco. The visualization provides comprchensive and intuitive overview of the simulation state and the environment at each timestep. The panel al the center shows the environment's grid, in which each human is represented as a colored circle depending on its cpidemiological state. The legend shown on the left of the grid (blue for susceptible, red for infectious, green for recovered, and gray for dead). The focal agent, whose actions are being learned by RL, shown as orange dot. with dark border. The small arrow, originating from the agent, shows its most recent. movement. direction and magnitude.
The upper panel on the right of the grid, shows the details related to the agent such current state, position, adherence level, and cumulative reward. While the panel below it shows the population distribution across all epidemiological states in numbers and absolute percentages (due to increased need of processing power for rendering this function is turned off by default).
50
20/912
Agent
Buscopitbll
40
Foetion
(23.r 24.5)
Movement
Monemrnn
Adhercnnee
0.58
Susoepllbbee Recovered Agent
Coordinats
Category
(Percentage)
Busceptible
25 (92.5%)
Infectious
12 (%)
Recovered
3.07.5%
Dead
(8.0%)
Figure 8: Sample render of the SIRS-D environment at step 30 of an episode
C Methodology Continued
C.1 Toroidal Distance Calculation
Throughout the simulation, distances between individuals calculated using the standard Euclidean distance metric, modified to account for the periodic boundary conditions of the toroidal grid. This ensures that the distance represents the shortest path, considering potential 'wrap-around" across grid edges.
The distance dii.j) between two individuals and j at positions (xr, and (æj, (5) respectively, within a grid of side length G, is calculated as:
d(i,j) / A wrap + A
(19)
where the wrapped differences along each axis are determined by:
ATwrap min(12; -zyl G las æzl)) Aywrap min(ly: 3y/ G 11y yy
(20)
22

---

Published in Transactions on Machine Learning Research (02/2026)
This formulation ensures that the distance along each dimension corresponds to the shorter of the two possible paths (direct or wrapped around the grid). This distance metrie is fundamental to calculating infection exposure, determining the nearest threat for the Greedy policy, and computing forces in the potential field reward function.
C.2 Environment Configuration Details
For completeness and to ensure the reproducibility of our experiments, this section provides detailed breakdown of the specific parameter values used to configure the STRS+D simulation environment described in section 3. Table 5 lists each parameter, its corresponding symbol used within the text (where applicable), its configured value, and a brief description of its role within the simulation dynamics. These paramcters define the epidemiological characteristics, spatial properties, and agent interaction rules governing the environment. used throughout our study.
Table 5: SIRS+D Environment Configuration Parameters
Parameter
Value
Type
Ronge
Destription
Time
312
Integes
[1,00)
Maxtrum number timostepe per opisodo.
Grid
30
Integee
[,,0)
Side lengt' of the square. torocial gnd (Volue Value).
Number Bumaus
40
infogrt
(1.00)
Number non opest iulividucas in the simulation.
Initisl Seseed
10
integge
[0.10)
Numberr initially infectdd individus (up to Number of Humomn)).
Infiction Rate
0.5
Boad
[0.
Buse intetiko inmesissiu rale
unil.
Initial Agrnt Adherence
0
flot
Agme's NPI adhermmce level the
an eposdd.
Pepureniid dwexy futtrr for
expoccrn lased
Lehality
[0
Per step poobibiity death for Infected inditiinals.
Immunity Loes Probabciit
0.25
Bont
1x 1]
Pre-step prolebciry transttionng from Recoureed (R) to Stscepetble (S).
Rooneyy Hate
0.1
Boat
[0
Pre peolubiitty of trazattionng from Infected (1) to Resovered (K).
Adberence Penalty Pictor
1
Boot
(1,00)
Multipllative factor for adherrcaee cost (tood) by some reward fumetions).
Allberenc Ellectivataoe
0.2
Baat
[0
Rosiluad testsem aoim facloe whsn sgent aillxrmux is 1.0 (egg, meaz 20% lala nomaine).
Movereent
continocur_roe randon
Categorical
Movement model for non-agett todendunle (o.g, consinuons raadonn).
Moseemnn
step
nun-agen individuel.
Visibiiity Badins
[-1] U oo)
's observation todras indicates grs vidlidity).
Retnfoction Count
5
Integge
00 oc)
Number S midonduule relnfected if I,
Sate Distaoe
10
Boat
(c)
Min distazoe for (relideecing S individuais reative to the sgrat.
Initial Agent Distmnes
5
Boot
(0)
Min distazee all individuals are placed from the agent Initialiaaiion.
Mas: Infection Distamor
10
Boat
(-1) oc)
Max dicance Ice in intividual to contrinute expoome ke infiaia)).
Rownnd Function Type
potential__te _feld
Categorica
Specifies roward finetton (egg. potential_ ffeld, sparse).
Ieentier
Ablation
Oategarica'
fenction ablicion strelisee (n-S) 1el1,, th)) Bendeting vinulization.
Bondec Made
None
staiin NoneType
(ego, arroy, Nozo]
C.3 Reward Function Implementation Details
This appendix provides the detailed mathematical formulations for key reward functions implemented in the CONTAGIONRL, environment. The goal of the agent, in the context of these reward functions, is maximize the cumulative reward, which designed to correlate with prolonging its susceptible state.
C.3.1 Constant Reward
This is the simplest reward structure, providing a binary signal based on the agent's health state. Let Sa denote the state of the agent. The reward Reonat at each timestep is defined as:
if Sg = Susceptible (1
Reonst
(21)
0 otherwise
C.3.2 Reduce Infection Probability Reward
This reward function directly incentivizes the agent to minimize its own probability of becoming infected. Let (a) be the agent's 's probability of infection at the current timestep, calculated considering its position, adherence, and proximity to infected individuals. The reward Rreduce inf is defined as:
( (1 - Pintect (a))2 if S. Susceptible
Ryeduce_ lnf
(22)
-5
if S. + Susceptible
The quadratic term (1- Pinfect (a))2 strongly rewards states with very low infection probability.
23

---

Published in Transactions on Machine Learning Research (02/2026)
C.3.3 Combined Reward
This function combines the incentive to reduce infection probability with a small constant bonus for remaining susceptible. Using (of as defined above, and Reonst (which is susceptible, 0 otherwise from the agent's perspective in this formulation if it's not infected):
{o.s
0.8 (1 Piffcct (a))² 0.1 S. = Susceptible
Brombo {
(23)
if S, Susceptible
The 0.1 term derives from 0.1 x Roonst where Reonst = 1 when the agent is susceptible.
C.3.4 Maximize Nearest Distance Reward
This reward encourages the agent to maintain certain distance from other individuals, primarily those who are susceptible infected. Let dmin be the Euclidean distance from the agent to the nearest human h whose state Sд E (Susceptible, Infected]. Let Da be the parameter maximum distance for beta calculation from the environment configuration. The reward Rmax dist is defined as:
0
if S. Infected
1
if S. / Infected and no relevant humans (S or I) exist if S. * Infected and dmia 2 D3 (and D3 > 0)
Rmax dist
(24)
1 max(0, dwin/ Dg)
if S. * Infected and dmin < Dg (and Dg > 0)
If D3 (interpreted as infinite or disabled threshold in some contexts, though the code uses it a divisor if positive), the behavior might implicitly rely normalization other logic not explicitly captured by this simplified D3 > 0 case. The provided code primarily gives reward based on dmin/Dg when dmin Dg and D3 0.
C.3.5 Potential Field Reward
The Potential Field reward function conceptualizes agent- human interactions system of forces, guiding the agent's movement based on the proximity and state of other individuals. It is composite reward with several terms:
Rpp = Whealth/health + "aatheroncnaadhhremaa + Wmovement move
(25)
In the implementation, the weights Whealth 0.1, Wadherence 0.2, and Wmovement 0.7.
1. Health Reward (Thealth): This is a binary reward for maintaining susceptible state.
[1 if S. Susceptible {
Thealth
(26)
otherwise
(This component can be ablated to O via the no health variant).
2. Adherence Reward (radherence)): This reward is directly proportional to the agent's NPI adherence level. Let aa be the agent's current adherence (a value [0, 1)).
(27)
Fadherence
(This component can be ablated 0 via the no adherence variant).
3. Movement Reward (rmove): This component rewards the agent for moving in alignment with suggested force vector F and optionally for matching its magnitude. Let the agent's current position be
= (Eas Va) and human j's position be P, (Iz, Mj). The shortest displacement vector on the toroidal grid from human j to the agent is Ap; = (Aæ,, Aw;), where:
Ar, (s. 2j G/2)
(mod G) G
(28)
Ay, (yo yj + G/2) (mod G) - G/2
(29)
24

---

Published in Transactions on Machine Learning Research (02/2026)
Here, G is the grid size. The squared distance is df (Ar,)² + (Ay,)² Edist, where Edist is small constant to prevent division by vero.
The force contribution from each human j depends on its state:
Wf
human Infected (1)
weighty
Ws if human is Susceptible (S) (can be 0 in Susceptible ablation)
(30)
0
otherwise
The scaling factor for the force from human j is scale; weight, (d p//2 The implementation uses 1, W, = 1.0, Ws = 0.5. The resultant force vector is F (Fz, F.y = Ej scale, The normalized force direction is = / + Enorm).
Let the agent's chosen movement voctor from the last action be (ax, ay). The directional alignment reward is:
a - 1,
Ydir = clip
(31)
||a
norrm
The magnilude matching reward is:
Tmag clip(1 |||||| mi]( 1.0))
(32)
The combined movement reward is (by default, with 3m = 0.25))
Tmove (1 Bm)rdir 3mmmag
(33)
This term is subject to ablations:
No Magnitude: Tmove Tdir
No Direction: Tmove Tmag
No Movement: Tmove 0
The small constants Cdist and Cnorm (10-8) are used ensure numerical stability.
C.4 Reward Function Information Scope
The summation in F iterates over all humans in the simulation irrespective of the agent's visibility radius parameter. Consequently, under POMDP configurations where the agent's observation is restricted to radius the force vector still incorporates contributions from humans beyond The observation function zeroes out features for humans outside the visibility radius, but the reward computation does not apply this filter. This asymmetry means the directional reward acts as an oracle-like gradient signal during training, informing the agent whether its action aligned with the globally optimal escape direction even when the agent could not observe all contributing humans. This is consistent with asymmetric actor-critic approaches where training-tine information may exceed test-time observations (Pinto et al., 2017).
25

---

Published in Transactions on Machine Learning Research (02/2026)
Table 6: Information scope of each reward function and its components. Local terms depend only on the agent's own state or action. Globol terms are computed from non-agent humans regardless of the agent's visibility radius (under the default MDP configuration with visibility_raceeeee _rciuus -1).
Reward Function
Component
Scope
Depends On
Constant
Alive bonus
Local
Agent state (Sa)
Alive check (1 Pint)²
Local
Agent state (S.)
Reduce Infection
Global
Agent pos., adherence, all infected positions
Alive bonus
Local
Agent state (S.) Agent state (S.)
Const. + Red. Inf.
Alive check
Local
(1 Pint)
Global
Agent pos., adherence, all infected positions
Alive check
Local
Agent state (S.)
Max Near. Dist.
dmin/ Da
Global
Agent pos., all $ & positions
Health (Fhealth)
Local
Agent state (Sa)
Adherence (Fadherence)
Local
Agent action (a)
Potential Field
Direction (raie)
Global
Agent action, all positions
Magnitude (Tmag)
Global
Agent action, all & I positions
C.5 Algorithmic Comparison Hyperparameters
The following tables summarize the hyperparameters used for the Proximal Policy Optimization (PPO), Soft Actor-Critic (SAC), and Advantage Actor-Critic (A2C) algorithms in our comparative study (Figure 2).
Table 7: PPO Hyperparameters
Hyperparameter
Value
Policy Type
Multi IntutonpyyPlicy
Policy Kwargs:
Net Arch (pi)
[256, 256, 256, 256]
Net Arch (vf)
[256, 256, 256, 256]
Activation Fn
ReLU
Ortho Init
True
Batch Size
2048
Steps
1024
N Epochs
Learning Rate
3e-4
Gamma
0.96
GAE Lambda
0.95
Target KL
0.04
Clip Range
0.2
Ent Coef
0.02
Normalize Advantage
True
Total Timestops
8,000,000
N Envs
4
26

---

Published in Transactions on Machine Learning Research (02/2026)
Table 8: SAC Hyperparameters
Hyperparameter
Value
Policy Type
MlpPolicy
Policy Kwargs:
Net Arch (pi)
[256, 256, 256, 256]
Net Arch (qf)
[256, 256, 256, 256]
Activation Fn
ReLU
Learning Rate
3e-4
Buffer Size
1,000,000
Batch Size
256
Tau
0.005
Train Freq
1
Gradient Steps
Ent Coef
auto
Gamma
0.96
Total Timesteps
8,000,000
N Envs
Table 9: A2C Hyperparameters
Hyperparameter
Value
Policy Type
MipPolicy
Policy Kwargs:
Net Arch (pi)
[256, 256, 256, 256]
Net Arch (vf)
[256, 256, 256, 256]
Activation Fn
ReLU
Ortho Init
True
N Steps
640
Gamma
0.96
GAE Lambda
0.95
Ent Coef
0.01
VF Coef
0.5
Max Grad Norm
0.5
Learning Rate
3e-4
Use RMS Prop
True
Normalize Advantage
True
Total Timesteps
8,000,000
Envs
C.6
Baseline Implementation Details
This section provides further detail on the implementation of the non-learning baseline policies used for comparative analysis, as described Section
C.6.1 Stationary Policy Implementation
The Slationary policy represents a minimal activity baseline. Its implementation is deterministic: at every timestep within an episode, the agent selects constant action vector a(t) = [0.0, 0.0, 0.0.. This corresponds to zero intended movement (Ar 0, Ay = 0) and zero NPI adherence (a = 0). The agent's position remains fixed, and its susceptibility is solely determined by the base infection rate and its exposure to infected individuals based the environment. dynamics.
27

---

where Published in Transactions on Machine Learning Research (02/2026)
C.6.2 Random Policy Implementation
The Random policy simmlates behavior devoid of environmental feedback or strategy. At each timestep t, the action a(t) = (Ar, Ay, a) is generated by sampling independently from uniform distributions:
Movement components Ar, Ay ~ 4(1111)
The adherence level a (O(0))).
This ensures that the agent explores the action space randomly without any directed objective, using the simmlation's internal pscudo-random number generator for consistency.
C.6.3 Greedy Distance Maximizer Policy Implementation
The Greedy policy emacts a deterministic, reactive heuristic focused on maximizing immediate distance from the nearest identified threat. The implementation logic proceeds as follows each timestep
Identify Nearest Threat: Let the agent's position at time be Pa(() = (4), Ma (1)). Identify the set of currently infected humans I(t). This step utilizes privileged knowledge of the true state S, E {S, R, D, Dj for all humans j. If I(t) is empty, the agent defaults to a stationary action m (0,0).
2. Target Selection: If I(t) is not empty, calculate the Euclidean distance d(pa (t), (t), (t)) for all j E I(t), where (t) is the position of infected human j and d(-,-) accounts for the toroidal grid geometry. Identify the single infected human Anearest (1) corresponding to the minimum distance:
Anearest (t) aag min (t). jEI(()
The exact position Piacarea (1) is privileged information.
Evaluate Potential Moves: Define a discrete of candidate movement vectors Amove- This sot includes the zero vector (0,0) and scaled unit vectors representing the maximum possible step in the eight cardinal and diagonal directions, e.g., {(0, 0), (#SM, 0), (0, SM is the maximum movement scale (typically 1.0).
Select Best Move: For each candidate movement m; = (Az,, E Amover calculate the agent's potential next position (t, + 1) by applying the movement to Pa (t) and considering the grid's periodic boundarics. Evaluate the distance from this potential position to the initially identified nearest. threat's current position Phacarca (l). Select the movement vector m² that maximizes this distance:
m
= arg max d(p'a,(t + 1), Preerret (t))
Ammovv
This calculation relies on knowing Phamarare (t).
Set Adherence: The adherence component of the action is deterministically set to the maximum value, 1.0.
6. Final Action: The resulting action for timestep / is a(l) (m", 1.0).
This implementation defines a simple, reactive strategy that exploits complete and accurate environmental state information maximize instantaneous separation from the nearest perceived threat, while also employing maximum protective adherence.
28

---

Published in Transactions on Machine Learning Research (02/2026)
D Results Continued
These figures and tables provide additional visualizations and experiments corresponding to those presented in the main Results section of the paper.
200
175
Mean Episodd Duration
Stationary
Random
Greedy
Algorithm
Figure 9: Mean episode durations and 95% bootstrapped confidence intervals (10,000 samples) for each agent. Values are aggregated across 3 training seeds and 100 evaluation episodes per seed. This figure summarizes the performance trends visualized in Figure 2.
Mean Episode Duration
125
Constant
Reduos Infection
Constant Reduce Inf. Reward Function
Max Nearest Distance
Potential Field
Figure 10: Mean episode durations with 95% bootstrapped confidence intervals (10,000 samples) for each reward function. These values correspond to the experiments shown in Figure 3, aggregating performance across 3 seeds and 100 evaluations per seed.
29

---

Published in Transactions on Machine Learning Research (02/2026)
Table 10: One-sided Mann Whitney U test results comparing each reward function to the Potential Field baseline (PF). Bonferroni correction applicd one-sided p-values. "Winner" denotes the reward function with significantly longer episode durations. Abbreviations: Const Constant, RedInfect Infection, Combo Constant Reduce Infection, MaxDist Max Nearest Distance, PF Potential Field. Significance: p < 0.05, p < 0.01, *** p 0.001, n.s. not significant. This table supports the significance annotations shown in Figure 3.
Reduce
Reward Function (2-sided)
(1-sided)
Sig (2)
Corrected (1)
Sig (1)
Winner
Const
3.38e-06
1.69e-06
6.76e-06
PF
RedInfect
0.00556
0.00278
**
0.0111
PF
Combo
0.00185
0,00092
0.00369
PF
MaxDist
1.04e-06
5.21e-07
***
2.08e-06
PF
200
Mean Episode Duration
No Magnitude
No Direction
No Movement
No Adherence
No Health
No Suscepiible Repulsion
Full Potential
Ablation Variant
Figure 11: Mean episode durations and 95% bootstrapped confidence intervals (10,000 samples) for each ablation variant in the Potential Field reward study. These aggregate values summarize the same evaluations shown in Figure 4, providing a direct comparison of central performance lendencies across reward configura- tions.
Table 11: One-sided Mann-Whitney U test results comparing each ablation variant to the Full Potential Field (FPF), which represents the original, unmodified reward function. Bonferroni correction is applied to one-sided p-values. "Winner" indicates the variant with significantly longer episode durations. Significance levels: 0.05, ** p 0.01, *** < 0.001, n.s. = not significant. Results support the differences shown in Figure 4.
Ablation Variant
(2-sided)
(1-sided)
Sig (2)
(1)
Corr
Sig (1)
Winner
No Magnitude
0.9571
0.5216
1.S..
I.S.
Direction
3.80e-11
1.90e-11
*8*
1.14e-10
44*
FPF
No Movement
6.93e-06
3.46e-06
2.08e-05
FPF
Adherence
3.88e-09
1.94e-09
***
1.16e-08
***
FPF
No Health
0.8242
0.5881
Susceptible Repulsion
0.1336
0.0668
n.s.
0.4007
n.S.
30

---

Published in Transactions on Machine Learning Research (02/2026)
Table 12: Statistical significance tests comparing RL agents trained with different visibility radius constraints in epidemic control. Agents were trained and evaluated with Full Visibility, Limited Visibility r=10, r=15, and r=20 (where r represents the radius within which the agent can observe infected individuals). Each condition used 300 episodes (3 sceds 100 episodes per seed). Statistical significance determined using Mann-Whitrer U lests with Bonferroni correction for 18 multiple comparisons (corrected a 0.002778). Tests compare trained agents against baselines (Stationary, Random, Greedy) and trained agents against each other to assess the impact of visibility constraints on performance.
Category
Comparison
Winner
U-stat.
p-value
Sig.
Full vs Stationary
Full
79,732.0
p < 0.001
***
Full vs Random
Full
66,639.0
p < 0.001
***
Full vs Greedy
Full
46,927.0
p = 1.000
n.8.
r= 10 vs Stationary
10
82,063.0
p < 0.001
1= 10 Random
r= 10
73,277.0
p 0.001
Trained vs
r= 10 vs Greedy
10
55,948.0
p < 0.001
Basclines
r=15 15 vs Stationary
1=15
83,000.0
P < 0.001
***
r= 15 vs Random
r-
74,393.0
p < 0.001
$**
r= 15 VS Greedy
r=15
59,913.0
P < 0.001
***
vs Stationary
r=20
82,436.0
p < 0.001
I=20 vs Random
r=20
72,773.0
P < 0.001
20 Greedy
r=20
57,229.0
P < 0.001
Full vs r=10
1=10
35,009.0
P < 0.001
$**
Full vs r=15
r=
31,504.0
p < 0.001
#**
Trained vs Trained
Full vs r=20
r=20
34.320.0
p < 0.001
r= 10 1=15
1=15
40,943.0
p 1.000
1=10 r-20
r=20
43,770.0
P 1.000
15 vs 2-
1=
47,696.0
P 1.000
n.8.
31

---

Published in Transactions on Machine Learning Research (02/2026)
E Learned Adherence Behavior Analysis
To investigate adherence usage of trained models across different reward formulations, we conducted an analysis of the learned adherence behaviors. This analysis examines if agents can independently discover the value of non-pharmaceutical intervention (NPI) compliance or if they require explicit reward tuning for discovering the importance of adherence.
Methodology. We measured the mean and standard deviation of adherence actions from trained policies across all reward functions and potential field ablation variants. Each policy was evaluated using the same experimental setup the main results. Adherence values were recorded at each timestep and aggregated across episodes.
Results. Table 13 shows the Icarned adherence usage across reward functions. Agents trained with simple reward functions like Constant, Reduce Infection, and Constant + Reduce Infection achieve maximal adherence (a= = 1.0) with minimal variance. This indicates that even without explicit adherence incentives, temporal credit assignment enables agents to discover the survival benefit of full NPI compliance through the sparse reward signal. The Max Distance reward function yields slightly lower mean adherence with higher variance, suggesting that spatial optimization partially competes with adherence learning in this formulation.
The Full Potential Field reward achieves maximal adherence, consistent with its explicit adherence component. However, when the adherence reward is removed from potential field reward function (PF: No Adherence), mean adherence drops to 0.41 with high variance of 0.47. This indicates essentially random adherence behavior. All other potential field ablations (No Direction, No Health, No Magnitude, No Movement and No Susceptible Repulsion) maintain maximal adherence despite their respective component removals.
These results reveal an asymmetry in adherence learnability based on reward complexity. Simple reward functions allow agents to implicitly learn maximal adherence through the survival signal alone. However, in complex multi-component objectives like the potential field reward, agents fail to discover adherence value without explicit incentivization. The diverse reward signals in the potential field formulation appear to dilute the implicit survival gradient that simpler rewards preserve.
The ContagionRI. platform includes configurable adherence penalty parameter (Table 5) to enable the future investigation of strategic trade-offs where agents must balance survival benefits against adherence costs.
Table 13: Learned adherence behavior across reward functions and potential field ablations. Mean and standard deviation computed over 300 evaluation opisodes sceds 100 episodes). PF denotes Potential Field.
Reward Function
Mean Adherence
Std Adherence
Constant
1.000
0.000
Reduce Infection
1.000
0.000
Constant + Reduce Inf
1.000
0.000
Max Distance
0.872
0.312
Potential Field
1.000
0.001
PF: No Adherence
0.406
0.472
PF: No Direction
1.000
0.000
PF: No Health
1.000
0.000
PF: No Magnitude
1.000
0.000
PF: No Move
1.000
0.000
PF: No S Repulsion
1.000
0.001
32

---

Published in Transactions on Machine Learning Research (02/2026)
F Partial Observability Feature Analysis
To further understand the counter-intuitive result that limited visibility agents outperform full visibility agents (Figure 5), we conducted a feature importance analysis measuring how strongly each trained model responds to observed human positions.
Methodology. We computed Spearman correlations (p) between the mean relative positions of visible humans and the agent's 's movement actions. Specifically, we measured correlations between the mean x- displacement of observed humans and the agent's x-movement (paz), and similarly for the y-dimension (Pay).
Results. Table 14 presents the action-position correlations by visibility radius. All correlations are negalive, indicating avoidance behavior where agents move away from observed humans. Limited visibility models exhibit substantially stronger correlations than full visibility models. Agents trained with r=10 visibility show correlations approximately 2-2.5x stronger than full visibility agents (p 0.001), demonstrating more decisive responses to observed threats.
Interpretation. The stronger correlations under limited visibility occur because full visibility dilutes attention across all humans, many of whom are distant and pose negligible infection risk. Limited visibility effectively filters the observation space to include only nearby individuals who have a larger impact on the effective transmission rate since transmission follows distance decay.
Limited visibility thus acts as an implicit attention mechanism, restricting observations to the radius where humans actually pose significant threat. Full visibility forces agents to process information about distant humans who have near-zero impact infection risk, increasing observation dimensionality without proportional informational benefit. These results suggest that appropriately scoped observations matched to the task's depeeneen mechanics produce more focused and effective policies, rather than requiring more expressive architectures to handle larger state spaces.
Table 14:
Spearman correlations between mean relative positions of visible humans and agent movement actions across visibility conditions. Negative values indicate avoidance behavior. All correlations significant at p 0.001.
Visibility
PAr
PAy
Limited (r=10)
-0.557
-0.540
Limited (r=15)
-0.441
-0.464
Limited (r=20)
-0.342
-0.438
Full Visibility
-0.220
-0.234
G Compute Resources
All experiments were conducted on Apple Silicon M4 chip workers (16-core CPU) with 48 GB LPDDR5 unified memory. We devoted approximately 40 CPU-hours to hyperparameter optimization. All training and evaluation were performed across 3 random seeds (with 100 runs per seed unless otherwise specified). Training times for each set of experiments were approximately as follows:
Comparison of different RL algorithms: 22 hours
Ablation studies of the potential field reward function: 13 hours
Comparison of reward functions: 14 hours
Comparison of POMDP and full observability (MDP): 15 hours
Constraints in Movement: 12 hours
Additional results and experiments (in appendix): 26 hours
33

---

Published in Transactions on Machine Learning Research (02/2026)
500
400
Mean Episode Duration (steps)
Agent Type Stationary
800
Random
Trained
Greedy
200
0.2
B
B 0.7
B 0.9
Infection Rate (.3)
Figure 12: Comparison of mcan episode durations across agents (Trained, Stationary, Random, Greedy) for different infection rates 3. Bars indicate the mean episode duration over 3 seeds (each evaluated for 100 episodes, totaling 300 episodes per agent), with 95% bootstrapped confidence intervals (10,000 resamples). The dashed red line marks the maximum possible episode duration. See Table 15 for detailed two-sided and one-sided Mann-Whitncy U test results with Bonferroni correction.
H Additional Results and Experiments
In the supplementary experiments presented here, we vary one environmental parameter of the epidemic time to simulate range of spatial epidemic scenarios constructible with CONTAGIONRI. This allows to explore both favorable and challenging conditions under which to stress-test agent learning, For each experiment, we provide two outputs: grouped bar plot illustrating mean episode durations with 95% confidence intervals, and a table reporting results from one-sided and two-sided Mann-Whiten Whitney U tests. All experiments in this section use PPO for training, with hyperparameters defined in Appendix C.5 and environment configurations Appendix C.2. Each table and figure corresponds to single modified environmental parameter specified their respective captions.
34

---

Published in Transactions on Machine Learning Research (02/2026)
Table 15: Results of two-sided and one-sided Mann Whitney U tests comparing Trained agents to baselines (Stationary, Random, Greedy) across four infection rates (3). Each comparison uses 300 evaluation episodes (100 per seed over 3 seeds). The one-sided p-values are Bonferoni-corrected within each group. "Winner" indicates the agent with significantly longer episode durations based on the corrected one-sided test. Significance levels: p << 0.05, ** < 0.01, *** p < 0.001, n.s. = not siguificant.
Baseline
p (2-sided)
p (1-sided)
Sig (2)
p (1) Corr
Sig (1)
Winner
Stationary
2.86 x 10-17
1.43 x 10 17
*4*
4.30 x 10-17
***
Trained
0.2
Random
6.87 10 24
3.43 10-24
***
1.03 10-23
***
Trained
Greedy
2.03 x 10-21
1.02 10-21
3.05 x 10
Greedy
Stalionary
1.01 x 10-16
5.06
10-17
1.52 x 10 10-16
Trained
0.5
Random
5.13 x 10-15
2.57
10 10-15
7.70 x 10-15
***
Trained
Greedy
1.29 x 10-46
6.45
10-47
1.94 x 10
Greedy
Stationary
9.82 10
4.91
10
1.47 10-15
Trained
0.7
Random
2.37 x 10
1.18 x 10
3.55 X 10 12
Trained
Greedy
1.56
10 10-27
7.79
10-28
***
2.34
10 10-27
***
Greedy
Stationary
5.46
27
2.73
10 27
***
8.19
27
****
Trained
0.9
Random
2.53 x 10
1.27
10
3.80 X 10 10-15
Trained
Greedy
6.74 x 10 23
8.37 X 10 23
1.01 X -2
***
Greedy
600
Mean Episode Duration (steps)
I
Agent Type
800
Stationary
Random
Trained
Greedy
200
100
Grid 30 x 30
Grid 40 x 40
Grid 50 50
Grid 60
Grid Size
Figure 13: Mean episode durations for Trained, Stationary, Random, and Greedy agents evaluated across environments with varying grid sizes. Each bar represents the mean across 3 seeds (cach with 100 evaluation episodes), and is accompanied by 95% bootstrapped confidence intervals (10,000 resamples). The red dashed line marks the maximum episode length. See Table 16 for statistical comparisons between Trained and baseline agents using both one-sided and two-sided Mann Whitney U tests with Bonferroni correction.
35

---

Published in Transactions on Machine Learning Research (02/2026)
Table 16:
One-sided and two-sided Mann-Whitncy U test results comparing Trained agents to bascline agents across varying grid sizes. For each grid configuration, we tested whether the Trwined agent achieves significantly longer episode durations than Stationary, Random, and Greedy baselines. Bonferroni correction is applied to the one-sided p-values. Significance levels: p 0.05, ** p 0.01, *** p< 0.001, n.s. not significant. "Winner" indicates the model with significantly longer duration after correction. "Mean Difl" refers to (Trained Baseline) mean.
Grid
Baseline
p (2-sid.)
(1-sid.)
Sig (2)
p (1) Corr
Sig (1)
Winner
Mean Diff
30
Stationary
1.96e-40
9.80e-41
2.94c-40
***
Trained Trained
9.12
Random
1.62e-24
8.13e-25
2.44e-24
7.56
Greedy 0.1482
0.0741
n1.8
0.2223
n.8.
33.27
40 x 40
Stationary 3.11e-29
1.55e-29
4.67e-29
***
Trained
26.81
Random Greedy
1.74e-17
8.71e-18
2.61e-17
Trained
21.66
6.17e-12
3.23e-12
888
9.70e-12
***
Greedy
61.05
50 x 50
Stationary 1.01e-16
5.06e-17 17
1.52e-16
Trained
101.40
Random Greedy
4.04c-16
2.02e-16
6.06e-16
***
Trained
114.08
2.07e-37
1.04e-37
3.11e-37
Greudy
225.88
60 60
Stationary
7.98e-16
3.99e-16
1.20e-15
Trained
142.10
Random 4.25e-24
2. 2.13e-24
88%
6.38e-24
***
Trained
176.46
Greedy 1.12e-24
5.57e-25
1.67e-24
Greedy
127.80
Table 17: Mann-Whitney U test results comparing Trained agents to Stationary, Random, and Greedy baselines under varying adherence effectiveness. Both two-sided and one-sided tests were performed. One-sided p-values were Bonferroni-corrected. "Winner" indicates the agent with significantly longer episode duration. "Mean Difl" is Trained minus Baseline. Significance: < 0.05, ** < 0.01, *** < 0.001, n.s. = not.
significant.
Adh. Eff. Baseline
p (2-sid.)
p (1-sid.)
Sig (2)
(1) Corr
Sig (1)
Winner
Mean Diff
0.1
Stationary
1.05e-35
5.25e-36
***
1.58e-35
Trained
172.85
Random
4.14e-33
2.07e-33
6.21e-33
Trained
177.92
Greedy
4.49e-34
2.24e-34
6.73e-34
888
Greedy
192.65
0.3
Stationary
1.07e-08
5.37e-09
1.61e-88
Trained
51.67
Random 2.36e-06
1.18c-06
3.55e-06
Trained
62.49
Greedy 1.22e-43
6.12e-44
1.84e-33
Greedy 249.12
0.5
Stationary
0.0180
0.0090
0.0270
Trained
25.03
Random 6.83e-06
3.42e-06
***
1.03e-05
***
Trained
51.82
Greedy 4.72e-42
2.36e-42
7.07e-42
Greedy
242.83
36

---

Published in Transactions on Machine Learning Research (02/2026)
Mean Episode Duration (steps)
Agent Type
Stationary
Random
Trained
Greedy
200
I
100
I
I
I
Adherence Eff. 0.1
Adherence Eff. 0.3
Adherence Eff. 0.5
Adherence Effectiveness (a.#)
Figure 14: Mean episode durations of Trained, Stationary, Random, and Greedy agents across varying levels of adherence effectiveness. Fach bar represents the mean of independent training runs, each evaluated over 100 episodes (300 evaluations per agent-type per sctting), with 95% bootstrapped confidence intervals (10,000 resamples). A red dashed line indicates the maximum episode length imposed by the environment. Statistical comparisons between the Trained agent and each baseline are provided in Table 17, using one-sided and two-sided Mann-Whitne U tests with Bonferroni correction.
Table 18: Mann Whitney U test results comparing Trained agents to Slationary, Random, and Greedy baselines across varying levels of the distance decay parameter. Both two-sided and one-sided tests were performed. Onc-sided p-valucs were corrected using Bonferroni correction. "Winner" denotes the agent with significantly longer episode durations after correction. "Mean Diff" reports the difference mean episode duration between Trained and baseline (positive = Trained better). Significance thresholds: p 0.05, ** p 0.01, *** p < 0.001, n.s. significant.
Dist. Decay
Baseline
p (2-sid.)
(1-sid.)
Sig (2)
p (1) Corr
Sig (1)
Winner
Mean Diff
0.15
Stationary
2.42e-26
1.21e-26
***
3.63e-26
***
Trained
8.44
Random
3.97e-18
1.99e-18
5.96e-18
Trained
17.90
Greedy
6.77e-14
3.39e-14
***
1.02e-13
***
Greedy
-96.67
0.30
Stationary
1.09e-24
5.44e-25
***
1.63e-21
***
Trained 113.78
Random
1.73e-21
6.64e-22
***
2.59e-21
***
Trained
117.54
Greedy
9.73e-35
4.86e-35
1.46e-34
*** Greedy
2213.42
0.45
Stationary
2.91e-21
1.46e-21
4.37e-21
Trained 147.33
Random Greedy 1.42e-15
2.84e-28
1.42e-28
***
4.26e-28
***
Trained
165.01
7.1le-16
2.13e-15
Greedy
-77.05
37

---

Published in Transactions on Machine Learning Research (02/2026)
Mean Episode Duration (steps)
Agent Type
300
Stationary
Random
Trained
Greedy
100
Distance Decay 0.15
Distance Decay 0.30 Distance Decay Factor (kaa)
Distance Decay 0.45
Figure 15: Performance comparison of Trained and baseline agents across varying values of the distance decay parameter, which controls how repulsion from infected individuals diminishes with spatial separation. Bars show mcan episode durations, with 95% bootstrapped confidence intervals (10,000 samples) calculated from per-seed means. The red dashed line indicates the maximum episode duration. For statistical analysis of whether Troined agents significantly outperform baselines, see Table 18.
38


Reviews:
[{"review_id": "OQaTnFGemC", "reviewer": "Reviewer_oFM4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper develops a new platform for spatial epidemic simulations. The special feature of the new platform is that it encodes an agent \nin the model, allowing its behavior to evolve over time upon interaction with the environment (while other individuals' behavior in the models is treated as static). On the platform, the agent learns the survival strategy via reinforcement learning algorithms, where different reward functions can be implemented. The paper provides numerical studies demonstrating what can be evaluated and compared \non the platform (e.g., different RL algorithms and different reward functions), drawing some preliminary conclusions based on the simulation results.", "claims_and_evidence_yn": "No", "claims_explanation": "1. Although the authors provide some discussion of the rationale behind the single-agent design, it remains unclear why this simplification would not substantially affect the overall conclusions. In particular, it is not obvious that the qualitative insights would carry over unchanged. Could the authors comment on whether---and to what extent---the framework could be extended to a setting with a finite number of interacting agents?\n\n2. The message of the paper appears to be a bit unclear to me. While the proposed ContagionRL framework is presented as a platform for evaluating and comparing different reinforcement-learning algorithms and reward functions, it is not evident how such comparisons are validated. For instance, if the single-agent formulation differs meaningfully from a setting in which each individual is modeled as an agent, the resulting conclusions could change. I wonder whether there is any notion of ground truth or external benchmark against which the proposed evaluations and comparisons are assessed.", "audience_yn": "Yes", "audience_explanation": "The idea of agent-based modeling in an epidemic simulation platform should be of interest to some audience.", "requested_changes": "Below are the questions/requested changes following up on my previous comments:\n\n1. Could the authors comment on whether---and to what extent---the framework could be extended to a setting with a finite number of interacting agents?\n\n2. Could the authors comment on whether there is any notion of ground truth or external benchmark against which the proposed evaluations and comparisons are assessed?\n\n3. I wonder if the heterogeneity of the population can be encoded into the model. For example, different characteristics of the agent could lead to other behaviors and different model outcomes.", "broader_impact_concerns": "NA.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Om3eORwjIU", "reviewer": "Reviewer_uG7X", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a framework for evaluating different individual-level behavioral strategies in epidemics. This is done in detail via a spatial agent-based model with dynamic interactions, i.e., agents move in space while potentially spreading the disease. It is assumed there is a single strategic agent who learns optimal policies for spatial movements via reinforcement learning (assuming an appropriate reward function is specified). The remaining agents move and interact in simple, natural ways. The authors then run extensive experiments in their framework, highlighting the impact of optimized decision-making, partial information and other properties.\n\nStrengths:\nThe proposed framework is fairly detailed and flexible. I could see it being useful for testing various hypotheses and running counterfactual experiments. The authors also highlight an interesting result that if strategic agents act on partial information rather than full information, they may be a higher chance of staying uninfected for longer. \n\nWeaknesses: \n- While the system described in the paper is a good building block for more complex experiments, it is in my opinion too simplistic to draw serious insights. A large part of what makes epidemics interesting and unpredictable is that agent behaviors *collectively* propel the disease forward in a population. However, the focus here is only on a single agent, whereas a multi-agent RL setup may be more realistic. \n- Given that key insights are drawn based on only one agent in a presumably large population, I wonder how necessary it is to simulate the entire large-scale system.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The paper is clearly written, and all insights are supported by appropriate data and plots.", "audience_yn": "Yes", "audience_explanation": "This paper studies an agent-based approach to a specific large-scale reinforcement learning problem. While the problem is highly specific and the framework is somewhat simplistic, it is a reasonable building block for more complex models and predictions.", "requested_changes": "List of changes required for informing my decision:\n- Please discuss how computational costs increase with the number of agents to simulate. Since we are only interested in the dynamics of one agent, I wonder if it is truly necessary to simulate the rest in detail, as this would be costly to compute. Perhaps a low-dimensional \"fluid limit\" could even be used in place of agent-based dynamics (except for the sole strategic agent). Please provide more context and information around these ideas. \n- Please discuss how one would make this model more realistic, i.e., account for more detailed physical constraints and incorporate a multi-agent learning framework.\n- Infections per timestep is a metric that is assessed in the paper. To my understanding, this captures the number of infections in the entire system per unit of time. If so, can you explain why this is used to assess the performance of a single strategic agent?\n\nMinor changes / questions to be addressed:\n- In your paper, $\\alpha$ represents the adherence to an NPI. If this always reduces the likelihood of transmission, why wouldn't a strategic agent always choose the largest value it can be? If it is always advantageous to choose the largest value, why have it as an adjustable parameter?", "broader_impact_concerns": "I have no concerns related to the broader impacts of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "dAJ5aU0juW", "reviewer": "Reviewer_cDcn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Recent works have studied epidemic simulations broadly within the context of Agent-Based Models (ABMs). These approaches apply heuristics as compared to general learning-based alternatives for learning fine-grained behavioral control. Towards this goal, the work presents ContagionRL, a Reinforcement Learning (RL) platform for studying spatial epidemic simulations through the lens of reward functions. Simulations define an agent which interacts with rule-based Gaussian humans and agents over a grid in order to survive and adhere to epidemic conditions. SIRS+D epidemiological model is employed for simulating dynamics with continuous control agents such as PPO, SAC and A2C considered in a continuous state and action space setup. At the core of the paper, the platform studies and integrates a range of reward functions (sparse to dense variants) in order to study the role of reward engineering in epidemic behavioral control. While sparse rewards such as Reduce Infection Prob are found to be effective, dense rewards such as Max Distance and the proposed Potential Field rewards provide a more informative signal. Evaluations assessing episode lengths and infections per timestep are conducted for various settings of rule-based agents, learned agents, reward functions, visibility radius, gride size and infection rate. Statistical tests demonstrate the suitability of Potential Field rewards along with the application of platform for studying agents corresponding to human movement patterns.", "claims_and_evidence_yn": "Yes", "claims_explanation": "* **Reward Function Dependency:** While the paper conducts a detailed analysis of the considered rewards, most reward functions characterize human-related pandemic behavior from a limited viewpoint. For instance, Alive and Reduce Infection Probability rewards, being sparse rewards, promote conservative and stagnant behavior. These functions simply force policies to stay away from infection. Note that from an epidemic viewpoint, this is the optimal policy. But from a social viewpoint, one would execute additional precautions or preparatory steps while moving closer to infection-prone regions. Trivially, sparse rewards (and their combinations) achieve optimal behavior by updating positions and moving away from risk-prone regions. That is, these functions emphasize more on the spatial information rather than the overall environment dynamics of the model.\n\n* **Adherence:** Following on the reward function dependency, authors note that the spatial information is the more important aspect. However, it remains unclear as to how adherence, which is a more realistic and social aspect of epidemics, ties to the simulation. Most reward functions do not take into account adherence levels and treat these as stale state and action space information. It remains unclear as to how agents exploit the adherence information (or if at all they do) if the learning signal stems from spatial components of the domain. This limits an emphasis on the more practical aspects such as social restrictions, eg - social distancing, masks, sanitation, etc. Figure 11 compares different components of the Potential Field Reward. We observe that while spatial aspects such as direction and movement are important, adherence is also important in the overall design and optimization of the reward objective. However, Potential Field reward is the only objective that exploits this information. Thus, it remains unclear as to how adherence is utilized within the platform and what role does it play during the learning process.\n\n* **Discrete RL:** Authors formulate the platform as a continuous control problem making use of common RL algorithms such as PPO, SAC and A2C. Additionally, heuristic baselines are also utilized to compare traditional rule-based approaches. It would be worth discussing the choice of the setup as a continuous control problem and how the platform could be customized or extended for formulating discrete problems. Similarly, authors could discuss algorithms (such as DQN, DDQN, Impala) which could be utilized within the epidemiological domains for fine-grained control.\n\n* **Comparisons of Partial-Observability:** _\"This counter-intuitive...observation spaces.\"_ Could the authors further explain this insight? It remains unclear as to how increased observation visibility may be attributed to state space noise or curse of dimensionality. In the case of noise, could the authors explain as to what is the importance of different components of the features present in the state representation? As in how is feature importance distributed across the state? This would help explain the role of different components of state features and how they benefit/deteriorate reward-based learning. Humans generally benefit from additional information during epidemics. Their actions become more deliberate, i.e- the policy distribution becomes sharper and peaks towards deterministic behaviors.\n\n## Minors\n\n* Page 6 last paragraph: of a each $\\rightarrow$ of each\n* Page 10 scond last paragraph: could you explain/study in detail the counter-intuitive result on limite visibility?\n* Figure 6: 100 evaluat ion $\\rightarrow$ evaluation\n* Figure 6: Trained agent significantly $\\rightarrow$ agents\n* Figure 14: Could you explain the intuition and result of the plot? It is unclear as to how episode length should vary for different adherence levels.", "audience_yn": "Yes", "audience_explanation": "The paper studies an important problem by both designing a platform for its simulation and illustrating the use of the platform using reward-based learning. From an application perspective, the work studies a practical problem receiving limited attention. From a theoretical and experimental standpoint, the work systematically dissects the contributions of the various components of the platform. I believe this is a detailed study in RL which might be of interest to both the machine learning and scientific community in general.", "requested_changes": "Below are the changes critical for acceptance-\n* **Adherence:** An analysis and discussion of adherence plays a key role in determining policy improvement towards epidemic rule compliance. Authors should study this component by taking into account 1-2 reward functions that take adherence into account. These functions can then be compared to Potential Field Reward and its various ablations. Additionally, authors should consider highlighting how agents learn policies corresponding to adherence. This can be achieved by showing that adherence, both in the state and action spaces, is adequately utilized and not simply dominated by spatial features.\n\nBelow are the changes that would strengthen the quality of work-\n* **Movement Patterns:** Similar to the workplace movement pattern, authors could highlight or discuss additional movement patterns that are commonly observed during epidemics. For instance, home to grocery or home to clinic are common 2-location patterns. The platform could also cover/discuss additional 3-location patterns and qualitatively analyze the behavior of agents in simulation.\n\n* **Discrete RL:** Customization in Rl environments has been a long-standing challenge. Adding support for discrete environments in the form of a customizable state and action space would be beneficial. Similarly, authors could highlight/discuss discrete RL algorithms that may be utilized in such cases.\n\n* **Partial Observability:** Authors should consider explaining the counter-intuitive result on limited visibility. This can be achieved by studying feature importance ofthe state representation and perhaps providing more expressivity to RL agents in order to adapt to a larger dimensional state space.", "broader_impact_concerns": "Authors should consider adding a broader impact statement discussing the significance of their work towards population-based studies and epidemic control from a public health perspective.", "additional_comments": "NA", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)