Imitation learning for combinatorial optimization and contact tracing
Keywords: combinatorial optimisation
deep learning
deep learning in public health
graph neural networks (GNNs)
COVID-19 digital contact tracing
mixed-integer linear programming (MILP)
contact tracing
imitation learning
artificial intelligence
proactive contact tracing (PCT)
linear programming
agent-based modeling (ABM)

The field of Imitation Learning (IL) has seen significant progress in recent years as researchers have applied this machine learning technique to various domains such as robotics, self-driving cars, healthcare, and game playing. Each domain has contributed to the advancement of the field by developing and applying new methods to solve the unique problems specific to their domain. In this thesis, we focus on IL in two domains that pose their own unique challenges.
The first application involves learning to imitate a highly accurate heuristic for mixed-integer linear programming (MILP) solvers, which although precise, is not practical due to its computational inefficiency. The second application involves the development of an IL framework to accurately predict the infectiousness of individuals through a smartphone application utilizing the newly developed Proactive Contact Tracing (PCT) framework, which overcomes the limitations of conventional contact tracing methods.
We design our IL frameworks based on the dynamics of a manageable environment (e.g., simulator), with the goal of transferring the learned models to larger, unseen environments. The development of these frameworks requires the consideration and resolution of several challenges. These challenges include incorporating domain-specific inductive biases, ensuring the robustness of models against distribution shifts, and designing models that are lightweight and suitable for deployment. By addressing these challenges, we hope to contribute not only to the advancement of IL, but also to the domains in which it is applied, bringing new and improved solutions to these respective fields.
Specifically, to imitate the expert heuristic of MILP solvers, we identified and addressed two key shortcomings of the existing IL framework. First, the proposed Graph Neural Networks (GNNs) are computationally expensive but highly accurate and their runtime performance degrades in the absence of GPUs. This setting may arise since MILP solvers are CPU--only. To address this, we proposed novel architectures that trade-off the expressivity of GNNs with inexpensive computations of multi-linear perceptrons, along with training protocols that make the models robust to distribution shifts. The models trained using these techniques resulted in up to 26% improvement in runtime. The second issue is the inability to capture the dependence between observations to train GNNs. Our research revealed a
lookback'' phenomenon that occurs frequently in the expert heuristic, where the best decision at the child node is often the second-best at the parent node. To incorporate this phenomenon in the loss function, we proposed a new loss function that imitates this heuristic more accurately, resulting in models with up to 15% improvement in running time.
Finally, during the COVID-19 pandemic, nations around the world faced a dilemma of whether to open up the economy or prioritize saving lives. In response, digital contact tracing applications emerged. However, to avoid violating user privacy, most apps relied on a quarantine-or-not interface with limited intelligence on the level of risk of the notification recipient. This approach led to alert fatigue, making users less likely to follow recommendations. To overcome these issues while maintaining user privacy and sophisticated risk estimation models, we proposed the Proactive Contact Tracing (PCT) framework. Our framework repurposes user communication to carry information about estimated risk in "risk messages". These messages, along with personal information (e.g., medical history or symptoms), are used in a risk estimation model to output risk messages sent to other users. Depending on estimated risk, graded notifications (e.g., exercise caution or avoid unnecessary behavior) are shown to the users. Using an agent-based model (ABM) and a simple interpretable rule-based model, we demonstrated that the rule-based PCT has a better economic-public health trade-off than the existing apps.
In follow-up work, we turned to deep learning to design a risk estimation model. While reinforcement learning would have been ideal, the computationally expensive ABM precludes its use. Instead, we employed an imitation learning framework to train deep learning models, specifically, we proposed several variants of set transformers. We also used domain randomization, collecting observations using several random instantiations of the ABM, to ensure that models were robust to assumptions baked into the ABM. Furthermore, we used iterative training to ensure the models remained robust to auto-induced distribution shifts. Overall, we showed that a deep learning-based PCT outperforms rule-based PCT. To finalize our proposal, we suggest an iterative procedure for app deployment and ABM calibration to bridge the gap from the ABM to real-world deployment.