Policy Trees in Reinforcement Learning

• Policy trees are binary structures where internal nodes apply state-dependent rules and leaves encode actions or skills, offering transparent decision-making in RL.
• They leverage differentiable splits and soft routing to enable end-to-end gradient-based training and maintain performance in high-dimensional environments.
• Empirical studies show that policy trees can match neural models in performance while ensuring rapid inference, model auditability, and robust safety verification.

A policy tree in reinforcement learning (RL) defines a hierarchical structure in which decisions are made via a sequence of "if-then" rules, typically organized as a binary tree whose internal nodes represent state-dependent branchings and whose leaves encode either discrete actions, distributions over actions, or high-level options/skills. Policy trees are pursued as an interpretable alternative to neural network policies in RL, with contemporary approaches extending from axis-aligned "hard" trees to fully differentiable "soft" variants and hierarchical skill-trees. Recent advances leverage policy gradients, imitation learning, distillation, and direct optimization to simultaneously achieve competitive performance and policy auditability in both discrete and continuous high-dimensional settings.

1. Fundamental Definition and Taxonomy

Policy trees are full binary trees where each internal node applies a decision rule to the agent's current state representation. Formally, an internal node computes a split predicate—axis-aligned or linear—e.g., "$s_j \leq v$", or $w^\top s + \phi$, and recursively routes the state to child nodes. Leaves specify output labels—actions (for discrete domains), distributions (via softmax over actions), or skill indices. In differentiable policy trees, node gating employs a sigmoid or similar, enabling smooth routing and stochastic mixtures over leaf actions. This architecture subsumes classical hard CART trees, soft decision trees (SDT), differentiable decision trees (DDT), mixed types (cascading, hierarchical), and symbolic trees as in SYMPOL (Marton et al., 2024), CDT (Ding et al., 2020), and SkillTree (Wen et al., 2024).

SkillTree introduces a hierarchical policy tree with a high-level differentiable decision tree $\pi_h(k|s)$ that selects among $K$ discrete skills, each represented by a continuous skill embedding $z \in \mathbb{R}^D$. A low-level neural policy $\pi_l(a|s, z)$ executes actions conditioned on both state and skill embedding for $h$ steps (Wen et al., 2024). This structure generalizes discrete-action trees and option-based trees and is critical for explainable decision-making in long-horizon continuous control.

2. Training Methodologies and Optimization

Policy trees are trained with methods varying by differentiability, decision type, and RL objective:

• Direct Policy-Gradient Optimization: SYMPOL (Marton et al., 2024), DDT (Silva et al., 2019), DTPO (Vos et al., 2024), and SkillTree (Wen et al., 2024) embed tree structure directly into the policy parameterization and optimize via PPO-style stochastic policy gradients. Differentiable splits enable backpropagation through threshold and feature-selection parameters, using straight-through estimators or soft rounding for non-differentiable components.

Example from SYMPOL (Marton et al., 2024):

$$L_{\text{CLIP}}(\theta) = \mathbb{E} \left[ \min( r_\theta(s,a) \cdot \hat{A}, \, \text{clip}(r_\theta(s,a), 1-\epsilon, 1+\epsilon) \cdot \hat{A} ) \right]$$

with end-to-end gradient flow through axis-aligned splits implemented as hardmax.

• Imitation Learning and Distillation: VIPER (Bastani et al., 2018), MSVIPER (Roth et al., 2022), Dpic (Li et al., 2021), Distill2Explain (Gokhale et al., 2024), and CDT (Ding et al., 2020) first train a high-performing neural (or ensemble) "expert" policy, then fit decision trees by minimizing imitation or reward-weighted losses. VIPER applies Q-weighted aggregation for robust state coverage; Dpic and MSVIPER incorporate advantage or multi-scenario sampling to improve fidelity and prevent performance degradation due to distribution shift.
• Regression-Formulated Policy Trees: In offline RL, decision trees can be framed as regression predictors for actions, conditioned on normalized return-to-go and timestep (RCDTP) (Koirala et al., 2024). XGBoost or similar methods fit ensembles of regression trees to minimize squared error between predicted and ground-truth actions.
• Conservative Q-Improvement: CQI (Roth et al., 2019) incrementally grows the policy tree, splitting leaves only when expected global cumulative reward gain $\Delta J$ exceeds a dynamically decaying threshold, yielding compact policy trees with a robust performance-size tradeoff.
• Hierarchical Skill Trees: SkillTree (Wen et al., 2024) initially fits the low-level policy and skill codebook via a VQ-VAE objective, then optimizes the high-level policy tree using a soft actor-critic loss penalized by KL divergence from a skill-prior, with explicit per-iteration updates for critic, entropy temperature, and policy parameters.

3. Differentiable Formulations and Gradient Flow

The expressivity and trainability of modern policy trees depend on soft routing mechanisms at internal nodes. Differentiable decision trees (DDT) replace hard splits with sigmoidal gating:

$p_\ell(x) = \sigma(\omega_j^i \cdot x + \phi_j^i)$

The probability of reaching a leaf is the product over the path (SkillTree (Wen et al., 2024), DDT (Silva et al., 2019)), forming a mixture over all leaves. Softmaxed leaf vectors yield action distributions. The entire tree structure is trained by propagating policy gradients (or imitation objectives) back through gating and leaf parameters, with the possibility of post-hoc discretization for interpretability.

Straight-through estimators allow hard, axis-aligned decisions within a differentiable framework (SYMPOL (Marton et al., 2024)), ensuring the learned parameters correspond exactly to the policy evaluated during training; this eliminates information loss observed in post-hoc conversion or discretization steps.

4. Explainability, Interpretability, and Verification

Policy trees provide transparent, interpretable decision-making at the path and leaf levels. Each internal split is a test over state-features, making paths auditable: "if $\omega_0 \cdot s + \phi_0 > 0.5$ then ... else ..." (SkillTree (Wen et al., 2024)). This facilitates rule extraction, analysis of feature importance, and direct mapping of policy actions to conjunctions of state predicates.

VIPER (Bastani et al., 2018) and MSVIPER (Roth et al., 2022) produce small axis-aligned trees suitable for formal verification. The structural constraints allow immediate safety, robustness, and stability guarantees via reachability analysis, SMT/SAT tools, and Lyapunov-based stability certificates. MSVIPER enables targeted modifications at the node/leaf level for correcting pathological behaviors (freezing, oscillation, vibration), providing efficiency metrics for policy improvement (e_O, e_R).

Feature-importance metrics (XGBoost's weight, gain, cover) directly quantify policy reliance on state features and time indices (RCDTP (Koirala et al., 2024)), while combinatorial skills in SkillTree are visualized via bar-graphs linking skill leaves to high-level subtask completion.

5. Empirical Evaluation and Performance Analysis

Across RL domains—classic control, robotic arm manipulation, battery energy management, navigation, and Atari—policy trees:

• Achieve performance matching or exceeding neural network baselines for moderate tree depths (d=3–6; leaf count L=4–16), e.g., CartPole, LunarLander, MountainCar, D4RL locomotion tasks (Koirala et al., 2024, Marton et al., 2024, Wen et al., 2024, Silva et al., 2019).
• Show rapid training and inference: sub-second tree fitting on CPUs, inference times $<2$ ms/step, outperforming transformer-based counterparts on both speed and explainability (Koirala et al., 2024).
• Retain performance when post-pruned or discretized, given sufficient tree depth and regularization, although information loss may occur if conversion is not end-to-end differentiable (CDT (Ding et al., 2020), SDT).
• In hierarchical option spaces, SkillTree matches skill-based neural networks, with transparent episodic traces of skill index selection and reuse (Wen et al., 2024).
• In domain-specific energy controllers, DDTs distilled from DQN outperform rule-based baselines by 20–25% in cost reduction (Gokhale et al., 2024).

Example learned tree (SYMPOL, MountainCarContinuous):

$$\pi(a|s) = \begin{cases} \mathcal{N}(a;-1.0,\sigma^2), & v\le0,\; x\le-0.10\ \mathcal{N}(a;-0.5,\sigma^2), & v\le0,\; x>-0.10\ \mathcal{N}(a;+0.5,\sigma^2), & v>0,\; x\le+0.10\ \mathcal{N}(a;+1.0,\sigma^2), & v>0,\; x>+0.10 \end{cases}$$

6. Limitations, Trade-offs, and Open Challenges

Policy trees offer immediate interpretability, formal verifiability, and fast inference, but are constrained in expressivity for highly complex or very high-dimensional continuous spaces unless augmented by feature-learning (CDT (Ding et al., 2020)), multi-level skill codebooks (SkillTree (Wen et al., 2024)), or cascaded trees. Discretization of soft trees may degrade test-time performance, unless direct optimization is applied (SYMPOL (Marton et al., 2024)).

Imitation learning-based policy trees can suffer instability in split choices and reproducibility unless advantage-weighting or scenario augmentation is employed (Dpic (Li et al., 2021), MSVIPER (Roth et al., 2022)). Convergence properties for tree RL depend on the exact splitting and evaluation criterion (CQI (Roth et al., 2019)), with explicit trade-off controls via split thresholds and decay rates.

Empirical findings recommend hard constraints on tree size for human auditability (leaf count L≈4–16 for simple domains, L≈64 for complex tasks), iterative rollback mechanisms to avoid forgetting, and advantage-weighted or Q-weighted training to focus on critical state-action decisions. Policy trees remain less expressive than deep neural networks for raw-image domains unless enriched by embedded learned features.

7. Hierarchical, Skill-Based, and Hybrid Tree Extensions

Recent work extends policy trees to hierarchical structures—SkillTree (Wen et al., 2024) combines differentiable trees at the high-level policy for skill selection with neural skill executors at the low-level, making discrete skill selection transparent in long-horizon control. Cascading Decision Trees (CDT (Ding et al., 2020)) integrate feature-learning subtrees before the main decision tree, increasing expressivity but maintaining interpretability.

Tree search-based planning algorithms (e.g., in AlphaZero) interleave finite-horizon policy trees with neural function approximators, where enhanced backup methods guarantee $\gamma^h$-contractivity and convergence (Efroni et al., 2018).

An encyclopedia table summarizing core models and training regimes:

Model/Framework	Primary Training Method	Explainability/Verification Features
SkillTree (Wen et al., 2024)	Differentiable tree + skill learning (SAC)	Skill-level decision audit, path-level transparency
SYMPOL (Marton et al., 2024)	End-to-end differentiable axis-aligned tree	Zero information loss, compact interpretable trees
VIPER/MSVIPER [(Bastani et al., 2018)/(Roth et al., 2022)]	Weighted imitation from DNN expert	Formal safety/stability verification (SMT/SOS), direct tree modification
DTPO (Vos et al., 2024)	Policy-gradient with regression-tree updates	Explicit size control, auditability via leaf count
CQI (Roth et al., 2019)	Conservative Q-improvement based tree growth	Direct reward-size tradeoff, succinct trees
CDT (Ding et al., 2020)	Feature-learning + soft decision tree	Expressive yet parameter-efficient, path auditability
RCDTP (Koirala et al., 2024)	Return-conditioned regression (offline RL)	Feature/importances, action histograms

Policy trees in RL constitute a rigorously studied and actively evolving alternative to opaque neural policies, achieving an overview of interpretability, auditability, and—through differentiability and hierarchical skill-space abstraction—competitive domain-level performance.