Optimal Decision Tree Policies for Markov Decision Processes

1 Introduction

Advances in reinforcement learning using function approximation have allowed us to train powerful agents for complex problems such as the games of Go and Atari [Schrittwieser et al., 2020]. Policies learned using function approximation often use neural networks, making them impossible for humans to understand. Therefore reinforcement learning is severely limited for applications with high-stakes decisions where the user has to trust the learned policy.

Recent work has focused on explaining opaque models such as neural networks by attributing prediction importance to the input features [Ribeiro et al., 2016; Lundberg and Lee, 2017]. However, these explanation methods cannot capture the full complexity of their models, which can mislead users when attempting to understand the predictions [Rudin, 2019]. Concurrently, there has been much work on interpretable machine learning in which the model learned is limited in complexity to the extent that humans can understand the complete model. Particularly decision trees have received much attention as they are simple models that are capable of modeling non-linear behavior [Lipton, 2018].

Decision trees are difficult to optimize as they are nondifferentiable and discontinuous. Previous works have used different strategies to overcome the hardness of optimizing trees: using assumptions or relaxations to make the trees differentiable [Gupta et al., 2015; Silva et al., 2020; Likmeta et al., 2020], reformulating the MDP into a metaMDP that exclusively models decision tree policies [Topin et al., 2021] or extracting trees from a complex teacher [Bastani et al., 2018]. While these methods are increasingly successful in training performant trees they do not offer guarantees on this performance.

Our work takes a first step at bridging the gap between the fields of optimal decision trees and reinforcement learning. Existing formulations for optimal decision trees assume a fixed training set with independent samples. This cannot be used in a dynamic setting where actions taken in one state influence the best actions in others. Instead, we formulate the problem of solving a Markov Decision Process (MDP) using a policy represented by a size-limited decision tree (see Figure 1) in a single MILP. We link the predictions of the decision tree policy to the state-action frequencies in the dual linear program for solving MDPs. The dual allows us to reason over policies explicitly, which results in a more efficient formulation. Our formulation for Optimal MDP Decision Trees, OMDTs, optimizes a decision tree policy for a given MDP and a tree size limit. OMDT produces increasingly performant policies as runtime progresses and eventually proves the optimality of its policy under the size constraint.

Existing methods for training size-limited decision trees in reinforcement learning such as VIPER [Bastani et al., 2018] make use of imitation learning, where a student tries to learn from a powerful teacher policy. We compare the performance of OMDT and VIPER on a variety of MDPs. Interestingly, we show that, when training interpretable size-limited trees, imitation learning performs significantly worse as capacity of the learned decision tree is wasted on parts of the state space that are never reached by the policy. Moreover, VIPER cannot prove optimality even if it identifies the optimal solution. Regarding the performance-interpretability trade-off, we show that decision trees of 7 decision nodes are enough to perform close to unrestricted optimal policies in 8 out of 13 environments. Such trees are orders of magnitude smaller than size-unrestricted trees created by methods that replicate the unrestricted policy such as dtcontrol [Ashok et al., 2020].

2 Background

2.1 Decision Trees

Decision trees [Breiman et al., 1984; Quinlan, 1986] are simple models that execute a series of comparisons between a feature value and a threshold in the nodes to arrive at a leaf node that contains the prediction value. Due to their simple descriptions, size-limited decision trees are easy to understand for humans [Molnar, 2020]. Particularly, size-limited decision trees admit simulatability [Lipton, 2018]: humans can use the model to make predictions by hand in reasonable time and decomposability: humans can understand each individual aspect of the model. The method proposed in this paper, OMDT, also offers algorithmic transparency: we can trust the learning algorithm to produce models that fulfill certain qualities such as global optimality. Therefore decision trees are an attractive model class when interpretable policies are required.

2.2 Markov Decision Processes

Markov Decision Processes (MDPs) [Bellman, 1957] are the processes underlying reinforcement learning problems. An MDP models a decision-making process in a stochastic environment where the agent has some control over the state transitions. An MDP can be described by a tuple ⟨S,A, P,R⟩, where S contains all states, A the set of actions an agent can take, Ps,s′ ,a the probabilities of transitioning from state s to state s′ when taking action a, and Rs,s′ ,a the reward the agent receives when going from state s to state s′ under action a. When solving an MDP, we want to find a policy π : S → A such that, when executing its actions, the expected sum of rewards (the return) is maximized. In this work we define policies (w.l.o.g.) as a mapping from states and actions to an indicator for whether or not action a is taken in state s: π : S × A → {0, 1}.

Value Iteration

When solving MDPs we generally discount future rewards in each step by a user-defined value of 0 < γ < 1 to ensure that the optimal policy will generate a finite return. The most common approach for optimally solving MDPs is by using one of many dynamic programming variants. In this work, we focus on value iteration. Value iteration finds a value Vs for each state s that holds the expected discounted return for taking optimal greedy actions starting from that state. These values can be found by iteratively updating Vs until the Bellman equation [Bellman, 1957]

is approximately satisfied. We will refer to this optimal solution found with value iteration as the unrestricted optimal solution as the computed policy can be arbitrarily complex.

3 RelatedWork

3.1 Learning Decision Tree Policies

Decision trees have appeared at various parts of the reinforcement learning pipeline [Glanois et al., 2021] for example in modeling the Q-value function or in modeling the policy. In this work, we are interested in modeling the policy with a decision tree as this gives us an interpretable model that we can directly use at test time.

Decision trees make predictions using hard comparisons between a single feature value and a threshold. The learned models can be discontinuous and non-differentiable, which makes optimization with gradients challenging. One line of research focuses on overcoming the non-differentiability of decision trees to allow for the use of gradient-based optimization methods. Gupta et al. [2015] train decision trees that contain linear models in their leaves that can be optimized with gradients but this results in models that are hard to interpret. Silva et al. [2020] first relax the constraints that decision nodes select one feature, that leaves predict one value, and that thresholds are hard step functions. Such relaxed trees are differentiable and can be trained with policy gradient methods. By discretizing the relaxed tree they end up with an interpretable model that approximates the relaxed model. However, the relaxed tree can get stuck in local minima and the discretized tree offers no performance guarantees. Likmeta et al. [2020] consider decision tree policies for autonomous driving tasks. To make the decision tree parameters easily optimizable, they fix the tree structure along with the features used in the decision nodes. By learning a differentiable hyper-policy over decision tree policies they are then able to approximately optimize the models with gradient descent.

With Iterative Bounding MDPs, Topin et al. [2021] reformulate the underlying MDP into one where the agent implicitly learns a decision tree policy. The method can be thought of as a tree agent learning to take actions that gather information and a leaf agent learning to take actions that work well given the gathered information of the tree agent. By reformulating the MDP, its decision tree policy can be optimized using differentiable function approximators and gradient-based optimizers.

In a separate line of work, the goal is to represent a specific, usually optimal, policy as a decision tree that is unbounded in size. These techniques have been developed for policies with a single goal state [Br´azdil et al., 2015] and as a tool for general controllers [Ashok et al., 2020]: dtcontrol.

Imitation Learning (VIPER)

Instead of directly optimizing a decision tree, one can also try to extract a decision tree policy from a more complex teacher policy using imitation learning. These imitation learning algorithms turn reinforcement learning into a supervised learning problem for which we have successful decision tree learning algorithms [Breiman et al., 1984; Quinlan, 1986]. DAGGER [Ross et al., 2011] (dataset aggregation) is an algorithm that iteratively collects traces from the environment using its current policy and trains a supervised model on the union of the current and previous traces. Since DAGGER only uses information on the predicted action of the teacher policy, it ignores extra information on Q-values that modern Q-learning algorithms provide. VIPER [Bastani et al., 2018] focuses on learning decision trees and improves on DAGGER by including Q-value information into the supervised learning objective. While VIPER generates significantly smaller decision trees than DAGGER, we will show that these trees are not yet optimal with respect to the tradeoff in size and performance.

3.2 Optimal Decision Trees

The standard algorithms for training decision trees in supervised learning are greedy heuristics and can learn trees that perform arbitrarily poorly [Kearns, 1996]. Therefore in recent years there has been increasing interest in the design of algorithms that train decision trees to perform optimally. Early works formulated training decision trees for classification and regression and used methods such as dynamic programming to find optimal decision trees [Nijssen and Fromont, 2007]. Mixed-Integer Linear Programming based formulations [Bertsimas and Dunn, 2017; Verwer and Zhang, 2017] have since become popular. These methods are flexible and have been extended to optimize performance under fairness constraints [Aghaei et al., 2019] or directly optimize adversarial robustness [Vos and Verwer, 2022]. Generally, the size of the tree is limited to provide regularization and aid interpretability, then the solver is tasked with finding a decision tree that maximizes training performance given the size limits. The field has since worked on increasingly efficient optimization techniques using a variety of methods such as MILP [Verwer and Zhang, 2019], dynamic programming [Demirovi´c et al., 2020; Lin et al., 2020], constraint programming [Verhaeghe et al., 2020], branch-andbound search [Aglin et al., 2020; Aglin et al., 2021] and Boolean (maximum) satisfiability [Narodytska et al., 2018; Hu et al., 2020; Schidler and Szeider, 2021].

4 OMDT: Optimal MDP Decision Trees

As a first step in bridging the gap between optimal decision trees for supervised and reinforcement learning, we introduce OMDTs: Optimal MDP Decision Trees. OMDT is a Mixed-Integer Linear Programming formulation that encodes the problem of identifying a decision tree policy that achieves maximum discounted return given a user-defined MDP and tree size limit. Our formulation can be solved using one of many available solvers, in this work we use the state-of-theart solver Gurobi 1 .

Intuitively the OMDT formulation consists of two parts: a (dual) linear programming formulation for solving MDPs and a set of constraints that limits the set of feasible policies to decision trees. Figure 2 summarizes OMDT's formulation in natural language. All the notation used in OMDT is summarized in Table 1.

4.1 Constraints

It is well known that MDPs can be solved using linear programming, the standard linear program is [Puterman, 2014]:

It is not easy to efficiently add constraints to this formulation to enforce the policies to be size-limited decision trees because it reasons abstractly over policies, i.e. by reasoning over the policy's state values. To create a formulation for decision tree policies, we resort to the standard dual program:

1 https://www.gurobi.com/

This program uses a measure xs,a of how often the agent takes action a in state s. This allows us to add efficient constraints that control the policy of the agent. Intuitively the program maximizes the rewards ∑s′ Ps,s′ ,aRs,s′ ,a weighted by this xs,a. The constraints enforce that the frequency by which a state is exited is equal to the frequency that the agent is initialized in the state p0(s) or returns to it following the discounted transition probabilities γPs,s′ ,a.

To enforce the policy to be a size-limited decision tree we will later constrain the xs,a values to only be non-zero when the agent is supposed to take action a in state s according to a tree policy. We will first introduce the variables and constraints required to model the decision tree constraints.

Modeling Decision Nodes

Our decision tree formulation is roughly based on OCT [Bertsimas and Dunn, 2017] and ROCT [Vos and Verwer, 2022], MILP formulations for optimal (OCT) and robust (ROCT) classification trees. In these formulations, the shape of the decision tree is fixed. Like ROCT, we describe a decision node m by binary threshold variables bm,j,k, indicating whether the kth threshold of feature j is chosen. 2 Unlike ROCT, we only allow one of these variables to be true over all features and possible thresholds:

We follow paths through the tree to map observations to leaves. In each node mwe decide the direction ds,m that the observation of state s takes (left=0 or right=1 of the threshold k). ROCT uses two variables per state-node pair to model directions ds,m to account for perturbations in the observations. Since we are optimizing for an MDP without uncertainty in the observations we only require one variable ds,m per statenode pair.

We further improve over ROCT by determining ds,m using only one constraint per state-node pair instead of a separate constraint per state-node-feature triple. For this, we pre-compute a function side(s, j, k) which indicates for each feature-threshold pair (j, k) and every observation s the side of k that s is on for feature j (left=0 or right=1), i.e. whether Xsj > k holds. This formulation is not limited to the predicates '≤' or '>' however and can be easily extended to other predicates in the pre-computation of side(s, j, k). The following then forces the direction ds,m to be equal to the direction of the indicated threshold:

2 In practice, the possible values for threshold k depends on the chosen feature j. We do not model this for convenience of notation.

The variables ds,m represent the direction of an observation's path at a decision node. Together, the d variables allow us to follow an observation's path through the tree which we use to identify the leaf that it reaches. Important in this formulation, compared to existing binary encodings, is that it requires no big-M constraints to describe these paths. This makes the relaxation stronger and therefore the solver gives much better bounds than using the big-M style formulations from ROCT.

Modeling Policy Actions

Decision leaves only have one set of binary decision variables: ct,a encoding whether or not leaf t predicts action a. We want each leaf to select exactly one action:

As mentioned before we can follow an observation's path through the tree by using their ds,m path variables. One can linearize an implication of a conjunction of binary variables as follows:

If an observation reaches leaf t and the leaf predicts action a then we want to force the policy πs,a to take that action in the associated state s. Using the aforementioned equivalence we add the constraint:

This constraint forces the agent to take the action indicated by the leaf. To prevent the agent from taking other actions that were not indicated we force it to only take a single action in each state (giving a deterministic policy):

Now we have indicators πs,a that mark what action is taken by the agent. To link this back to the MDP linear programming formulation that we use to optimize the policy, we set the xs,a variables. We need to set xs,a = 0 if πs,a = 0, else xs,a should be unbounded. We encode this using a big-M formulation:

M should be chosen as small as possible, but larger or equal to the largest value that xs,a can take. We use the fact that we are optimizing the MDP using discount factor γ to compute an upper bound on xs,a and set M = 1 1-γ , proof is given in the appendix.

4.2 Complete Formulation

The runtime of MILP solvers grows worst-case exponentially with respect to formulation size so it is important to limit the scale of the formulation. The number of variables in our formulation grows with O(|S||J||TD|+|A||TL|+|S||A|) which

a ∈ A set Set of actions in MDP

s ∈ S set Set of states in MDP

i=1..|S| set Observation and state indices

j ∈ J set Set of feature indices

k = 1..K set Indices of all possible feature thresholds

m ∈ TD set Set of decision nodes in the tree

t ∈ TL set Set of leaves in the tree

A(t)

set Set of ancestors of leaf t

Al(t)

set

... that have t in their left path

Ar (t)

set

... that have t in their right path

Table 1: Summary of notation used in OMDT.

follows from their indices in Table 1. The number of constraints grows with the order O(|S||TD|+|S||A||TL|) as it is dominated by the constraints that determine ds,m at each node (Equation 4) and constraints that force πs,a according to the tree (Equation 6). We summarize OMDT below:

s.t.

5 Results

We present experiments comparing the performance of OMDTs with VIPER and dtcontrol. Viper uses imitation learning to extract a size-limited decision tree from a teacher policy and dtcontrol learns an unrestricted tree that exactly copies the teacher's behavior. To provide a fair comparison we have trained VIPER and dtcontrol with an unrestricted optimal teacher by first solving the MDP with value iteration and then extracting all Q values, both methods ran with default parameters. We also implemented and ran experiments on interpretable Differentiable Decision Trees [Silva et al., 2020] but excluded these models from our analysis as they did not outperform a random policy. The full code for OMDT and our experiments can be found on GitHub 3 . All of our experiments ran on a Linux machine with 16 Intel Xeon CPU cores and 72 GB of RAM total and used Gurobi 10.0.0 with default parameters. Each method ran on a single CPU core.

5.1 Environments

For comparison we implemented 13 environments based on well-known MDPs from the literature, the sizes of these MDPs are given in Table 2. All MDPs were pre-processed such that states that are unreachable from the initial states are removed. We briefly describe the environments below but refer to the appendix for complete descriptions.

In 3d navigation the agent controls a robot in a 5×5×5 world and attempts to reach from start to finish with each voxel having a chance to make the robot disappear. blackjack is a simplified version of the famous casino game where we assume an infinite-sized deck and only the actions 'skip' or 'hit'. frozenlake is a grid world where the agent attempts to go from start to finish without falling into holes, actions are stochastic so the agent will not always move in the intended direction (e.g. the action 'up' will only not send the agent 'down'). inventory management models a company that has to decide how many items to order to maximize profit while minimizing cost. system administrator refers to a computer network where computers randomly crash and an administrator has to decide which computer to reboot. A crashed computer has an increased probability of crashing a neighboring computer. tictactoe vs random is the well-known game of tic-tac-toe when played against an opponent that makes random moves. In tiger vs antelope the agent attempts to catch an antelope that randomly jumps away from the tiger in a grid world. traffic intersection describes a perpendicular intersection where traffic flows in at different rates and the operator decides when to switch the traffic lights. xor is an MDP constructed with states randomly distributed on a plain, the agent gets 1 reward for taking the action according to an XOR function and -1 for a mistake. The XOR problem is notoriously difficult for greedy decision tree learning algorithms.

5.2 Performance-Interpretability Trade-off

It is often assumed that there is a trade-off in the performance and interpretability of machine learning models [Gunning and Aha, 2019], since interpretable models necessarily lack complexity but this assumption is not always true [Rudin, 2019]. We aim to answer whether the performance-interpretability trade-off occurs in a variety of MDPs by training size-limited decision trees and comparing their performance to the optimal solutions that were not restricted in complexity. We visualize the normalized return of depth 3 OMDTs and unrestricted dtcontrol trees in Figure 3. Returns were normalized such that 0 corresponds to a random policy and 1 to an optimal one. Since small deterministic decision tree policies are limited in the number of distinct actions an optimal tree can perform worse than a random policy. Experiments were repeated 3 times and runs were limited to 2 hours. We consider an OMDT optimal when the relative gap between its objective and bound is proven to be less than 0.01%.

3 https://github.com/tudelft-cda-lab/OMDT

While it is debatable what the precise size limits are for decision trees to be interpretable [Lipton, 2018] we use trees of depth 3 which implies that a tree has at most 8 leaves. Note that this also limits the number of distinct actions in the policy to 8. We find that in all environments, OMDTs of depth 3 improve on the performance of random policies, and in 8 out of 13 environments the policy gets close to optimal. Decision trees trained with dtcontrol always achieve the optimal normalized return of 1 since they exactly mimic the optimal policy. However, dtcontrol produces large trees that are not interpretable to humans. When run on 3d navigation for example, dtcontrol produces a tree of 68 decision nodes which is very complex for humans to understand. OMDT produces a tree of 7 decision nodes which performs equally well.

Overall, our results demonstrate that for small environments there is no performance-interpretability trade-off: simple policies represented by size-limited trees perform approximately as well as the unrestricted optimal policy.

5.3 Direct Optimization versus Imitation Learning

The above conclusion holds when the policy is learned to optimality under the constraint that it has to be a small tree, e.g., using OMDT. Techniques such as VIPER enforce the size constraint but aim to imitate the unrestricted optimal policy. We now show that this comes at a cost when the unrestricted policy is too complex to be represented using a small tree.

VIPER trains trees by imitating high Q values of the optimal policies while OMDT directly maximizes expected return. In Table 2 we list the normalized return (0 for random policies, 1 for optimal policies) for VIPER and OMDT with respectively 5 minutes and 2 hours of runtime. After 5 minutes OMDT improves performance over random policies but often needs more time to improve over VIPER. After 2 hours OMDT's policies win on 11 out of 13 environments. For instances with large state space such as tictactoe vs random OMDT needs more than 2 hours to improve over VIPER.

Shortcomings of Imitation Learning

Overall, given sufficient runtime, OMDT produces better policies than VIPER. This cannot be easily solved by giving VIPER more runtime but is an inherent problem of imitation learning. To illustrate this, we investigate the results on the frozenlake MDPs as Table 2 demonstrates that imitation learning can perform far from optimal in these environments.

In theory, imitation learning performs optimally in the limit [Ross et al., 2011] but this result requires the student policy to be as expressive as the teacher policy. This is not the case for size-limited decision trees. When VIPER learns a policy for frozenlake 12x12 it tries to imitate a complex policy using a small tree that cannot represent all teacher policy actions. This results in VIPER spending capacity of its decision tree on parts of the state space that will never be reached under its student policy. In Figure 4 we visualize the paths that the agents took on 10,000 runs and indicate the policies with arrows. VIPER creates leaves that control action in the right section of the grid world (indicated in red). The optimal teacher policy often visits this section but the simple student does not. By directly optimizing a decision tree policy using OMDT, the policy spends its capacity on parts of the state space that it actually reaches. As a result, VIPER cannot prevent actions that send its agent into holes on the left part of the grid world (indicated in red). OMDT actively avoids these.

5.4 Runtime

Runtime for solving Mixed-Integer Linear Programming formulations scales worst-case exponentially which makes it important to understand how solvers operate on complex formulations such as OMDT. We solved OMDTs for a depth of 3 for a maximum of 2 hours and display the results in Table 2. The table compares the runtimes of VIPER and solving OMDT to optimality. If the solver does not prove optimality within 2 hours we denote it as 'timeout'. We also denote the number of possible decision tree policies computed as:

|TB|possible splits ×|TL||A| . It provides an estimate of howmany decision tree policies are possible and shows that enumerating trees with brute force is intractable.

OMDT solves a simple environment such as Frozenlake 4x4 (16 states, 4 actions) to optimality within 2 seconds but runtime grows for larger environments such as inventory management (101 states, 100 actions) which took on average 2533 seconds. VIPER needs roughly 2250 seconds of runtime for every MDP and runs significantly faster on some MDPs. This is because VIPER spends much time evaluating policies on the environment and some environments quickly reach terminal states which results in short episodes. While OMDT was able to prove optimality on only 7 out of 13 environments within 2 hours, OMDT finds good policies before this time on 12 out of 13 environments.

6 Conclusion

We propose OMDT, a Mixed-Integer Linear Programming formulation for training optimal size-limited decision trees for Markov Decision Processes. Our results show that for simple environments such as blackjack, we do not have to trade off interpretability for performance: OMDTs of depth 3 achieve near-optimal performance. On Frozenlake 12x12, OMDT outperforms VIPER by more than 100%.

OMDT sets a foundation for extending supervised optimal decision tree learning techniques to the reinforcement learning setting. Still, OMDT requires a full specification of the Markov Decision Process. Imitation learning techniques such as VIPER can instead also learn from a simulation environment. Therefore, future work should focus on closing the gap between the theoretical bound supplied by OMDT and the practical performance achieved by algorithms that require only simulation access to optimize interpretable decision tree policies in reinforcement learning. Additionally, future work can incorporate factored MDPs into OMDT's formulation to scale up to larger state spaces.