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: This paper proposes a new method to learn (stationary) interpretable policies using soft decision trees in partially observed settings. The soft decision tree structure is extended to allow for recursion over time, and account for policy decisions based on history of collected data. An algorithm is presented to optimize the parameters of the soft decision tree as well as the structure/topology of the tree. The algorithm mainly proceeds by splitting nodes and locally optimizing the parameters of the the associated probability representation of the soft node, and recursively split (if local optimization does not improve validation performance) and fixed as leaf otherwise. A global update step is then used after topology is fixed followed by pruning low probability paths in the trees. Experimental validation on surveys with clinicians demonstrate reasonable interpretability and improved prediction performance on imitating clinician policy.