ODTlearn: A Package for Learning Optimal Decision Trees for Prediction and Prescription

Summary

ODTLearn is an open-source Python package that provides methods for learning optimal decision trees for high-stakes predictive and prescriptive tasks based on the mixed-integer optimization (MIO) framework proposed in (Aghaei et al., 2021) and several of its extensions. The current version of the package provides implementations for learning optimal classification trees, optimal fair classification trees, optimal classification trees robust to distribution shifts, and optimal prescriptive trees from observational data. We have designed the package to be easy to maintain and extend as new optimal decision tree problem classes, reformulation strategies, and solution algorithms are introduced. To this end, the package follows object-oriented design principles and supports both commercial (Gurobi) and open source (COIN-OR branch and cut) solvers.

Automated data-driven predictive and prescriptive methods are increasingly being used in high-stakes domains to inform and support decision-making. In such settings, these tools should be: (a) accurate (to minimize erroneous predictions/prescriptions that may negatively affect the populations on which they are deployed), (b) interpretable (so that predictions and decisions are transparent, accountable, and easy to audit), (c) flexible (i.e., possible to easily augment with domain specific constraints such as capacity and/or fairness constraints), and (d) robust (to ensure high-quality solutions even under adversarial shifts between training and deployment data).

Despite their popularity, decision trees (Breiman et al., 1984) are not necessarily well suited for data-driven decision-making in high-stakes domains. While the structure of decision trees makes them easy to interpret, they are typically constructed using heuristics and may yield suboptimal solutions. Furthermore, with heuristic-based decision trees, it is not immediately apparent how to incorporate relevant side information into the tree construction process or how to make a tree robust to adversarial shifts between training and deployment data. Optimal decision trees retain the interpretability of heuristic decision trees while providing the flexibility to model the types of problems decision-makers face and deliver optimal, trustworthy solutions.

Our ODTlearn Python package provides methods for fitting provably optimal decision trees using mixed-integer optimization (MIO) for various problem types and settings commonly encountered by practitioners in high-stakes settings. ODTlearn implements the FlowOCT modeling and solution paradigm proposed by Aghaei et al. (2021), which is significantly faster and provides better out-of-sample performance than previous MIO-based algorithms. This approach generalizes beyond standard classification problems to problems involving imbalanced datasets (e.g., by optimizing weighted accuracy or worst-case accuracy, by constraining recall or precision, or by balancing sensitivity and specificity). In addition to the core algorithm, we implement its generalization to learn optimal fair decision trees (FairOCT) that optimize accuracy while satisfying arbitrary domain specific fairness constraints such as statistical parity, conditional statistical parity, or equalized odds as proposed in (Jo et al., 2022). Experiments demonstrated that these methods outperform popular heuristic-based algorithms for learning fair trees. For users deploying decision trees in settings with potential distribution shifts between training and testing, we implement the MIP formulation proposed in (Justin et al., 2021) (RobustOCT) for learning optimal robust classification trees. Finally, we implement the MIP formulation proposed in Jo et al. (2022) (FlowOPT) for learning optimal prescriptive trees from observational data. This framework can be used to design treatment assignment policies in the form of decision trees, being highly interpretable while offering a tunable degree of personalization. Importantly, the learned trees can also be constrained to satisfy domain specific requirements such as budget constraints (e.g., limited amount of treatments) or fairness constraints (e.g., conditional statistical parity in allocation).

There are numerous high-quality implementations of heuristic-based methods for learning decision trees in R and Python, including scikit-learn (Pedregosa et al., 2011), Rpart (Therneau et al., 1997), and Caret (Kuhn et al., 2008) among many others. As discussed in Section 1, these methods are fast, but lack the modeling flexibility and optimality guarantees of MIO-based methods.

Several different R and Python packages provide implementations of methods for constructing optimal decision trees using approaches other than mixed-integer optimization. The evtree (Grubinger et al., 2014) Python package implements an evolutionary algorithm that searches the space of possible decision trees to find globally optimal decision trees. PyDL8.5 (Aglin et al., 2021) is an open-source Python package that implements the DL8.5 (Aglin et al., 2020) algorithm in C++ for learning optimal decision trees through the use of itemset mining techniques. The gosdt Python package provides a C++ implementation of Generalized Optimal Sparse Decision Trees (GOSDT) (Lin et al., 2020; McTavish et al., 2022), which employs a branch-and-bound method with dynamic programming, similar to DL8.5. GOSDT can directly model continuous features and construct decision trees for various objectives, including accuracy, AUC, and partial area under ROC convex hull. While the MIO formulations used in ODTlearn provide the flexibility to incorporate side constraints, they are generally slower than the methods used in evtree, DL8.5, and GOSDT, especially for larger datasets or deeper trees. These methods employ specialized algorithms, such as evolutionary algorithms or branch-and-bound with dynamic programming, which can often find optimal trees more quickly than general-purpose MIO solvers. However, it is not immediately obvious how to incorporate side constraints into the tree construction process for these methods, which is a strength of the MIO approach.

The policytree package (Sverdrup et al., 2020) provides an R implementation of Policy Trees (PT) (Zhou et al., 2022), which learn personalized treatment assignment policies from observational data using decision trees and doubly robust estimators to handle confounding. PT uses a recursive partitioning approach, allowing it to learn trees that handle continuous and discrete covariates without pre-discretization and in less time than the FlowOPT MIO formulation. However, PT’s recursive partitioning approach may not converge to the global optimum and cannot easily incorporate complex constraints.

There are also open-source implementations of other MIO formulations for learning optimal classification trees. The authors of the BinOCT MIO formulation (Verwer and Zhang, 2019) provide an open-source implementation of their method. Additionally, Tang (2021) provides open-source implementations of the BinOCT, OCT (Bertsimas and Dunn, 2017), and FlowOCT (Aghaei et al., 2021) MIO formulations. While both projects provide valuable implementations of these methods for technical users and researchers, these projects are less accessible to a general audience with limited knowledge of MIO. Neither project distributes its code as a Python package, and they are not indexed on PyPi. Furthermore, their implementations rely exclusively upon proprietary MIP solvers such as CPLEX or Gurobi, which require paid or academic licenses to use with large-scale problems.

In comparisons on simulated data, the FlowOCT and FlowOPT formulations solve problem instances faster and are more accurate out-of-sample compared to existing MIO formulations. In Aghaei et al. (2021), the authors implement FlowOCT, BinOCT, and OCT in Gurobi to directly compare the in-sample and out-of-sample performance of the methods. When comparing the speed at which each method solves a fixed number of problem instances, they find that FlowOCT is 29 times faster than the other two formulations at solving problems with balanced decision trees. This significant speedup is attributed to the stronger linear optimization relaxation of the FlowOCT formulation, which leads to more effective pruning of the branch-and-bound tree during the solving process. Regarding out-of-sample accuracy, FlowOCT improves accuracy compared to BinOCT and OCT by up to 7% and 21%, respectively. In Jo et al. (2021), the authors compare FlowOPT to Gurobi implementations of the MIO formulations proposed in Kallus (2017) (K-PT) and Bertsimas et al. (2019) (B-PT). FlowOPT outperforms K-PT and B-PT in terms of the out-of-sample probability of correct treatment assignment across various experiment designs.

Interpretable AI (Interpretable AI, 2022) offers a proprietary Julia (Bezanson et al., 2017) package with a Python wrapper that provides methods for learning optimal classification, prescriptive, and regression trees. These implementations use the heuristic algorithm based on local search proposed by Bertsimas and Dunn (2019) (LST). On the one hand, LST is much faster than MIO formulations like FlowOCT, requiring only seconds to find a local optimum. On the other hand, MIO formulations are more flexible and able to solve instances to provable optimality.

In summary, ODTlearn distinguishes itself from existing implementations of optimal decision tree methods in several key ways. First, it provides open-source, user-friendly implementations of state-of-the-art MIO formulations for learning optimal decision trees, including FlowOCT and FlowOPT, which have been shown to outperform other MIO formulations in terms of solution speed and out-of-sample performance. Second, ODTlearn is distributed as a Python package indexed on PyPi, making it easily accessible to a broad audience. Third, ODTlearn supports open-source and proprietary MIP solvers, allowing users without access to commercial solvers to still benefit from the package. Finally, the MIO formulations in ODTlearn provide the flexibility to incorporate custom constraints and objectives, enabling users to tailor the learned decision trees to specific requirements of the problem they wish to solve.

The software architecture of ODTlearn is motivated by the optimal decision tree literature in which researchers have recently proposed numerous new problem classes, reformulation strategies, and solution approaches. With this in mind, we have created a class structure for our package, summarized in Figure 1, that follows the SOLID principles of object-oriented programming for developing software (Martin, 2003). These principles emphasize structuring classes that are easy to maintain and extend. The Single Responsibility Principle (SRP) states that each class should have a single responsibility or purpose. The Open-Closed Principle (OCP) requires that users should be able to extend a class’s behavior without modifying its existing code. According to the Liskov Substitution Principle (LSP), derived classes should be able to replace their base classes without affecting correctness. The Interface Segregation Principle (ISP) asserts that users should not be forced to depend on methods they do not use. Finally, the Dependency Inversion Principle (DIP) suggests that software design should rely upon abstract interfaces rather than concrete implementations. These principles establish practices for developing software with considerations for maintaining and extending it as the project grows.

Following the dependency inversion principle, all classes in ODTlearn are derived from our abstract base class OptimalDecisionTree. This class provides a standardized interface for the two types of trees currently supported (classification and prescription) while keeping their implementation details separate to follow the interface segregation principle. The OptimalDecisionTree abstract class allows our class structure to follow the open-closed principle as we can create the derived classes OptimalClassificationTree and OptimalPrescriptiveTree with problem-specific methods for traversing and visualizing the decision tree. Next, the children of OptimalClassificationTree and OptimalPrescriptiveTree implement methods for creating the decision variables, the constraints, and the objective function necessary for constructing the optimization problem of interest. The separate classes for each of the variations of the MIO formulations ensure that each of the derived classes has a single job, obeying the single responsibility principle. Finally, the classes in the third and fourth levels of Figure 1 implement user-facing methods such as fit and predict required by their parent classes. This ensures that our low-level classes respect the Liskov substitution principle by adhering to the basic requirements of their parent class. Thus, our adherence to the SOLID principles ensures that researchers and practitioners building upon ODTlearn can easily augment it with more features (e.g., different objectives or additional constraints) or even build new types of trees.

ODTlearn is built upon the Python MIP package (Santos and Toffolo, 2020), which provides a unified interface for implementing MIO formulations with callbacks across both open-source and commercial solvers. By leveraging Python MIP’s consistent interface, our package is able to support both the open-source CBC solver and the commercial Gurobi solver without the need for solver-specific code. This design choice allows us to focus on implementing the core functionality of ODTlearn while still supporting a wide range of users with different solver preferences and access to commercial licenses. The use of Python MIP is particularly important for our BendersOCT formulation, which relies on callbacks to implement the Benders’ decomposition algorithm. Callbacks are essential for the efficient solution of large-scale problems using BendersOCT, and Python MIP’s interface allows us to implement these callbacks in a solver-independent manner. Furthermore Python MIP interacts directly with the solver’s native dynamic loadable library (DLL) interface, ensuring that any computationally intensive steps are performed by the solver and not in Python.

For fitting optimal classification and prescriptive trees, we provide users with an interface that follows a common fit-predict structure. Figure 2 provides a code snippet demonstrating how to use the ODTlearn API to fit a fair optimal classification tree. Once an optimal decision tree has been learned, users can employ the build in plot tree function to visualize the tree.

Releases of the ODTlearn package are available via PyPI at https://pypi.org/project/odtlearn. The package source code and documentation are hosted on GitHub (https://github.com/D3M-Research-Group/odtlearn). Collaboration in the form of discussions, feature requests, or bug reports is made possible through the GitHub issue and pull request workflow. We have implemented continuous integration through GitHub Actions to ensure backward compatibility and quickly identify any code regressions. Our documentation includes installation instructions, a user guide, an API reference, and downloadable example notebooks demonstrating each of the classification methods implemented in the package. The documentation is hosted via GitHub pages at https://d3m-research-group.github.io/odtlearn. The package is distributed under the GPL-3.0 license and makes use of several core libraries within Python’s scientific computing ecosystem: scikit-learn (Pedregosa et al., 2011), numpy (Harris et al., 2020), and pandas (McKinney et al., 2010).

Based on the paper, I can implement the following specific improvements to an AI system:

1. Optimal Decision Tree Learning via Mixed-Integer Optimization (MIO)

• Replace heuristic tree-building (e.g., CART, ID3) with the FlowOCT formulation for classification, which guarantees provably optimal trees (not just locally optimal ones) and provides 29× faster solving than prior MIO methods (BinOCT, OCT).

• The improved system can now solve classification problems with weighted accuracy, worst-case accuracy, recall/precision constraints, or balanced sensitivity/specificity—all within a single optimization framework, rather than requiring ad-hoc post-processing.

2. Fair Classification Trees (FairOCT)

• Integrate the FairOCT formulation to enforce arbitrary fairness constraints (statistical parity, conditional statistical parity, equalized odds, equal opportunity) directly into the tree optimization, not as a post-hoc repair.

• The improved system can now output a decision tree that simultaneously maximizes accuracy and satisfies a user-specified fairness bound (e.g., fairness bound=1), which is impossible with standard scikit-learn trees that ignore fairness.

3. Robust Classification Trees (RobustOCT)

• Implement the RobustOCT MIP formulation to handle adversarial distribution shifts between training and deployment data.

• The improved system can now guarantee a minimum accuracy level even under worst-case perturbations of feature distributions, making it suitable for high-stakes deployment where data drift is expected (e.g., medical diagnosis, credit scoring).

4. Prescriptive Trees (FlowOPT) for Treatment Assignment

• Add the FlowOPT formulation to learn optimal treatment assignment policies from observational data (not just predictions), with support for inverse probability weighting (IPW), direct method (DM), and doubly robust (DR) estimators.

• The improved system can now output an interpretable decision tree that prescribes the best action (e.g., which drug to prescribe, which loan to approve) for each individual, while respecting budget constraints (limited treatments) and fairness constraints (conditional statistical parity in allocation)—a capability not present in standard classification or regression trees.

5. Solver Flexibility and Accessibility

• Use the Python MIP interface to support both open-source (CBC) and commercial (Gurobi) solvers, eliminating the dependency on proprietary solvers that prior MIO implementations (e.g., BinOCT, OCT) required.

• The improved system can now run optimal tree learning on any machine without a paid license, making it accessible to smaller organizations and researchers with limited budgets.

6. Object-Oriented, Extensible Architecture

• Adopt the SOLID-principled class hierarchy (e.g., OptimalDecisionTree → OptimalClassificationTree → FlowOCT, FairOCT, RobustOCT; OptimalPrescriptiveTree → FlowOPT variants) so that new problem classes, reformulations, or solution algorithms can be added without modifying existing code.

• The improved system can now be extended by a developer to add, for example, a regression tree or a tree with capacity constraints, by simply subclassing the appropriate base class and implementing the create variables, create constraints, and create objective methods—no need to rewrite the fitting or prediction logic.

7. Built-in Visualization and API Consistency

• Provide a unified fit/predict/plot tree interface across all tree types, so the improved system can be used as a drop-in replacement for scikit-learn's DecisionTreeClassifier but with optimality and constraint guarantees.

• The improved system can now produce a visualizable, interpretable tree (via matplotlib) that is provably optimal under the specified objective and constraints, enabling direct auditability by domain experts.

What the improved AI system can do specifically:

• For a hospital: Learn a fair, optimal treatment assignment tree that maximizes patient health outcomes while enforcing equalized odds across racial groups and a limited total treatment budget—something no heuristic tree can guarantee.

• For a bank: Learn a robust credit-scoring tree that maintains a minimum approval accuracy even if the applicant population shifts (e.g., due to economic downturn), with a provable worst-case performance bound.

• For a hiring platform: Learn an optimal classification tree that maximizes hiring accuracy subject to statistical parity constraints, with a certificate of optimality (i.e., the solver proves no better tree exists within the depth limit).

• For a researcher: Run all of the above on a laptop with only the open-source CBC solver, and extend the code to add a new fairness metric or a new objective (e.g., F1-score) in under 50 lines of code, without breaking existing functionality.