Optimizing Interpretable Decision Tree Policies for Reinforcement Learning

TL;DR

The paper tackles the lack of interpretability in reinforcement learning policies by introducing DTPO, a method that directly optimizes interpretable decision tree policies using gradient-based signals. Building on regression-tree heuristics and a PPO-inspired framework, DTPO iteratively improves a single decision tree to produce probabilistic action outputs, while a neural critic guides learning. DTPO is evaluated against VIPER and neural-policy baselines across a suite of control tasks and discrete MDPs, showing competitive performance and, in some cases, superior results when gradient-based optimization favors policy-gradient methods. The work demonstrates that small, human-interpretable trees can achieve near-optimal performance, offering a tangible, verifiable alternative to opaque neural policies with potential for safety-critical applications.

Abstract

Reinforcement learning techniques leveraging deep learning have made tremendous progress in recent years. However, the complexity of neural networks prevents practitioners from understanding their behavior. Decision trees have gained increased attention in supervised learning for their inherent interpretability, enabling modelers to understand the exact prediction process after learning. This paper considers the problem of optimizing interpretable decision tree policies to replace neural networks in reinforcement learning settings. Previous works have relaxed the tree structure, restricted to optimizing only tree leaves, or applied imitation learning techniques to approximately copy the behavior of a neural network policy with a decision tree. We propose the Decision Tree Policy Optimization (DTPO) algorithm that directly optimizes the complete decision tree using policy gradients. Our technique uses established decision tree heuristics for regression to perform policy optimization. We empirically show that DTPO is a competitive algorithm compared to imitation learning algorithms for optimizing decision tree policies in reinforcement learning.

Figures (3)

• Figure 1: High-level overview of the DTPO algorithm. The tree is initialized as a single leaf with equal probability for each action and iteratively refined by optimizing the loss with regression tree heuristics on batches of environment experience. In the end, we round the leaf values to obtain an interpretable deterministic policy.
• Figure 2: Undiscounted returns of policies with varying decision tree sizes. Each tree size was run with 6 random seeds, and best performing policies are highlighted. DTPO and VIPER perform similarly on average, but their performance varies depending on the environment. Simple environments such as Frozenlake4x4 reach optimal performance with policies of at most 8 leaves, while more complex problems like CartPoleSwingup require at least 128 leaves to reach approximately optimal performance.
• Figure 3: DTPO policies for the simple CartPole-v1 and much harder CartPoleSwingup environments, as environments increase in complexity, we need larger trees to achieve satisfactory performance. The explicit policy representation allows us to inspect them. This allows us to see, e.g., the symmetry in the CartPoleSwingup policy (notice the highlighted subtrees). While the tree for CartPoleSwingup is possibly too large to interpret exactly, one can still understand its parts and verify its properties. In contrast, a small PPO neural network policy typically consists of thousands of parameters and cannot easily be visualized.