Interpretable Reinforcement Learning for Healthcare

Reinforcement-learning agents can learn effective policies while making it difficult to understand why they choose a particular action. That is a serious limitation in a high-stakes setting such as healthcare, where a practitioner needs to inspect and trust a model’s behavior.

For this project, we approximated a trained black-box policy with interpretable classifiers. Our main focus was the decision set: a collection of independent if-then rules that maps a state directly to an action. Unlike a decision tree, a decision set does not require someone to remember every earlier branch in a hierarchy before interpreting a prediction.

Approach

We studied two ways to collect training examples from an expert policy:

1. A supervised approach that records state-action pairs generated by the expert.
2. DAgger, an imitation-learning algorithm that repeatedly adds states visited by the learned policy and asks the expert which action it would take.

We first validated the idea in a small grid-world environment, where a decision set recovered the optimal policy. We then evaluated decision sets, decision trees, Explainable Boosting Machines, logistic regression, and Extra Trees in a simulated HIV treatment-scheduling environment.

Takeaway

DAgger substantially improved the learned policies because it addressed the distribution shift created when an imitation policy visits states that were rare in the expert’s original trajectories. The more flexible tree-based models came closest to the expert’s reward. Decision sets sacrificed some performance, but produced individually readable rules with at most two predicates—a concrete example of the tradeoff between predictive performance and human interpretability.

This was a team project with Nam Luu Nhat and Nicholas Stern.

We also released pynterp, an open-source Python package providing some of the interpretable algorithms developed for this work, including decision sets.

Read the full project report.