Improving counterfactual reasoning with kernelised dynamic mixing models

Abstract

Simulation-based approaches to disease progression allow us to make counterfactual predictions about the effects of an untried series of treatment choices. However, building accurate simulators of disease progression is challenging, limiting the utility of these approaches for real world treatment planning. In this work, we present a novel simulation-based reinforcement learning approach that mixes between models and kernel-based approaches to make its forward predictions. On two real world tasks, managing sepsis and treating HIV, we demonstrate that our approach both learns state-of-the-art treatment policies and can make accurate forward predictions about the effects of treatments on unseen patients.

Fig 1

In our model-mixing approach (a), we create a simulator that chooses between parametric (discrete POMDP) and non-parametric (kernel) approaches for performing the forward simulation and use this simulator for planning. In contrast, earlier work (b) solved for a policy using either the POMDP or the kernel, and then chose between policies at test time. Given that both the POMDP and the kernel both have their respective weaknesses, we expect policies derived from just one to be less optimal than those derived from a model that can dynamically mix between both.

Fig 2. Illustration of dynamics for the toy example.

The optimal sequence of actions for a type A variant is to initially take action 1 or 2, followed by action 1. For type B variants, the optimal sequence of actions is to first take actions 1 or 2, followed by action 2.

Fig 3. Simulating the viral load in an HIV patient when the viral load is below detection limits (indicated by 0).

KDM can detect the occurrence of blips at 12 and 30 months, unlike a MoE. No treatment change should be administered here.

Fig 4. Comparison of predictive log-likelihood across baselines for HIV for a typical test patient.

KDM’s predictions are more accurate across the forward time steps.

Fig 5. Box plot of viral load predictions across 3000 test patients under baselines over a 30-month horizon.

KDM’s predictions are closer to the ground truth than POMDP or kernel predictions.

Fig 6. Simulating the SpO₂ of a sepsis test patient under baselines over a 20-hour horizon.

Counterfactual predictions of SpO₂ levels are more accurate using KDM than existing baselines.

Fig 7. Comparison of predictive log-likelihood across baselines for sepsis for a typical test patient.

KDM’s predictions are more accurate across the forward time steps.

Fig 8. Box plot of SpO₂ predictions across 3000 test patients under baselines over a 20-hour horizon.

KDM’s predictions are closer to the ground truth than POMDP or kernel predictions.

Fig 9

Distributions of frequencies of non-zero IS weights for (a) HIV and (b) sepsis respectively. Our treatments are fairly consistent with those in the data sets.