POETREE: Interpretable Policy Learning with Adaptive Decision Trees

1 INTRODUCTION . Different approaches to integrating , analysing and acting upon clinical information leads to wide , unwarranted variation in medical practice across regions and institutions ( McKinlay et al. , 2007 ; Westert et al. , 2018 ) . Algorithms designed to support the clinical decision-making process can help overcome this issue , with successful examples in oncology prognosis ( Beck et al. , 2011 ; Esteva et al. , 2017 ) , retinopathy detection ( Gulshan et al. , 2016 ) and radiotherapy planning ( Valdes et al. , 2017 ) . Still , these support systems focus on effectively replacing physicians with autonomous agents , trained to optimise patient outcomes or diagnosis accuracy . With limited design concern for end-users , these algorithms lack interpretability , leading to mistrust from the medical community ( Laı̈ et al. , 2020 ; The Royal Society , 2017 ) . Instead , our work aims to better describe and understand the decision-making process by observing physician behaviour , in a form of epistemology . These transparent models of behaviour can clearly synthesise domain expertise from current practice and could thus form the building blocks of future decision support systems ( Li et al. , 2015 ) , designed in partnership with clinicians for more reliable validation , better trust , and widespread acceptance . In addition , such policy models will enable quantitative comparisons of different strategies of care , and may in turn lead to novel guidelines and educational materials to combat variability of practice . The associated machine learning challenge , therefore , is to explain how physicians choose clinical actions , given new patient observations and their prior medical history – and thus describe an interpretable decision-making policy . Modelling clinical diagnostic and treatment policies is particularly challenging , with observational data involving confounding factors and unknown evolution dynamics , which we wish to then translate into a transparent representation of decision-making . State-ofthe-art solutions for uncovering demonstrated policies , such as inverse reinforcement learning ( Ng & Russell , 2000 ) or imitation learning ( Bain & Sammut , 1996 ; Piot et al. , 2014 ; Ho & Ermon , 2016 ) , fall short of interpretability – mainly proposing black-box models as behaviour representations . In contrast , transparent policy parametrisations based on visual decision boundaries ( Hüyük et al. , 2021 ) , high-level programming syntax ( Verma et al. , 2018 ) or outcome preferences ( Yau et al. , 2020 ) can explain learned behaviour , at the cost of restrictive modelling constraints and obscure dynamics . Among possible interpretable representations , decision trees offer the flexibility to capture highdimensional environments and , as we demonstrate , can model time-series through recurrence . In addition , decision trees are highly familiar to the medical community , with many clinical guidelines having this form ( Chou et al. , 2007 ; McCreery & Truelove , 1991 ; Burch et al. , 2012 ) . Cognitive research suggests human decision-making does follow a hierarchical , branching process ( Zylberberg et al. , 2017 ) , making these models straightforward to understand . As a result , our algorithm Policy Extraction through decision Trees ( POETREE ) represents clinical policies as decision trees , evolving over time with patient information – as illustrated in Figure 1 . Contributions The main contributions of this work are as follows : ( i ) A novel framework for policy learning in the form of incrementally grown , probabilistic decision trees , which adapt their complexity to the task at hand . ( ii ) An interpretable model for representation learning over timeseries called recurrent decision trees . ( iii ) Illustrative analyses of our algorithm ’ s expressive power in disambiguating policies otherwise unidentifiable with related benchmarks ; integration of domain knowledge through inductive bias ; and formalisation and quantification of abstract behavioural concepts ( e.g . uncertain , anomalous or low-value actions ) – thus demonstrating its effectiveness in both understanding and replicating behaviour . 2 PROBLEM FORMALISM . Our goal is to learn an interpretable representation of a demonstrated decision-making process to understand agent behaviour . The clinical context requires offline learning using observational data only , as experimenting with policies would be both unethical and impractical . Finally , as previous observations and actions affect treatment choice , we are concerned with a partially-observable decision-making environment . We assume access to a dataset of m demonstrated patient trajectories in discrete time D = { ( zi1 , ai1 , . . . , ziτi , a i τi ) } m i=1 , where τi is the length of trajectory i . We drop index i to denote a generic draw from the population . At timestep t , the physician-agent observes patient features , denoted by random variable Zt ∈ Z , and chooses treatment or diagnostic action At ∈ A , where A is a finite set of actions . Let zt and at denote realisations of these random variables . In contrast to traditional Markovian environments , the agent ’ s reward function R , state space S and transition dynamics are all inaccessible . Let f denote a representation function , mapping the history of patient observations and actions to a representation space H , and ht = f ( z1 , a1 , ... zt−1 , at−1 ) = f ( z1 : t−1 , a1 : t−1 ) ∈ H represent the patient history prior to observing a new set of features zt . Following the partially-observable Markov Decision Process ( POMDP ) formalism , information in { ht , zt } can be combined to form a belief over the inaccessible patient state st . In contrast to recent policy learning work on POMDPs , which uncover mappings from belief to action space ( Sun et al. , 2017 ; Choi & Kim , 2011 ; Makino & Takeuchi , 2012 ; Hüyük et al. , 2021 ) , however , our aim is to highlight how newly-acquired observations zt conditions action at , to ensure decisions are based on interpretable , meaningful variables . As a result , agent behaviour must be represented as a mapping from observation to action space . An adaptive decision-making policy is a stationary mapping π : Z × H × A → [ 0 , 1 ] , where π ( at|zt , ht ) encodes the probability of choosing action at given patient history ht and latest observation zt – thus ∀t , ∑ at∈A π ( at|zt , ht ) = 1 . We assume trajectories in D are generated by an agent following a policy πE , such that at ∼ πE ( ·|zt , ht ) . The goal of our work is to find an interpretable policy representation π which matches and explains the behaviour of πE . 3 INTERPRETABLE POLICY LEARNING WITH DECISION TREES . 3.1 SOFT DECISION TREES . Motivated by the interpretability and pervasiveness of trees in the medical literature ( Chou et al. , 2007 ) , Policy Extraction through decision Trees ( POETREE ) represents observed decision-making policies as trees . Our model architecture is inspired by soft decision trees – structures with binary nodes , probabilistic decision boundaries and fixed leaf outputs ( Irsoy et al. , 2012 ) – for their transparency and modelling performance ( Frosst & Hinton , 2017 ) . This supervised learning framework can be adapted for policy learning by setting input variables as x = zt or as { ht , zt } ( for fully- or partially-observable environments ) , and targets as y = at ( Bain & Sammut , 1996 ) . Inner nodes Illustrated in Figure 2a , the path followed by normalised input x ∈ RD is determined by gating functions encoding the probability of taking the rightmost branch at each non-leaf node n : pngate ( x ) = σ ( xTwn + bn ) , where { wn ∈ RD ; bn ∈ R } n∈Inner are trainable parameters and σ is the sigmoid function . This design choice is motivated by its demonstrated flexibility and expressivity ( Silva et al. , 2020 ) , and can be made interpretable by retaining the largest component of wn to form an unidimensional axisaligned threshold at each node , as illustrated in Figure 1 and detailed in Appendix C. Model simplification can be made more robust through L1 regularisation , encouraging sparsity in { wn } n∈Inner . The path probability of each node , Pn ( x ) , is therefore given by the product of gating functions leading to it from the tree root . Leaf nodes In our tree architecture , leaves define a fixed probability distribution over K categorical output classes . Leaf l encodes relative activation values âlt = softmax ( θ l a ) , where θla ∈ RK is a learnable parameter . For ease of com- putation and interpretation , the output of the maximum-probability leaf âlmaxt , where lmax = arg maxl P l ( x ) , is taken as overall classification decision , following Frosst & Hinton ( 2017 ) . 3.2 RECURRENT DECISION TREES . To account for previously acquired information , our algorithm must jointly extract and depend on a representation of patient history ht : Table 1 contrasts mappings learned by our recurrent decision trees ( RDTs ) and RNNs for different time-series tasks . As in Figure 2b , RDTs take history ht and new observation zt as inputs , and output both predicted action at and subsequent history embeddings ht+1 from decision leaves . These are computed as hlt+1 = tanh ( θ l h ) , where θ l h ∈ RM is an additional leaf parameter – alternative leaf models are discussed in Appendix B . Finally , as third leaf output with parameters θlz ∈ RD , our model also predicts patient evolution , or observations at the next timestep z̃t+1 , effectively capturing the expected effects of a treatment on the patient trajectory as a policy explanation ( Yau et al. , 2020 ) . A single tree structure is used at all timesteps : unfolded over time , our recurrent tree can be viewed as a sequence of cascaded trees ( Ding et al. , 2021 ) with shared parameters Θ = { { wn , bn } n∈Inner ; { θla , θlhθlz } l∈Leaf } and topology T . 3.3 TREE GROWTH AND OPTIMISATION . Optimisation objective The policy learning objective is to recover the tree topology T and parameters Θ which best describe the observation-to-action mapping demonstrated by physician agents . Our loss function combines the cross-entropy between one-hot encoded target at ∈ { 0 , 1 } K and each leaf ’ s action output , weighted by its respective path probability under input { ht , zt } : L ( { ht , zt , at } ; T , Θ ) = − ∑ l∈Leaf P l ( ht , zt ) ∑ k at , k log [ âlt , k ] ( 1 ) where âlt , k is the output probability for action class k in leaf l. Additional objectives are necessary for the multiple outputs to match intended targets in Table 1 : regularisation term Lz̃ is designed to learn patient evolution z̃t+1 as expected by the acting physician . It minimises prediction error on zt+1 and ensures the policy is consistent between timesteps by constraining predicted observations to lead to similar action choices as true ones , under the new history ht+1 : Lz̃ ( { ht , zt , zt+1 } ; T , Θ ) = δ1 ‖zt+1 − z̃t+1‖2 + δ2 DKL ( π ( ·|ht+1 , zt+1‖π ( ·|ht+1 , z̃t+1 ) ( 2 ) where { δ1 , δ2 } are tunable hyperparameters , weighting fidelity to true evolution and to demonstrated behaviour . For a set of training trajectories D , and an overall loss L = L+ Lz̃ , our policy learning optimisation objective becomes : arg min T , Θ L ( D ) ( 3 ) The probabilistic nature of our model choice makes it fully-differentiable , allowing optimisation of parameters Θ through stochastic gradient descent and backpropagation through the structure for a fixed topology T. Optimising with respect to T , however , requires a more involved algorithm . Algorithm 1 : Tree growth optimisation . 3 . Global optimisation and path pruning inner node suboptimal leaf 1 . Initialisation ✗✗ ✔ ✔✔ ✔ ✔ Local optimisation 2 . For each suboptimal leaf : a . Split node b. Label as optimalor ✗✗ ✗ optimal leaf ✗✔ Keep best-performing in validation 1 . Initialise T to inner node with two suboptimal leaves ; Optimise Θ via gradient descent of loss L ; while ∃ suboptimal leaves in T do Split suboptimal leaf into inner node with two leaves ; 2a . Locally optimise Θ via gradient descent of L ; if validation performance improves then retain split ; else 2b . Retain leaf as optimal ; 3 . Globally optimise Θ via gradient descent of L ; Prune branches in T with low path probability P l ( Dval ) . Tree growth Following the work of Irsoy et al . ( 2012 ) and Tanno et al . ( 2019 ) , our decision tree architecture is incrementally grown and optimised during the training process , allowing the model to independently adapt its complexity to the task at hand . Algorithm 1 summarises the growth and optimisation procedure . Starting from an optimised shallow structure , tree growth is carried out by sequentially splitting each leaf node and locally optimising new parameters – the split is accepted if validation performance is improved . After a final global optimisation , branches with a path probability below a certain threshold on the validation set are pruned . Further details on the tree growth procedure are given in Appendix B . Complexity Policy training can be computationally complex due to the multiple local tree optimisations involved and significant number of parameters in comparison to other interpretable models ( Hüyük et al. , 2021 ) . Still , at test time , computing tree path probabilities for a given input is inexpensive , which is advantageous to allow fast decision support . A detailed comparison of runtimes involved with our method and related work is proposed in Appendix D .

Summary:   The authors argued that many methods failed the merits of interpretability in some important areas, e.g. clinical decision-making. Thus, this paper proposed a (soft) tree-based method for synthetic clinical datasets in the matter of interpretability. The authors model the clinical decision process as a partially observable Markov Decision Process (POMDP), which naturally fits the assumption of medical diagnosis.