Paper: Contextualized Policy Recovery: Modeling and Interpreting Medical Decisions with Adaptive Imitation Learning (ICLR-2024)

Reviewer comment (Writing):
the authors appear to be using behavior cloning and imitation learning interchangeably.

Author response:
2. __Behavior cloning and imitation learning are used interchangeably__

Thank you for alerting us to this ambiguity. Behavior cloning is a specific implementation of imitation learning based on supervised learning that emphasizes black-box policy performance, while imitation learning methods in general focus on policy reasoning but are not tied to a specific optimization scheme. CPR spans both areas, emphasizing both policy reasoning and performance. Based on prior literature terminology [1] and CPR’s interpretable design, we call CPR a medical imitation learning method and have updated the manuscript to be more consistent.

Reviewer comment (Writing):
(minor) typo in Eq. 1 $P_w(\theta|c)$

Author response:
__Misc__
- Math typos have been fixed and clarified.

Reviewer comment (Theory):
how did you realize the stochasticity around $\theta$ suggested in Eq. 1 if it's parameterized only by an RNN?

Author response:
3. __Realization of the stochasticity around $\theta$__

Thanks for the suggestion! We have clarified the problem setup in Section 3.1, and we conducted an additional experiment to discuss the robustness of our model with stochastic $\theta$. We find that CPR’s novel context-specific policy representation is significantly more robust to noise than black-box methods, with the context-specific policy serving as both an interpretability mechanism and a regularizer (Figure 5). We believe this result will be of interest to the representation learning and interpretable ML communities.

Reviewer comment (Experiments):
why all the experimented environments have a binary action space? It'd be good to see how the interpretability and accuracy generalize to multi-class predictions with a linear model.

Author response:
4. __All the experimented environments have a binary action space.__

The multi-action setting is possible with CPR, but we design our study with an emphasis on binary actions to enable fair comparison with recent baseline methods [1, 2], which have significant reproducibility issues. 

The most notable methods for medical imitation learning either require expert knowledge for implementation on new tasks [2] or are completely inaccessible [1] with dead links to supposedly open-source code. For these reasons, we have been unable to reproduce these supposedly SOTA baselines for new medical imitation learning tasks, which might extend to a multi-action setting. Furthermore, there are major discrepancies between these works in their reported results (see the INTERPOLE results in Table 2 from both [1] and [2]), and in many cases both are outperformed by a simple logistic regression baseline, and their reported comparison against a vanilla RNN is not rigorous, since we achieve a far stronger RNN baseline with minimal hyperparameter tuning (our Table 2 RNN vs. the BC baselines in [1] and [2]). This inconsistency and lack of reproducibility is a major concern for the fast-growing medical imitation learning community. In response, we have made CPR open-source and designed it to be task-agnostic, allowing an easy and accessible baseline for future work. The lack of reproducible baselines also limited our comparison with prior works to datasets where results were published and preprocessing code was standardized. There is an immediate and pressing need to publish open-source and adaptive tools to foster growth in this community. Alongside CPR’s positive results, our negative results on the reproducibility of prior works and CPR’s open-source and task-agnostic code are a timely and important contribution to this community.

Reviewer comment (Evaluation):
as the weights of the regression model are parametrized by an RNN, there is inductive bias between consecutive timesteps. So I would envisage it's possible that e.g. a feature is found to be important at a certain step not least because it was important at the previous step, but not that it's actually still important (or its current importance being exaggerated). Would you think this should be a problem and why?

Author response:
5. __Inductive bias of RNN and LSTM__

Thank you for raising this point. Our core framework splits black-box and glass-box models into modular components to enable context-specific interpretability. As a secondary consequence, any recurrent model can be substituted for the black-box model. We explore the RNN and LSTM for proof of concept, and our results strongly support that CPR is SOTA in terms of both accuracy and interpretability with this implementation. We primarily focus on validating the utility and performance of our novel interpretability mechanism, the context-specific policy. Extensive exploration of inductive biases in RNNs and LSTMs would distract from this aim, and we leave this to future work given that our results already strongly support CPR’s SOTA performance in terms of both accuracy and interpretability.

Reviewer comment (Novelty):
While the contribution is simple and easy to understand, I believe the contribution is a bit incremental than one would like.

Author response:
We provide an objective and domain-agnostic empirical evaluation of CPR against prior works on standardized success metrics in Table 2, and these quantitative results indicate extremely strong performance with +7% and +22% AUROC on canonical tasks for medical imitation learning. We then qualify and validate CPR’s performance on these metrics by comparing CPR’s interpretable context-specific policies against medical literature on treatment policies. In this follow-up, we present a variety of data-driven evaluations that directly utilize CPR’s mechanism of interpretability (Figures 2, 3, 4) to discover nuanced treatment policies and these treatment policies that are strongly supported by medical literature. Through rigorous quantitative and qualitative evaluation, we show that the context-specific policies generated by CPR are both medically meaningful and highly predictive. We directly utilize the learned context-specific policy representations to provide nuanced data-driven interpretations of medical decisions, which are extensively validated by medical literature (Figures 2, 3, 4). CPR recovers treatment heterogeneity relating to age and gender-based risk-factors (Fig. 4) [3], the context dependence of prior treatments (Fig. 2) [4], and catastrophic outlier scenarios such as kidney failure (Fig. 3) [5]. Our interpretable framework reveals driving factors of medical decisions that were missed by prior medical imitation learning methods [1, 2] and would be impossible to discover with these methods due to sub-par model alignment, lack of context dependence, and low accuracy.

We further improve our study’s relevance to the general machine learning community with a new simulation study and empirical evaluation, detailed below and available in the re-uploaded manuscript. To summarize, our work has significance to both the representation learning, interpretable ML, and medical informatics communities, and our results strongly support that CPR is SOTA in terms of both accuracy and interpretability for medical imitation learning.

Reviewer comment (Experiments):
I think its a strong assumption to make that imitation learning is the best thing to do in MIMIC and ADNI datasets. Healthcare biases in EHR data are fairly well-reported [1], although one might hypothesize it is less so for ICU data. Without clear discussion of why this is a reasonable choice, the relevance and applicability diminishes a bit, in my view.

Author response:
Electronic record keeping is ubiquitous in healthcare, but data collection has vastly outpaced the development of methods to analyze, quantify, and provide actionable data-driven insights about clinical operations. To address this, medical imitation learning aims to quantify and explain the decisions physicians make during treatment, providing an unprecedented opportunity to identify biases in treatment, unexplained variability among patients, and improve clinical operations. For precisely the reasons the reviewer mentions (EHR and treatment biases, domain knowledge necessity) as well as heterogeneity among patients, medical applications are perhaps the most difficult yet important use case of imitation learning. Methods designed specifically for _medical_ imitation learning are needed to address the unique nuances and constraints of medical data [1, 2]. Although the reviewer rightfully points out systemic biases in EHR data collection, the MIMIC and ADNI datasets used in our evaluation have become canonical tasks for medical imitation learning methods based on their well-understood clinical workflows, known biases and nuances to treatment, and the general consensus among clinicians that real policies can be improved for these tasks. When biases in treatment are present, CPR can recover these biases. CPR effectively detects age and gender-driven biases in treatment (Figure 4) which are validated by medical literature [3, 4, 5].

Reviewer comment (Experiments):
Why use MIMIC-II when we're already in a world with updated version of MIMIC-IV with a larger number of patients.

Author response:
Misc:
- MIMIC-II was a typo in the original version, we study CPR on MIMIC-III.

Reviewer comment (Experiments):
I don't clearly understand the ADNI MRI task. What is the goal and what is the desired outcome? Can the authors add more context in whether these cognitive impairment rating decisions tend to be accurate in the data?

Author response:
- The ADNI MRI task is predicting when a physician will prescribe an MRI for diagnosing Alzheimer’s. In medical imitation learning, we do not directly model the patient outcome. Instead, we aim to understand when a physician makes a decision and why. MRIs are expensive, time-consuming, and uncomfortable, and over-prescription and slow hospital operations and place patients under unnecessary stress. Data-driven evaluation of MRI prescription policies seeks to improve hospital operations and patient experiences.

Reviewer comment (Experiments):
I would suggest adding additional empirical evaluation in more traditional imitation learning tasks to show utility without requiring significant domain knowledge to interpret all results.

Author response:
Given that both primary baselines were published recently at ICLR, and that both baselines only contain studies with real observational data and a similar simulation to our previous result, we believe that it would be inappropriate and unfair to reject this work on the basis that it does not include a more domain-agnostic evaluation. However, we do agree with the reviewer that performing objective assessment of medical imitation learning methods on a known MDP would push the field forward and set an excellent example for future work. To foster the development of interpretable imitation learning methods and promote comparisons with black-box models, we design a simulation study to compare CPR against black-box models and post-hoc interpretability methods. We evaluate CPR and black-box methods in terms of action-matching accuracy and their recovery of known policy parameters. In this simulation, we find that CPR’s SOTA performance and interpretability is driven by its explicit representation of the context-specific policy, which in certain cases improves over the performance of even black-box models, a significant result for the representation learning and interpretable ML communities. This simulation study is in the new Section 4.2 with results in Figure 5, but for the reviewer’s convenience we summarize below.

Reviewer comment (Experiments):
I would also consider improving the empirical evaluation altogether. Suppose clinicians given you some domain constraints that need to be satisfied (e.g., antibiotics dose should not exceed a certain amount, although this is not the exact decision point here). Can the authors demonstrate that using contextualized policy learning more easily satisfies such constraints (without explicitly adding them) compared to existing baselines and RNN/LSTM parametrizations of the policy class? I think this would add some interesting evaluation than is currently presented in the paper.

Author response:
- Domain constraints would be excellent in practice, but in this study we design CPR to be domain-agnostic within medical imitation learning. We show that this approach recovers nuanced and medically valid treatment regimes with minimal prior bias, allowing us to validate CPR’s interpretations against known best-practices. We believe this minimal-bias approach to data-driven policy modeling is the most adaptable and convincing framework for our study.

Reviewer comment (Experiments):
I still don't know why use MIMIC-III over MIMIC-IV.

Author response:
Misc:
- MIMIC-II was a typo in the original version, we study CPR on MIMIC-III.

Reviewer comment (Evaluation):
I also don't believe methods like LIME can actually provide real interpretability of treatment policies in a way that's actionable to clinicians.

Author response:
We aim to test CPR in known environments (simulations) to understand how its learned policy representation deviates from true policies. Toward this end, we provide an improved simulation and a comparison with post-hoc interpretability methods. CPR incorporates both a deep learning component and a statistical modeling component to introduce a novel mechanism of interpretability, the context-specific linear policy. Our approach differs substantially from prior interpretable methods by including a deep learning component. Naturally, we explore if CPR's explicit linear policy representation is key to its performance and interpretability, or if accurate and robust context-specific linear policies can also be recovered from black-box policy models using post-hoc interpretation methods. To test this, we simulate a heterogeneous, action-dependent Markov decision process (MDP) and evaluate CPR versus black-box baselines on their recovery of true simulation parameters: the true action probability and the true coefficients of a context-specific linear policy. While CPR explicitly generates these context-specific linear coefficients, black-box models implicitly model these coefficients as feature gradients (i.e. linear coefficients in a first-order Taylor expansion). Akin to popular post-hoc interpretability methods like LIME by (Ribeiro et al., 2016), we leverage the differentiability of RNNs $\Phi(x_t, h) \rightarrow a_t$ to recover the implicit context-specific linear policies $\theta$. $$\widehat{\theta} = \frac{\partial}{\partial x_t}\Phi(x_t, h)$$ We generate data with a known heterogeneous MDP governed by a true context-specific policy $$P(a_t = 1|x_t, x_{t-\tau}, a_{t-\tau}, t) = 1 / (1 + \exp(-\theta \cdot x_t + \epsilon_a))$$ $$\theta = x_{t-\tau} \cdot (2 a_{t-\tau} - 1) + \frac{t}{T} + \epsilon_{\theta}$$ where $T$ is the sequence length or maximum timestep, and $\tau$ is a time lag between the current policy and a dependence on past observations $x_{t-\tau}$ and actions $a_{t-\tau}$. We simulate $N$ total sequences, drawing observations $x_t \sim \operatorname{Unif}[-2, 2]$, decision noise $\epsilon_a \sim N(0, \sigma_{a}^2)$, and policy noise $\epsilon_{\theta} \sim N(0, \sigma_{\theta}^2)$ at each timestep. We compare the policy models learned by CPR and RNN in terms of the Pearson's correlation coefficient (PCC) between estimated and true action probabilities and context-specific policy coefficients. We hold out 15% of trajectories at random for evaluation. Results for total patients $N=200$, time-lag $\tau = 2$, decision noise $\sigma_a = 0$, and $\sigma_{\theta} = 0$, and sequence length $T = 10$ are below, presented below as mean (std).

|              | PCC θ            | PCC P(a)        |
|--------------------|------------------|-----------------|
| CPR explicit policy model | 0.830 (0.036)    | 0.848 (0.029)   |
| RNN implicit policy model | 0.203 (0.021)    | 0.706 (0.087)   |

On this known heterogeneous and action-dependent MDP, CPR's explicit policy representation not only improves its representation of MDP parameters but increases overall performance versus a black-box model with an unstructured policy representation. It is exciting to us that CPR's policy representation is not only more accurate than implicit policy representations in unstructured black-box models, but that CPR's representation choice also serves as a regularizer under the right circumstances to improve prediction accuracy as well.

Reviewer comment (Experiments):
I do believe adding more clinician annotations will make the study interesting and the empirical evaluation more credible.

Author response:
We developed CPR in collaboration with physicians to make accurate and realistic assumptions about clinical decision making. Through this collaboration, we identified the importance of context in interpreting medical observations, and the lack of existing methods that encode context-dependence in decision policies. In CPR, we address this by designing a machine learning architecture with a novel interpretable representation of medical decisions: the context-specific policy, which we rigorously evaluate using real and simulated data as well as medical domain knowledge.