[Paper Review] Knowledge Infused Policy Gradients for Adaptive Pandemic Control

Kaushik Roy, Qi Zhang+2 more|arXiv (Cornell University)|Feb 11, 2021

Advanced Bandit Algorithms Research15 references4 citations

TL;DR

This paper proposes Knowledge Infused Policy Gradients (KIPG), a reinforcement learning framework that integrates structured public health knowledge—such as soft and hard constraints on lockdown policies—into adaptive pandemic control. By modeling non-homogeneous, relational features through kernelized aggregates and enforcing functional constraints, KIPG enables sample-efficient, interpretable policy learning that outperforms knowledge-agnostic baselines in balancing health and economic outcomes during simulated pandemics.

ABSTRACT

COVID-19 has impacted nations differently based on their policy implementations. The effective policy requires taking into account public information and adaptability to new knowledge. Epidemiological models built to understand COVID-19 seldom provide the policymaker with the capability for adaptive pandemic control (APC). Among the core challenges to be overcome include (a) inability to handle a high degree of non-homogeneity in different contributing features across the pandemic timeline, (b) lack of an approach that enables adaptive incorporation of public health expert knowledge, and (c) transparent models that enable understanding of the decision-making process in suggesting policy. In this work, we take the early steps to address these challenges using Knowledge Infused Policy Gradient (KIPG) methods. Prior work on knowledge infusion does not handle soft and hard imposition of varying forms of knowledge in disease information and guidelines to necessarily comply with. Furthermore, the models do not attend to non-homogeneity in feature counts, manifesting as partial observability in informing the policy. Additionally, interpretable structures are extracted post-learning instead of learning an interpretable model required for APC. To this end, we introduce a mathematical framework for KIPG methods that can (a) induce relevant feature counts over multi-relational features of the world, (b) handle latent non-homogeneous counts as hidden variables that are linear combinations of kernelized aggregates over the features, and (b) infuse knowledge as functional constraints in a principled manner. The study establishes a theory for imposing hard and soft constraints and simulates it through experiments. In comparison with knowledge-intensive baselines, we show quick sample efficient adaptation to new knowledge and interpretability in the learned policy, especially in a pandemic context.

Motivation & Objective

• Address the lack of adaptive, interpretable, and knowledge-aware reinforcement learning in pandemic control, where policies must balance public health and economic resilience.
• Overcome limitations in prior knowledge infusion methods that fail to handle non-homogeneous feature counts, partial observability, and structured expert knowledge.
• Develop a framework that enables both hard and soft constraints (e.g., closing high-risk venues or cautious lockdowns) to be embedded as functional constraints in the policy gradient process.
• Ensure policy interpretability by learning structured, linear basis functions with transparent feature weighting, avoiding post-hoc explanation methods.
• Enable rapid adaptation to new knowledge—such as emerging hotspots or updated guidelines—through principled constraint infusion without retraining from scratch.

Proposed method

• Model the world as a multi-relational graph where entities (e.g., workplaces, shops) and their interactions are represented as features with non-homogeneous counts.
• Use kernelized aggregation functions to compute hidden variables that represent latent, non-homogeneous feature counts (e.g., number of people in a hotspot), treating them as latent states.
• Infuse structured public health knowledge as functional constraints—both hard (e.g., mandatory closure of high-interaction zones) and soft (e.g., cautious lockdowns)—into the policy gradient objective.
• Formulate the policy gradient update with constraints via a Lagrangian relaxation, enabling joint optimization of the reward and constraint satisfaction.
• Enforce smoothness and interpretability in the neural network policy by leveraging linear basis functions and structured input representations, ensuring human-readable feature importance.
• Apply the framework in a simulated pandemic environment with dynamic, non-stationary conditions and partial observability of infection rates.

Experimental results

Research questions

• RQ1 Can knowledge infusion via functional constraints improve sample efficiency in learning adaptive pandemic control policies compared to standard RL?
• RQ2 How does the integration of both soft and hard constraints—representing public health guidelines—impact policy performance and convergence speed?
• RQ3 To what extent can the model handle non-homogeneous, relational features (e.g., varying numbers of people in different locations) while maintaining interpretability?
• RQ4 Can the policy be made interpretable *during* learning, rather than through post-hoc explanation, by structuring the policy as a linear combination of interpretable basis functions?
• RQ5 How does the model adapt to new knowledge (e.g., new hotspots or updated guidelines) without catastrophic forgetting or extensive retraining?

Key findings

• KIPG demonstrated significant sample efficiency in learning effective pandemic control policies, especially when new knowledge was introduced during training.
• The model achieved faster convergence and better performance than knowledge-agnostic baselines, particularly in balancing competing objectives like lives saved and economic impact.
• Interpretable policy structures were learned directly through linear basis functions, enabling human experts to assess feature importance and validate policy decisions.
• The framework successfully handled both soft constraints (e.g., cautious lockdowns) and hard constraints (e.g., mandatory closures of high-risk venues) through functional constraints in the optimization process.
• The use of kernelized aggregates for latent feature counts allowed the model to manage partial observability and non-stationarity in pandemic dynamics effectively.
• Policy weights at the input layer were found to be aligned with domain knowledge, indicating that the agent learned to prioritize high-risk locations and behaviors as expected by public health experts.