Flattening the Curve Through Reinforcement Learning Driven Test and Trace Policies

Abstract

An effective way of limiting the diffusion of viruses when vaccines are unavailable or insufficiently potent to eradicate them is through running widespread “test and trace” programmes. Although these have been instrumental during the COVID-19 pandemic, they also lead to significant increases in public spending and societal disruptions caused by the numerous isolation requirements. What is more, after the health measures were relaxed across the world, these programmes were unable to prevent substantial upsurges in infections. Here we propose an alternative approach to conducting pathogen testing and contact tracing that is adaptable to the budgeting requirements and risk tolerances of regional policy makers, while still breaking the high risk transmission chains. To that end, we propose several agents that rank individuals based on the role they possess in their interaction network and the epidemic state over which this diffuses, showing that testing or isolating just the top ranked can achieve adequate levels of containment without incurring the costs associated with standard strategies. Additionally, we extensively compare all the policies we derive, and show that a reinforcement learning actor based on graph neural networks outcompetes the more competitive heuristics by up to 15% in the containment rate, while far surpassing the standard random samplers by margins of 50% or more. Finally, we clearly demonstrate the versatility of the learned policies by appraising the decisions taken by the deep learning agent in different contexts using a diverse set of prediction explanation and state visualization techniques.

A Appendix

1.1 A.1 Performance Analysis

We compare the mean total elapsed time for running epidemics using each of our testing agents in Table A1. These results corresponds to the wall clock time recorded on an average Windows machine equipped with an Intel i7-7700 CPU, an NVIDIA RTX 3060 GPU and 32 GB of random access memory.

Table A1. Average wall clock time per epidemic during evaluation. Configuration: Barabási-Albert networks of 2000 nodes, an average degree of approximately 3, and a daily testing budget of $k=2$.

1.2 A.2 Epidemic Modelling

All our epidemic models rely on the SEIR compartmental formulation, but the diffusion process remains bound by the interaction network configuration.

The multi-site mean-field models considered in this work rely on exponential waiting times sampled via Gillespie’s algorithm to obtain subsequent events, with the state transition probabilities defined as follows:

$$\begin{aligned} \begin{gathered} p(S \rightarrow E) = b \, w_j \mathop {}\!\mathcal {4}t \\ p(E \rightarrow I) = e^{-1} \mathop {}\!\mathcal {4}t \\ p(I \rightarrow R) = \rho \mathop {}\!\mathcal {4}t, \end{gathered} \end{aligned}$$

(2)

where b is the base transmission rate, $w_j$ are time-dependent edge weights, e is the exposed state duration, $\rho$ is the recovery rate, while $\mathop {}\!\mathcal {4}t$ is a time interval.

In contrast, the agent-based model loops through every node i at every time step, executing the appropriate transition events when one of the normally-distributed samples ($d_i$ or $r_i$) decreases to 0. Concurrently, every edge j is visited to check whether an infection event occurs over that connection, according to the transmission probability defined in Eq 2.

1.3 A.3 Algorithmic Details for the Proximal Policy Optimization

We start by reminding the reader about some general reinforcement learning quantities and relations:

$$\begin{aligned} \begin{gathered} \hat{A}_t^{(\gamma ,0)}(a_t;\theta ) = \delta _t^\gamma (\theta ) = R_t + \gamma V(s_{t+1};\theta _T) - V(s_t;\theta ) \\ \hat{A}_t^{(\gamma ,1)}(a_t;\theta ) = G_t^\gamma - V(s_t;\theta ) \\ \hat{A}_t^{(\gamma ,\lambda )}(a_t;\theta ) = \sum _{l=0}^{T} (\gamma \lambda )^l\delta _{t+l}^\gamma (\theta ) \\ r_t^{ORIG}(\theta ) = \frac{\pi (a_t|s_t;\theta )}{\pi (a_t|s_t;\theta _k)} \quad r_t^{SARSA}(\theta ) = \frac{\pi (a_t|s_{t+1};\theta )}{\pi (a_t|s_t;\theta _k)} \end{gathered} \end{aligned}$$

(3)

In Eq 3, $R_t$ is the reward obtained by the agent after taking action $a_t \sim \pi (a|s_t;\theta _k)$ and transitioning from state $s_t$ to $s_{t+1}$. The value of a given state s is approximated using a neural network $V(s,\theta )$, which, together with $R_t$ and the discount factor $\gamma$, determines the TD-error $\delta ^\gamma$; $\theta _k$ parameterizes the acting policy, $\theta _T$ is a delayed state of $\theta _k$ that parameterizes the regression target in online learning [69], $r_t^{ORIG}(\theta )$ denotes the ratio between a policy parameterized by a given $\theta$ and the acting policy, while $r_t^{SARSA}(\theta )$ represents an alternative formulation for the latter that replaces the numerator with the policy of $\theta$ evaluated at the next state $s_{t+1}$. Finally, the $\hat{A}_t^{(\gamma ,\lambda )}$ terms represent different forms of the advantage function, as given by [87], with the special cases $\lambda =0$, when the advantage is equal to the TD-error, and $\lambda =1$, when the minuend of the RHS equation is the discounted return of the episode, $G_t$.

We rewrite the Proximal Policy Optimization (PPO) equations in terms of the quantities above in Eq 4, where $\mathcal {E}$, $c_1$ and $c_2$ are hyperparameters, clip(.) is a function that clips its argument to the specified range, transform(.) is a function that modifies the gradient descent update according to a specific optimizer (e.g. Adam [47]), while $\mathcal {H}_t(\theta )$ is an entropy regularizer [31]. In contrast to the original formulation, Eq 5 describes our proposed modification of PPO to allow for optimizing the objective in a memory-efficient online manner. In particular, we rewrite the loss terms using the one-step advantage function $\hat{A}_t^{(\gamma ,0)}(a_t;\theta )$, and introduce an intermediary operation that accumulates the gradients of our modified loss using a unified eligibility trace [100], in a similar fashion to the methodology employed by [50], obtaining a backward-view approximation of the generalized advantage estimate $\hat{A}_t^{(\gamma ,\lambda )}$ in the process [87]. We note that, by setting $r_t=r_t^{SARSA}$, we can eliminate the requirement of storing $s_t$ in memory for the subsequent timestamp, while retaining the benefits of ratio clipping. This works well empirically since major shifts between $s_{t+1}$ and $s_t$ are not common in our environment. Based on previous work and our own assessment, we set $\gamma =0.99$, $\lambda =0.97$, $\mathcal {E}=0.2$, $c_1=0.5$, $c_2=0.01$, and update the target value network every 5 episodes across all our experiments.

$$\begin{aligned} \begin{gathered} \mathcal {L}^{CLIP}_t(\theta ) = \min [r_t(\theta ) \hat{A}_t^{(\gamma ,\lambda )}(a_t;\theta ), \text {clip}(r_t(\theta ), 1 - \mathcal {E}, 1 + \mathcal {E}) \hat{A}^{(\gamma ,\lambda )}_t(a_t;\theta )] \\ \mathcal {L}_t^{VF}(\theta ) = [\hat{A}_t^{(\gamma ,1)}(a_t; \theta )]^2 \quad \mathcal {H}_t(\theta ) = - \sum _{a \in A} \pi (a|s_t;\theta ) \log \pi (a|s_t;\theta )\\ \mathcal {L}^{PPO}_t(\theta ) = \mathbb {E}_{t}[-\mathcal {L}^{CLIP}_t(\theta ) + c_1 \mathcal {L}^{VF}_t(\theta ) - c_2 \mathcal {H}_t(\theta ))] \\ \theta _{k+1} = \mathop {\mathrm {arg\,min}}\limits _\theta {\mathcal {L}^{PPO}_t(\theta )} \end{gathered} \end{aligned}$$

(4)

$$\begin{aligned} \begin{gathered} \mathcal {L}^{OCLIP}_t(\theta ) = \min [r_t(\theta ) \hat{A}^{(\gamma ,0)}_t(a_t;\theta ), \text {clip}(r_t(\theta ), 1 - \mathcal {E}, 1 + \mathcal {E}) \hat{A}_t^{(\gamma ,0)}(a_t;\theta )] \\ \mathcal {L}_t^{OVF}(\theta ) = [\hat{A}_t^{(\gamma ,0)}(a_t; \theta )]^2 \quad \mathcal {H}_t(\theta ) = - \sum _{a \in A} \pi (a|s_t;\theta ) \log \pi (a|s_t;\theta )\\ \mathcal {L}^{OPPO}_t(\theta ) = -\mathcal {L}^{OCLIP}_t(\theta ) + c_1 \mathcal {L}^{OVF}_t(\theta ) - c_2 \mathcal {H}_t(\theta ) \\ E_t = \gamma \lambda E_{t-1} + \frac{\nabla _{\theta _k}\mathcal {L}^{OPPO}_t(\theta _k)}{s} \; \text {, with} \; s=\delta _t^\gamma (\theta _k) \; \text {or} \; s=1 \\ \Delta \theta _k = \text {transform}(\delta _t^\gamma (\theta _k)E_t) \\ \theta _{k+1} = \theta _k - \Delta \theta _k \end{gathered} \end{aligned}$$

(5)

1.4 A.4 Supporting Figures

Fig. A1.

Block diagram of our control framework. The Agent is passed as a parameter to the Simulator, and every time the latter samples enough events for the conditions to be met, a call to the control(.) method of the first is performed. The aforementioned function performs some preprocessing steps, and then calls control_test(.) and control_trace(.), which are responsible for the actual node ranking and are specific to each type of agent. Combinations of agents can be selected with the MixAgent.

Fig. A2.

Infection control performance on different network architectures of 1000 nodes and a daily testing budget of $k=2$. Uncertainties shown as boxplots.

Fig. A3.

Epidemic curves for different network and epidemic seeds. These correspond to multiple 5000 nodes Barabási-Albert networks featuring a mean degree of 3, with a testing budget of $k=1\%$. Here, two versions of the RL agent are displayed: one trained for 50, and one trained for 200 episodes. The y-axis limit is set to 3200 to facilitate the comparisons, yet the random agents perform poorer than this level.

Fig. A4.

Infection control performance on different static network architectures with varying budgets. The uncertainties are shown as boxplots.

Fig. A5.

Infection control performance on different static network architectures and sizes, with a budget of $k=2$. Uncertainties are shown as boxplots.

Fig. A6.

t-SNE plots of the node hidden states and dendrogram corresponding to their hierarchical clustering into 10 groups. As can be observed, the agent mostly groups detected (blue) nodes in a region of the space, while the new undetected infections (red) are predicted to appear within the risk regions on the right. Recent negative results are plotted as dark green. The dendrogram on the right displays the cardinality and the infection probability associated with each cluster.

1.5 A.5 Control Framework

The logic behind our epidemic control framework in the continuous-time simulation scenario is outlined in Algorithm 1. The class hierarchy of the agents, together with their logic, can be consulted in Algorithm 2. Refer to Table A2 for details about the variables involved in these.

Table A2. Legend for the control framework pseudocode.