Consider the problem of an epidemic spreading in the population, without any known cure or vaccination procedure. To contain the disease and prevent it from spreading, it becomes critical to detect infected carriers and isolate them; see Fig. 1 for an illustration. As the epidemic spreads, the demand for tests outgrows their availability, and not all potential carriers can be tested. It becomes necessary to identify the most likely epidemic carriers using limited testing resources. This raises a major question: How can we rank candidates and prioritize testing to prevent the disease from spreading?
This prioritization problem is an important example of a family of problems: learning to control diffusive processes over networks through nodal interventions. Other examples include opinions spreading on social network, product adaption, viruses inflicting computer networks and cascades of failures in server farms. In all these cases, the dynamics of the system can be steered using interventions that modify the states of a (relatively) small number of nodes. For instance, infected people can be asked to self-quarantine, preventing the spread of a disease, at-risk computers can be patched by security updates, and users may be selected and be exposed to new information to influence their opinion.
The problem of controlling the dynamics of a system using localized interventions is very hard, and for several reasons. First, it requires to make decision in a continuously changing environment with complex dependencies. Second, to solve the problem one must assess the potential downstream ripple effect for any specific node that becomes infected, and balance it with the probability that the node indeed becomes infected. Finally, models must handle noise and partial observability.
We pose the problem of controlling a diffusive process on a temporally evolving graph as a sequential decision making problem in the context of a partially-observed Markov decision process. We then formulate the problem of selecting a subset of nodes for dynamical intervention as a ranking problem, and design an Figure 1 : A viral infection process on a graph and an intervention aimed to stop its spread. Here, graph nodes represent people and edges represent interactions. At t = 1 only two people are infected (red). At t = 2 several interactions resulted in new exposed people (yellow); At t = 3 the blue node was selected to be quarantined to stop the viral spread. This paper presents a framework for learning how to select which nodes should be quarantined.
actor-critic RL algorithm to solve it. We use the observed changes of nodes states and connections to construct a temporal multi-graph, which has time-stamped interactions over edges, and describe a deep architecture based on GNNs to process it. To the best of our knowledge, this work is the first work that consider Deep RL in the context of a temporally evolving graph.
The current work combines three research fields: dynamic processes on temporal graphs, deep graph learning and reinforcement learning. Combining these three into a cohesive model raises several new challenges. Most importantly, the model must learn to handle two types of dynamics: learn to infer the potential risk of not removing a node, and learn to predict the probability that a node becomes infected. As explained below, these two processes operate on different time scales. To address this issue, our architecture contains two separate GNN modules, taking as input a multi-graph over the nodes, where edges are time-stamped with the time of interactions. Also, we show below that combining RL with temporal graphs requires to stabilize how information is aggregated from neighbors when updating nodes hidden states, and control how actions are sampled during training to ensure sufficient exploration. We show empirically the beneficial effects of these components. This paper demonstrates that combining RL with GNNs provides a powerful approach for controlling spreading processes on graphs. In the context of COVID-19 spread, we show that using the RL-GNN approach increases the fraction of healthy by 25% and allows for confining the spread of an epidemic 30% more often, and 3× times more often that using non-learned approaches. This paper makes the following contributions:
(1) A new framework for controlling the dynamics of diffusive processes over graphs. Namely, learning to perform local interventions to steer the global dynamics of a graph-based dynamical system.
(2) A new architecture for this problem, and a way to train a decision-making agent using reinforcement learning to prioritize interventions on the temporal multi-graph.
(3) An observation of the interplay between the dynamics of graph states and how information flows over the graph for a decision making agent, which motivates the design of our deep network architecture. (4) A set of benchmarks and strong baselines for this problem. This includes statistics collected from real-world contact tracing data for COVID-19. Our RL approach achieves superior performance over these datasets, often significantly.
Our work is in the intersection of a few disciplines. We tackle the problem of controlling a dynamic process by considering it as ranking problem on a temporal graph. As a particular example, we address the problem of a controlling a viral epidemic spreading on a social graph.
Deep Learning on graphs. Graph neural networks (GNNs) are deep neural networks that can process graph-structured data. GNNs became very popular and were shown useful for solving a variety of tasks including social network analysis [25, 10] and molecule property prediction [14, 8] . Perhaps the most popular GNN models are Message Passing Neural Networks (MPNN) [14, 17, 51] , which operate by repeatedly updating the feature vector of each node by aggregating information from its neighbourhood. Several works combine recurrent mechanisms with GNNs to learn temporal graph data, e.g., [16, 58, 11, 55] tackled the traffic forecasting problem . [26, 22] proposed a graph-structured RNN for coarse spatial prediction of epidemic spread. Unlike this work, these works model the epidemic spread and do not try to intervene with the diffusive process. More generally, several recent studies [29, 43, 28, 39] tackle a setup in which both nodes and edges vary over time, with applications in social network analysis and other fields. Further information can be found in [23] .
Ranking on graphs. The problem of ranking on graphs is a fundamental problem in Computer Science, in which the task is to rank the nodes of a given graph according to some criteria. It has various applications such as web page ranking [38, 3] and knowledge graph search [53] .
Reinforcement learning and graphs. Recently, a surge of work combining Reinforcement Learning and graphs emerged. These works can be split into two main categories: leveraging graph structure for general RL problems (e.g., [56, 21] ), and applying reinforcement learning methods for graph problems. Our work falls into the latter. An important line of work uses Reinforcement Learning to solve NP-hard combinatorial optimization problems defined on a graph [59, 5] . Another common application is the usage of RL for path searching in a knowledge graph [53, 6] . Reinforcement learning was also shown in a few other graph problems, such as chemical reaction prediction [7] .
Dynamic processes on graphs. Modelling diffusive processes is an active research field. Key models such as SIR (Suscpetible-Infected-Removed) and SIS (Suscpetible-Infected-Suscpetible) [36] to the recent SEIR (Suscpetible-Exposed-Infectious-Removed) COVID-19 epidemic model [27, 30] have proven useful in modelling the spread of contagions. The application of these models is extended, and varies from early epidemic detection [33] , to influence maximization [24] and network security [15] . Recently, researchers have also used deep learning to forecast epidemic growth, e.g., [2, 42] .
The problem of node manipulation (e.g., vaccination) for controlling epidemic processes on graphs was intensively studied [20, 31] . This problem is often addressed in the setup of the fire-fighter problem and its extensions [13, 50, 48] . Common approaches include include developing centrality measures designed to highlight bottleneck nodes [54] , or using spectral methods for allocating resources [46, 41, 37] . Alternative line of research [34, 4] developed heuristics for the same task.
Our setup differs from these work in three important aspects. First, we do not assume a node can be vaccinated or immunized against the epidemic. In a realistic scenario, once a vaccination is discovered it will be rapidly distributed to the masses. For example, in the context of a computer malware, distributing the vaccine can be easily applied to every node by means of a software patch. In other cases, such as human epidemics, it is likely that factors such as underlying personal health condition (e.g, if a person is at increased risk) will be the dominant factor in prioritizing vaccination. Therefore, here we consider the problem at its pre-vaccination stage, where every node is susceptible to the disease.
Second, we do not assume a node can be quarantined or disconnected from the graph without justification, namely, without a positive test result. It may be tempting to isolate a high-degree node from the network, like putting a bus-driver in quarantine or disconnecting a hub computer from the network, but we view this solution as problematic. Nodes perform required network functionality. Removing a load balancer without a proper replacement will substantially impair the computer network functionality. Quarantining a bus driver due to her connectivity pattern will result in either replacing the bus driver with another driver with the same interactions pattern, or the transportation network functionality will deteriorate. In other words, a preemptive node removal would either not affect the network connectivity or impair the network functionality. Therefore, we assume that a node can be put in isolation only after it tests positive for infection. Finally, most previous works considered a setup in which a single vaccination decision is taken at t = 0 (with a possible delay execution), based only on the information available at t = 0 . In our setup, the agent needs to decide at each time step t on its testing policy, given the information available at t ≥ 0, which include the results of past actions. Furthermore, the agent is aware to the fact that it may take additional action at later times. The agent needs to balance between retrieving information (for better informed future decisions), maximizing the probability that the immediate intervention will be successful, and optimizing the long-term goal that depends the stochastic epidemic evolution, current action and future actions. The agent needs to optimize the trajectory of test allocations, where each decision is taken under different information setting. In contrast to other work, the agent need not only know where to test, but also when to test each node, in a partially observed setup with little information.
We begin with an example to illustrate the trade-offs of the problem. A natural algorithmic choice would be to act upon nodes that are most likely infected. The following example shows why this approach is suboptimal.
We form a time-varying graph from a list of interactions between nodes at various times. If u, v interacted at time t then the edge (u, v) exists at time t. Each interaction is characterized by a transmission probability p e (t). If a node was infected at time t and its neighbor was healthy, then the healthy node is infected with probability p e (t).
For the purpose of this example, assume we can test a single node at odd timesteps. If the node is identified as infected, it is sent to quarantine and cannot further interact with other nodes. Otherwise, we do not perturb the dynamics and it may interact freely with its neighbors. Our goal is to minimize the number of infected nodes.
Consider the "two stars" network in Figure 2 . The left hub (node v 1 ) has m 1 neighbors, and m 2 nodes are attached to the right hub v 2 . At t = 0, only the edge e = (v 1 , v 2 ) is present with p e (t = 0) = p. Then, for all t ≥ 1, all edges depicted in Figure 2 exist with p e (t) = 1. Assume that this information is known to the agent, and that at t = 1 it is revealed that node v 1 was infected at t = 0.
In this example, we clearly should test either v 1 or v 2 . We can compute the expected cost of each option exactly. Alternative I: Test v 2 . With probability p, v 2 becomes infected at t = 1, and we block the epidemic from spreading. However, we forfeit protecting v 1 neighbors, as all of them will be infected in the next step. With probability 1−p test is negative, and we fail to affect the dynamics. At t = 2 node v 2 will get infected and at t = 3 all of v 2 's neighbors become infected too, ending up with a total of (m 2 + 1) infections. The expected cost in choosing to test v 2 is (1 − p) · m 2 + m 1 . Alternative II: Test v 1 . We block the spread to v 1 's neighbors, but sacrifice all m 2 neighbors of v 2 with probability p. The expected cost in choosing v 2 is p · m 2 .
The decision would therefore be to test for v 2 if 2p ≥ 1 + m 1 /m 2 . This example illustrates that an optimal policy must balance two factors: the probability that the dynamics is affected -that a test action yields a "positive", and the future consequences of our action -the strategic importance of selecting v 1 vs. v 2 , expressed by the ratio m 1 /m 2 . A policy targeting likely-infected nodes will always pick node v 1 , but since it only focuses on the first term and ignores the second term, it is clearly suboptimal. An immediate question arise: How can we develop methods that address both terms? It is difficult Figure 3 : Schematic of our approach. The Ranking module receives as input a sequence of graphs and outputs scores over nodes. Scores are then used to sample actions, selecting nodes for intervention. Here, the person circled in blue is selected for quarantine and its connections are cancelled (dashed blue lines). The downstream effect on epidemic progression is then fed as a loss to the ranking module.
to measure the strategic importance of a node in a large dynamical process with constant external interventions directly. Instead, one may use simulations and try to learn from the collected experience and the interplay between dynamics and actions a method that finds the optimal policy end-to-end and internally weighs the two terms optimally. This calls for a reinforcement learning framework.
We start our discussion with a concrete example focusing on the spread of epidemics, and then present the general formalism. Section 7 provides additional use-cases and applications of the general framework.
We consider the recent COVID-19 outbreak that spreads through social contacts. The temporal graph G is defined over a group of nodes (people) V, and its edges E(t) are determined by their daily social interactions. An edge (u, v) between two nodes exists at time t iff the two nodes interacted at time t. Each of these interactions is characterized by various features e uv (t), including its duration, distancing and environment (indoors or outdoors). Additionally, each node v is characterized with features ζ v (t) (e.g, age, sex etc.).
The SEIR model. We follow the widely-used SEIR model [30] . Every node (person) can be in one of the following states, Y = S, L, I, R, namely: susceptible -a healthy, yet uninfected person (S state), exposed/latent -infected but cannot infect others (L state), infectious -may infect other nodes (I state), or removed -selfquarantined and isolated from the graph (R state). The state of a node v ∈ V at time t is a random variable denoted by ST v (t)
Node state dynamics. A healthy node can become infected with a probability that depends on its interactions with its neighbors. Once infected, transitioning from Exposed/Latent to Infected is defined by a probabilistic process. A node becomes emoved if it is selected for self-quarantine (tests positive); see Fig. 1 for an illustration.
Formally, let I(t) ⊂ V be the set of infectious nodes at time t, and similarly L(t), R(t) and S(t) be the sets of latent(exposed), removed and susceptible (healthy) nodes. Each active edge at time t, e ∈ E(t), carries a transmission probability p e (t). Denote the set of impinging edges on node v with an infectious counterpart at time t by E v (t) = {e ∈ E(t)|e = (v, u), SV u (t − 1) = I} . The probability of a healthy node to remain healthy at time t is 1 − e∈Ev(t) (1 − p e (t)), otherwise it becomes infected, but still in a latent state. We denote the time of infection of node v as T v . A node in a latent state will stay in this state at time t if t < T v + D v , where D v is a random variable representing the latency period length, otherwise its state changes to infectious.
The testing intervention h(·) changes the state of a node. If infected or exposed, its state is set to R, otherwise it remains as it is.
Optimization goal and action space. The objective is to minimize the spread of the epidemic, namely, minimize the number of infected people over time, I STv(t)∈{L,D} .
Assuming that testing is limited to a fixed capacity of k tests per day, the optimization goal becomes min t,v γ t I STv(t)∈{L,D} , where γ ∈ (0, 1] is a discount factor representing the relative importance of the future compared to the present. We used γ = 0.99 throughout the paper. Of course, this is but one particular optimization goal and other goals that weigh different states differently, or add additional aspects are possible.
The action space consists of all possible selections of a subset a(t) of k nodes a(t) ⊂ V . Even for moderate graph, with ∼ 100 − 1000 and small k the action space |V| k is huge. Observation space. At each time t, the agent is exposed to all past interactions between network nodes {E(t )|t < t}. In addition, we are given partial information on the nodes state. The agent is provided with information on a subset of the infectious nodes at t = 0. At every t > 0, the agent observes all past test results. Formally, for every v ∈ a(t) we observe if the node was infected, ST v (t) ∈ I(t) ∪ L(t), or not.
Consider a diffusive process on a temporal graph G(t) = (V, E(t)) whose structure changes in time. V is the set of nodes and E(t) = {e uv (t)} is the set of edges at time step t. Each edge e uv (t) is associated with features φ uv (t) which may vary in time, and each node v is characterized with features ζ v (t).
The state of a node v ∈ V is a random variable ST v (t) which can have values in Y = {y 1 , y 2 , ..}. The node's state ST v (t) depends on the interactions between v and its neighbors at time t and on the state of those neighbors, We aggregate all neighborhood information to a random variable
Additionally, ST v (t) depends on the previous state ST v (t − 1) and on node features. In other words,
At each step, we may select a subset A(t) of k nodes, and change the state of any selected node v ∈ A(t). Selecting nodes and setting their states defines the space of actions for the agent, and plays the role of a knob for controlling the global dynamics of the process over the graph. Formally we apply a transformation h(·) on
The optimization objective should be invariant to permuting (relabeling) of the nodes. We assume it depends only on the total number of nodes in state i,
where future evaluation are weighted by a discount factor γ ≤ 1. Additionally, the agent may be subject to constraints written in a similar manner i f i (c 1 (t), c 2 (t), ..) ≥ y i (t). 
Our approach is depicted in Figure 3 . The RL agent receives as input the node and edge features of the temporal graph, and processes them using its ranking module. A probability distribution over graph nodes is generated from the ranking module's scores, and the agent samples a subset of k nodes for testing. Namely, the scores encode the agent's policy. We use Proximal Policy Optimization algorithm (PPO, [49] ) as to optimize our agent. We sequentially apply the suggested action, log the (state, action) tuple in an experience replay buffer, and train our model based on the PPO loss term. Next, we describe the ranking module and sampling procedure.
Overview Our GNN-RNN module serves to update the internal representation of a node v, denoted h v (t), and its score s v (t) (Figure 4 ). This score is later used for selecting nodes to be acted on. The ranking module is applied in a recurrent fashion at each time step. The ranking module contains two GNNs: (1) E, which updates the epidemic state, and (2) I, which updates the information state. It also contains two other networks, G and F , which update node representations and node scores by using the epidemic state and information state as well as the previous node representations.
Input. The input to the ranking module consists of three feature types (See Figure 4 ): (1) Static node features ζ s v (t): topological graph centralities (betweeness, closeness, eigenvector and degree centralities) and random node features. (2) Dynamic node features ζ d v (t) : All test results that were performed up the current timestamp (including positive and negative test results). We denote all nodes features as a concatenation
(3) Edge features and the structure of the temporal graph E(t): All previous interactions up to the current step, including the transmission probability for each interaction. Figure 4 illustrates the basic data flow in the ranking module.
Epidemic GNN. The spread of epidemic through point contact is modeled by a GNN E. As the epidemic cannot spread by more than one hop per step, it is sufficient to model the spread with a single GNN layer.
Formally, denote by p vu the probability of transmission during the interaction between (v, u) at time t. For each v, the output of E(·) is a feature vector denoted by e v (t):
where M is multilayer perceptron (MLP). Rather than considering the probability as an edge feature, this module mimics the epidemic transition rule to accelerate learning.
Information GNN. The score of a node is affected both by the propagation dynamics, and by the information available to the agent. One may hope that since the former has a known timescale (days), on a short time scale (single day) the score of node would only be affected by its neighboring nodes. This, however, is not the true because information can propagate long distance in the graph almost instantaneously. As a simple example, consider nodes in a connected chain of (untested) nodes and note that they are statistically dependent. As a result, revealing the state of one node immediately affects the distribution over all nodes in the chain.
With this consideration in mind, we designed an information GNN, I, which represents the information state of each node. As discussed above, updated information on a node u a few hops away from node v may abruptly change our beliefs on the state of v. Furthermore, this change may occur even if v and u did not interact in the last time step but rather a while ago. To update the information state, we construct a cumulative multi-graph G where the set of edges between nodes v and u at time t are all the interactions that occurred during the last τ steps,
, are the interaction delay t − t and the transmission probability p v,v (t ). The information features are the output of k-layer GNN; the l th layer is:
are the final node features. The value of τ , the information window size, was 7 in all our experiments.
Score and hidden state update. For every node we hold a hidden state h v (t), which is updated following
After updating the new node hidden state, we use them to calculate the node score using a neural network F ,
Here, F is an MLP, while G can be either an MLP or recurrent module such as GRU.
Sampling. Given the score per node s v (t), we sample without replacement k nodes iteratively. We (1) map the score of n nodes to a probability distribution using a score-to-probability distribution function, (2) sample a node, and (3) adjust the distribution by removing its weight. We repeat this process k iterations.
RL framework. The action space of choosing a subset of k nodes out of n nodes is large even for small n and k. Using an action-value approach like Q-learning implies that an approximate value is assigned to every possible action, but the action space size is prohibitively too large for action-value methods. Instead, we use a policy-gradient algorithm and model the problem as a ranking problem. The algorithm learns to rank nodes using a parameterized model, and then uses a sampling procedure to choose a subset of k nodes.
Many on-policy gradient algorithms use entropy to define a trust region. Computing the entropy requires summing k |V| terms at each step, and it is computationally expensive. A more scalable solution is the unbiased entropy estimator of [57] , but the variance of that estimator is high. As an alternative, PPO trust region is not based on an explicit evaluation of the entropy, and performed better in our experiments. We also evaluated A2C, which did not perform as well as PPO in our experiments.
Critic module. PPO, as an actor-critic algorithm, requires a critic module to estimate the value function in a given state. We construct the actor using an architecture that is similar to the ranking module, but apply to element-wise max operation on the rows (representing the nodes) of the input to the score module F ( Figure  4 ). This reduces F 's input to a single row of features, and the output is then a scalar rather than a vector. Importantly, the critic is parametrized by a different set of weights than the ranking module (actor).
Score-to-probability distribution. Usually, node scores are converted to a distribution over actions using a softmax. This approach is problematic for our case because node probabilities decay exponentially with their scores, leading to two major drawbacks. It discourages exploration of low-score nodes, and limits sensitivity to the top of the distribution, instead of at the top-k selected. Instead, we define the probability to sample an action a i to be
where {x i } is the set of scores and a constant. By not using an exponential as in softmax, the probability differences between low scoring nodes and high scoring nodes become less extreme. Furthermore, the parameter controls the initial exploration ratio. In standard DNN initialization schemes (e.g., XAVIER), the initial value of x i is expected to be in [-1,1]. If 1 than the dominant term in Eq. (3) is . This promotes exploration initially, as all actions are likely to be sampled in the early training stages.
Normalization in scale-free networks. RNN are well-known to suffer from the problem of exploding or vanishing gradients. This problem is exacerbated in a RNN-GNN framework used for RL algorithms, because they may be applied for arbitrary long episodes, causing internal state to grow unbounded. This problem is particularly severe if the underlying graph contains hubs (highly connected nodes). One approach to alleviate this problem, is by including an RNN like a GRU module, where the hidden state values pass through a sigmoid layer. As the magnitude of the input grows, gradient become smaller and training slows down.
Scale-free networks contain with high probability "hub" nodes that have high-degree, namely O(n) neighbors. The presence of these hubs further aggravates this problem. As a simple case, consider a star graph with a large number of nodes. In a GNN framework, it receives updates from a large number of neighbors and its internal state increases in magnitude. The next time that the GNN module is applied (e.g., at the next RL step), the growing internal state increases the magnitude of the internal state of its neighbors. This leads to a positivefeedback loop that causes the internal state representation to diverge. Since RL algorithms may be applied for arbitrary long periods, the internal state may grow unbounded unless corrected.
This problem can be solved by directly normalizing each node hidden state. We have experimented with various normalization methods, and found that L 2 normalization worked best, as shown in the next section.
Transition probabilities. In the case of the COVID-19 pandemic, the transition probabilities can be estimated using the interaction properties, such as duration and inter-personal distance, using known epidemiological models. This was done by the government agency which provided our contact tracing network (see below). Alternatively, one can learn the transmission probability as a regression problem from known interactions (e.g, using data from post-infection questioning). Finally, if this information is not accessible, it is possible to omit the epidemic model E from the proposed framework and use only the feature vector created by the information module I.
We compare methods from three categorizes: A) programmed; B) supervised learning (SL) C) Reinforcement learning (RL). Each experiment was performed with at least three random seeds. There exist additional approaches designed for the vaccination problem, but, as discussed in Section 2, they are not appropriate for the setup studied here. More details about experiments and implementation, including network architecture, are given in Appendix A.
A. Programmed baselines. Most countries currently prioritize their testing based on fixed procedures determined in advance and not learned form data. We compare with two such methods to rank nodes. (1) Infected neighbors. Rank nodes based on the number of known infected nodes in their 2-hop neighborhood (neighbors and their neighbors). Each node v is assigned a tuple (I (1) v , I (2) v ), and tuples are sorted in a decreasing lexicographical order. A similar algorithm was used in [33, 32] to detect infected nodes in a noisy environment.
(2) Probabilistic risk. Each node keeps an estimate of the probability it is infected at time t − 1. To estimate infection probability at time t, beliefs are propagated from neighbors, and dynamic programming is used to analytically solve the probability update. See Appendix B for details. (3) Degree centrality. In this baselines high degree nodes are prioritized. This is an intuitive heuristic and it is used frequently [47] . It was found empirically to provide good results [48] . (4) Eigenvalue centrality. Another common approach it to select nodes using spectral graph properties, such as the eigenvalue centrality (e.g., [41, 54] ).
The main drawback of these programmed algorithms is that they do not exploit all available information about dynamics. Specifically, they do not use negative test results, which contain information about the true distribution of the epidemic over network nodes.
B. Supervised learning. Algorithms that learn the risk per node using features of the temporal graph, its connectivity and infection state. Then, k nodes with the highest risk are selected. (5) Supervised (vanilla). We treat each time step t and each node v i as a sample, and train a 3-layer deep network using a cross entropy loss against the ground truth state of that node at time t. The input of the DNN has two components: A static component described in Section 5.1, and a dynamic part that contains the number of infected neighbors and their neighbors (like #1 above). Note that the static features include the, amongst other features, the degree and eigenvector centralities. Therefore, if learning is successful, this baseline may derive an improved use of centralities based on local epidemic information. (6) Supervised (+GNN). Like #5, but the input to the model is the set of all historic interactions of v i 's and its d-order neighbours and their time stamps as an edge feature. The architecture is a GNN that operates on node and edge features. We used the same ranking module as our GNN framework, but the output probability is regarded as the probability that a node is infected. (6) Supervised (+weighted degree). Same as #6, but the loss is modified and nodes are weighted by their degree. Indeed, we wish to favour models that are more accurate on high-degree nodes, because they may infect a greater number of nodes. (8) Supervised (+weighted degree +GNN). Like #6 above, using degree-weighted loss like #7.
C. RL algorithms: RLGN is our algorithm described in Section 5. The input to (9) RL-vanilla is the same as in (#1) and (#6) above. Correspondingly, the GNN module of described in Section 5 is replaced by a DNN similar to (#6), while the rest of the RL algorithm remains intact.
Training. We train the RL and SL by generating random networks, and selecting for each instance a random subset of m 0 initially infected nodes. We propagate the epidemic until it spans at least k 0 infected nodes (for at least t 0 steps), and randomly detect a subset of the infected nodes of size < k 0 . The simulation then follows the agent-network dynamics described in Section 4.
Evaluation Metrics. The end goal of quarantining and epidemiological testing is to minimize the spread of the epidemic. As it is unreasonable to eradicate the epidemic using social distancing alone, the hope is to "flatten the curve", namely, to slow down the epidemic progress. We use two success metrics: (1) % healthy: The percent of nodes kept healthy throughout the simulation. (2) %contained: The probability of containing the epidemic. This was computed as the fraction of simulations having cumulative infected nodes smaller than a fraction α. We focus on this metric because it captures the important notion of the capacity of a health system. In the 2-community setup, where each community has half of the nodes, a natural choice of α is slightly greater than 0.5, capturing those cases where the algorithm contains the epidemic within the infected community. In all the experiments we set α = 0.6. The only exception is the three-communities experiments, in which we set the bar slightly higher than 1/3, and fixed α = 0.4.
We study three types of networks which differ by their connectivity patterns.
(1) Community-based networks have nodes clustered into densely-connected communities, with sparse connections across communities. We use the Stochastic Block Model (SBM, [1] ), for 2 and 3 communities. (2) a matrix P of size m × m, which represents the edge probabilities between nodes in different communities, namely, the matrix entry P i,j determines the probability of an edge (v, v ) between v ∈ C i and v ∈ C j . The diagonal elements in P are often much larger than the off-diagonal elements, representing the dense connectivity in a community, compared to the intra-connectivity between communities.
(2) Preferential attachment (PA) networks exhibit a node-degree distribution that follows a power-law (scale-free), like those found in many real-world networks. We use the dual Barbarsi-Albert model [35] , which allows for continuously varying the mean node degree.
Generating temporal graphs. Static networks generated using PA or SBM are converted to a temporal graph by first selecting a random subset of edges E(t) at each time step t, and then assigning to each edge a transmission probability q e (t) sampled uniformly U [a, b].
(3) Contact-tracing networks. We received anonymized high-level statistical information about real contact tracing networks that included the distribution of node degree, transmission probability and mean number of interactions per day, collected during April 2020. Fig. 6 (a) presents the degree distribution in this data, and the transmission probability is presented in Fig.  6 (b). The latter was derived based on the contact properties, such as the length and the proximity of the interaction. On average, 1.635 ± 0.211 interactions with a significant transmission probability were recorded per-person per-day. We generated random networks based on these distributions using a configuration model framework [36] . The fitted model for the degree distribution is a mixture of a Gaussian and a power-law distribution P (degree = x) = 0.47 · N (0.41, 0.036) + 0.53 · Beta(5.05, 20.02).
The fitted model for the transmission probability is a mixture of a Gaussian and a Beta distribution
(4) large-scale datasets. We consider graphs of different sizes and sources: (A) A research collaboration network (CA-GrQc), representing social interactions between researchers. In addition to social interaction, information also diffuses on this network [44] . (B) The Montreal dataset (Montreal), based on WiFi hotspot tracing, collected in Montreal, Canada over three years, representing a proximity network [19] . (C) a compartmentbased synthetic network, derived from first principles, using a population census in Portland, Oregon (Portland) [52, 9] . (D) An email network (Enron), representing a computer network on which malwares spread, (E) An online social network. We used the recent GEMSEC-RO [45] , presenting friendship relations in the Deezer music service in Romania. This accommodates for a graph on which preferences and product adaption (songs) propagate. All these networks have been extensively used in previous work, in particular epidemiological studies, as key networks models [48, 54, 19, 18, 52, 9] . Table 1 summarizes the various datasets, spanning graphs with 5K nodes to graphs with more than 100K nodes.
In a first set of experiments, we compared 7 approaches for learning and inference on 3 network types. The results reported in Table 2 and 3 show that RLGN outperforms all baselines in all network types. A video highlighting the results can be found online 1 .
To gain insight into this result, we first look more deeply into the case of 3-community networks. Figure  5 (a) traces the fraction of contained epidemics and Figure 5 (b) the fraction of infected nodes during training. The supervised learning algorithm detects substantially more infected nodes (right panel) than RLGN, but these tend to have lower future impact on the spread, and it fails to contain the epidemic (left). A closer look shows that RL, but not SL, successfully learns to identify and neutralize the critical nodes that connect communities and prevent the disease from spreading to another community.
To further understand the solutions learned for PA networks, consider the following two extremes. First, when a network is very sparsely connected, it would be easy to cut long infection chains, and both approaches are expected to be successful. At the other extreme, for densely connected networks, there are no critical nodes, because there are many paths between any two nodes. To study this effect we generated networks with the preferential-attachment mechanism, while varying the mean node degree. This degree can also be viewed in terms of the R 0 coefficient, the mean number of nodes infected by a single diseased node. The greater R 0 , the more difficult it is to contain the epidemic. Fig. 7 (a) quantifies this effect, showing the percent of contained epidemics as a function of the R 0 coefficient. RL has a significant advantage over supervised+GNN for a range of R 0 values between 2.0 and 2.9. Finally, Figure 7 (b) depicts a robustness analysis of RLGN for variations in the epidemiological model. One of the most difficult quantities to assess is the probability for infection per social interaction. Figure 7(b) shows that the trained model can sustain up to ∼ 40% deviation at test time in this key parameter.
We have tested the robustness of our results to the underlying graph size. Specifically, we compare the two best algorithms RLGN (#8) and SL+GNN (#4), using graphs with various sizes, from 300 nodes to 1000 nodes. Table 4 compares RLGN with the SL+GNN algorithm on preferential attachment (PA) networks (mean degree = 2.8). We provide results for various sizes of initial infection i 0 and number of available tests k at each step. The experiments show that there is a considerable gap between the performance of the RL and the second-best baseline. Furthermore, RLGN achieves better performance than the SL+GNN algorithm with 40%-100% more tests. Namely, it increases the effective number of tests by a factor of ×1.4 − ×2.
We also tested the sensitivity of the results to the relative size of the initial infection. Table 4 shows results when 4% of the the network was initially infected, as well as for 7.5% and 10%. The results show that RLGN outperforms the baselines in this wide range of infection sizes.
We further evaluated how models trained on medium-size graphs generalize to perform inference on much larger graphs. Specifically, we trained RLGN and SL+GNN (three model initialization for each) on a preferential attachment network with 1000 nodes and evaluated its performance on various larger graphs.
We investigated the progression of the epidemic under either RLGN or supervised+GNN algorithms. Figure  8 shows that the epidemic spread speed is slower under the RLGN policy in all graph. In general, there are two extreme configuration regimes. First, the "too-hard" case, when the number of tests is insufficient to block the Table 4 : A comparison between RLGN and SL+GNN (the best baseline). RLGN performance is highlighted. The number of additional resources needed to surpass the RLGN performance in a given metric is also highlighted. In many cases, even using SL+GNN with twice as many resources than RLGN performs worse than RLGN. The evaluation was performed on a preferential attachment network with mean degree 2.8. The number of nodes is indicated at the top of each 0.91 ± 0.03 810 L 2 norm. 0.93 ± 0.02 500 Table 6 : Training time and fraction of contained epidemic for three normalization schemes. The L 2 normalization scheme is fastest and achieves the best performance.
epidemic, and second, the "too-easy" case when there is a surplus of tests such that every reasonable algorithm can contain it. The more interesting case is the intermediate regime, where some algorithms succeed to delay the epidemic progression, or block it completely, better than other algorithms. Fig. 8(a) illustrates the case where the number of tests is insufficient for containing the epidemic, for all algorithms we tested. In Fig. 8(b) , the number of tests is insufficient for SL to block the epidemic. However, with same number of tests, RLGN algorithm successfully contains the epidemic. Fig. 8(c) presents an interesting case where RLGN slows down the epidemic progression and reduces the number of total number of infected node, compared with SL, but RL does not contain it completely. Table 5 compares the performance of the RLGN and the best baseline (SL) on the large-scale datasets. We included the centralities baselines (#3,#4) in the comparison as they are heavily used in epidemiological studies. Table 5 shows that the RLGN consistently performs better than the baselines, and the gap is clearly statistically significant. Finally, we evaluated the performance of RLGN on a Preferential Attachment network with 50, 000 nodes (mean degree = 2.8), as this random network model is considered a reasonable approximation for many other real world networks. We found that RLGN successfully contained the epidemic in all 150 evaluation episodes, while the SL+GNN was unable to block the epidemic even once. The mean percentile of healthy nodes at the end of the episode was 51 ± 1 for RLGN, while for the SL+GNN it was only 21 ± 2, a difference of more than 15 STDs. 6.7 Ablation Studies 6.7.1 Mapping scores to action distribution.
We compare the performance of our score-to-probability function (calibrated-scores) to the popular softmax (Boltzmann) distribution. In practice, in most instances, we were unable to train a model using the softmax distribution as the neural network weights diverge. Fig. 9 presents the training curve in one of the few instances that did converge. It is clear that the model was not able to learn a useful policy while using the calibrated-scores probability function resulted in a corresponding value of more than 0.75.
We compared the suggested normalization to a number of other alternative normalization methods. (1) Applying a sigmoid layer after the hidden state update module G. (2) Replace the hidden state update module with a GRU layer. (3) Apply L 2 normalization to each feature vector h v (t) (similarly to [17] ) (4) Normalize the feature vector matrix by its L 2 norm. These four normalization schemes span three different types of normalization: single-feature normalization (1+2), vector normalization (3), and matrix normalization (4). Figure 9 : The fraction of contained epidemics during training on a preferential attachment model with 200 nodes and a mean degree 2.8. For non-normalized mapping, only one of the three seeds in the softmax distribution simulation completed training due to numerical instability. No stability issues were observed when using the calibrated scores normalization scheme described by Eq. Eq. (3). Table 6 presents the score after training and the number of training steps required to complete training. Method (4) was unstable and training did not converge, therefore it was omitted from the table. The main reason for the training time difference is that without normalization, the DNN weights' magnitude increases. In a GRU module, or with a direct application of a sigmoid layer, the features pass through a sigmoid activation function. When the magnitude of the input to this layer is large, the gradient is very small due to the sigmoid plateau. This substantially slows down the learning process.
Our experiments showed the information module has a critical role in improving the performance of the RL-GNN framework. We performed an ablation study by removing it completely from our DNN module, keeping only the epidemic module. The full DNN module achieved a contained epidemic score of 0.77 ± 0.06, while the ablated DNN module corresponding score was 0.62 ± 0.10, a degradation of more than 20%.
The approach and model discussed in this paper can be applied to important problems other than epidemic control.
Influence maximization (IM). Unlike epidemic control, in IM, the decision maker the objective is to maximize the spread of some opinion on a social network graph. They act to select nodes to be influenced by presenting information, actively contacting them, or sending coupons etc. Given a budget, the IM agent can only contact a limited number of nodes over a period of time. Influence spreads over the social networks similar to the model described above. The overall cost has two terms: cost of spending budget on influencing individual nodes, and benefit from influencing a give portion of the network. Local effects such as high degree nodes are important to IM.
Fake news detection and confinement. Consider a social network where fake news can be maliciously distributed, and spread over the network. A decision maker can verify the authenticity of items, but only verify a limited number of items per a time period. Once a fake item is detected, it is erased from the network. The objective is to minimize the total number of nodes that observe fake items. The main difference from the epidemic problem is that once a fake item is discovered, it can be erased from the entire network. The trade-off is that the decision maker does not want to waste inspection resources on items that are likely to die out. The fake new control and confinement problem can be therefore naturally cast within the framework outlined here.
Epidemic Control: Beyond Node Selection. While this paper addressed the important problem of deciding which nodes to examine given a certain umber of tests, we can also ask what if not only the subset of nodes to be tested, but also the number of tests. to be performed can be controlled. Increasing the number of tests will help us contain the epidemic, however, it may be expensive to increase the number of tests and economy of scale may not work here. Another form of control is imposing a partial quarantine that reduces the number of interactions and resulting infections. Such a quarantine would come with a cost, of course. Both variations, and others, can be easily modeled within the framework presented in this paper. The crucial aspects of the model is having a objective that is cumulative reward structure, an action space that depends only on the information states, and a relatively short horizon for the relevant information to propagate so that the GNN approach works. Other important variations to the test kit allocation problems such partial quarantines, or using less accurate but more abundant tests can be easily modeled within the framework presented in this paper. The crucial aspects of the model is having an objective that has a cumulative reward structure, an action space that depends only on the information states, and a relatively short horizon for the relevant information to propagate so that the GNN approach works.
This paper shows that combining RL with GNNs provides a powerful approach for controlling spreading processes on graphs. In the context of COVID-19 spread, we demonstrated that using an RL+GNN approach allows us to confine the spread of an epidemic that is approximately 30% more contagious (i.e., R 0 that is 30% higher) with the same resources as a standard supervised learning-based approach. In addition, our results indicate that prioritizing tests using RL on temporal graphs can increase the number of healthy people by 25% and contain the epidemic 30% more often than supervised approaches and 2.5× more often than non-learned baselines using the same resources.
The architecture of the ranking module is shared by algorithms #4, #6 and #8 with slight variations indicated below.
Input. We encode the dynamic node features ζ d v (t) as a one hot vector of dimension 4. Each of the first three elements corresponds to one of the three mutually exclusive options, which depends on the action and node state in the previous step: untested, tested positive, tested negative. The last entry indicates whether a node was found positive in the past, namely, if it is quarantined and disconnected from the graph. The static node features, ζ s v (t), are as described in the main paper, topological graph centralities (betweenness, closeness, eigenvector, and degree centralities) and random node features. The graph centralities were calculated using NetworKit. The total number of node features is 9.
Epidemic GNN. This module M e is composed of a single graph convolutional layer. The input features are the last time step node features. The number of output features is 64.
Information GNN. Each message passing module M l contains one hidden layer, where the number of hidden features is 64. After both the hidden and last layer we apply a leaky ReLu layer with leakage constant 0.01. After aggregating the result using the addition aggregation function, we apply an additional MLP with one layer (linear+ReLu) on the resulting feature vector. The number of output features is 64. We experimented with the numbers of stacked modules l (layers). We found that l = 3 performed slightly better than l = 2 but training was considerably slower because the batch size had to be reduced. We therefore used l = 2 in all experiments reported.
Hidden state update. The hidden state MLP G is composed of a single linear layer follows by a ReLu activation layer. To keep the resulting hidden feature vector (of dimension 64) norm under check, an additional normalization scheme is then applied. This module was replaced by a GRU layer in the ablation studies.
Output layer. The last module is a single linear layer, with an output dimension as the number of the nodes in the graph.
Learning framework . We used Pytorch [40] and Pytorch Geometric [12] to construct the ranking module. We used ADAM with default parameters as our optimizer.
We train the RL and SL by generating random networks and initializing each network by selecting for each instance a random subset of m 0 infected nodes. We propagate the epidemic until it spans at least i 0 infected nodes (for at least t 0 steps), and randomly detect a subset of the infected nodes of size k 0 < i 0 . At each step, in all algorithms but RL, we pick the top k rated nodes. Each of these nodes is tested, and if detected is positive it is effectively removed from the graph. Otherwise, it is not modified. In RL, we perform the same procedure during the evaluation phase, while during training we sample k nodes using the score-to-probability distribution.
Each model was training for at most 1500 episodes, but usually, training was completed after 1000 episodes. Each episode contained 1024 steps, collected by 4 different workers. As our network contains a recurrent module, we propagate each sample in the buffer for three steps, in a similar fashion to R2D2.
For each setup we described, at least three models were trained using different seeds, and the results are the average over the performance of all models. The errors are the standard deviation of the mean. over at least 100 evaluation episodes for each model. Each episode lasted for 25 steps, each corresponds conceptually to a day. The transition time from the latent to the infectious state was normally distributed with a mean of two steps and a standard deviation of 1 step, corresponding to real-world values. The advantage was calculated using the Generalized Advantage framework with parameters γ = 0.99, λ = 0.97. Table 7 presents the simulation parameters used in the main paper. We shall make the repository and code available online.
In this appendix We describe our tree model baseline (algorithm #1). Consider an epidemic propagating on a tree, and assume there is a single initially infected node ("patient zero"). In this case, there is a single path from the infection source to every node in the graph and we can we can analytically solve for the probability a node is infected, given that the root of the tree was infected at time t 0 . This model is useful when the underlying network is locally a tree, i.e, that for every new infected node v there is w.h.p just one node which may have infected it.
We start with a simple case.
Let us first consider a simple model in which the epidemic spreads on a tree like structure with a single epidemic source, a.k.a. patient-zero, as the root. For now, let us assume there is no latent state. Our goal is to calculate the probability that a node n will be infected at time T For every node j there is a single path from the node to the root, denoted by r. Let us assume the path is {y 0 = r, y 1 , y 2 , ..y n−1 , y n = j}. Assume that in [0, T ] a sequence of interactions between node y n and y n−1 occurred at discrete times (t 1 , t 2 , ...t m ), and that each interaction is characterized by an infection probability (p 1 , p 2 , ...p m ). We evaluate F n (T ) by induction. For abbreviation, we write ST y i (t) = Y i (t) and denote the event ST r (0) = I as A.
Our key result is that The state of node n at the time of interaction m is a function of its state at penultimate interaction time F n (t m−1 ), the interaction transmission probability p m , and the predecessor node n − 1 state at time m, F n (t m−1 ).
The first term is the probability to get infected at the m interaction, and the second term is the probability to get infected before hand. We shall now prove this result.
Proof. We can write the conditional probability using a graphical model decomposition and obtain
since if the ancestor node is not in an infectious state, the decedent can not be infected. Denote the indicator that interaction l was able to transmit the epidemic as I l . We have, As, for an infection event to take place at it must be that node y n−1 was infected at t l , node y n was healthy beforehand, and that the interaction resulted in an infection. We can now write this as p l Pr (Y n (t l−1 ) = H|Y n−1 (t l ) = I, A) F n−1 (t l ) F n−1 (t m ) = p l (1 − Pr (Y n (t l−1 ) = I|Y n−1 (t l ) = I, A)) F n−1 (t l ) F n−1 (t m )
The transition from the first line to the second is due to the independence of the interaction infection probability with the history of the participating parties. The third line is Bayes' theorem. If a node is infected at time t l , it will be infected later on at t m , as expressed in line 4. The last line is the complete probability formula.
We rewrite Pr (Y n (t l−1 ) = I|Y n−1 (t l ) = I, A) as The transition from the first line to the second line is a complete probability transition. The third line is due to the fact that if y n−1 was not infected at time t l , clearly y n could not be infected before t l . We have F n (t m ) = Pr (Y n−1 (t m ) = I|A) = m l=1 p l 1 − F n (t l−1 ) F n−1 (t l ) F n−1 (t l ) F n−1 (t m ) F n−1 (t m ) = m l=1 p l (F n−1 (t l ) − F n (t l−1 )) Therefore, given F n−1 (t l ) for all l ∈ {1..n − 1} and F n (t l ) for all l ∈ {1..n}, we can directly calculate the infection probabilities, given the initial condition: F i (0) = δ i,0 .
We can write the partial density function of F i (t l ) as f i (t l ) = F i (t l ) − F i (t l−1 ), and obtain: f n (t m ) = p m (F n−1 (t m ) − F n (t m−1 )). This allows us to write this with an intuitive formulation F n (t m ) = F n (t m−1 ) + p m (F n−1 (t m ) − F n (t m−1 )) = p m F n−1 (t m ) + F n (t m−1 ) (1 − p m )
The first term is the probability to get infected at the m interaction, and the second term is the probability to get infected before hand.
We now discuss the case where a node can be in a latent state. The main difference is that the complement of the infectious state is composed of two states, healthy H, and latent L. We shall denote all the non-infecting states as H + = {H, L} and all the infected states as I + = {I, L}, and sometime abuse the notation by writing S i (t) = H + . We denote the transmission delay from the latent to infectious state as L(τ ).
As before, we are interested in the probability that Pr Y n (T ) = I + |S r (0) = I
The derivation below shows that, similar to the previous case, we can solve for this probability using dynamic programming. The end result is that Therefore, as before, given F n−1 (t m ) and q n (t i ) for all i < m, we can propagate and calculate q n (t m ) and F n (t m ). where we kept the definition of F j (t). Therefore, almost identically, 
We have, q n (t m ) = p m (F n−1 (t m ) − F n (t m−1 ) − Pr (Y n (t l−1 ) = L|A)) .
The transition from the latent state to the infected state follows:
F n (t l ) = Pr (Y n (t l ) = I|A) = t i <t l L(t i − t l )q n (t i )
while Pr (Y n (t l ) = L|A) = t i <t l (1 − L(t i − t l )) q n (t i ).
Therefore, given F n−1 (t m ) and q n (t i ) for all i < m, we can propagate and calculate q n (t m ) and F n (t m ).