Source-linked AI summary

Temporally Interpretable Differentiable Decision Trees

Eisuke Hirota, Aarav Sane, Rohan Paleja

arXiv:2610.10367v1cs.LGcs.RO

TL;DR

Sequential decision policies need transparency that helps people understand an agent’s short-term plans, not just its model structure. The paper adds action-chunked policy gradients and information-theoretic tree restructuring to DDTs. Warm-started trees match neural-network policies in three environments and come close in a fourth, while using up to 80% fewer parameters.

  • Problem

    In sequential decision-making, transparent models still need to show short-term plans that help users and engineers understand agent behavior.

    Source: §1 · p. 2
  • Method

    The paper introduces Temporal Ensemble and Temporal Prediction policy gradients for action-chunked DDTs and uses ITTR to restructure trees during training.

    Source: §1 · p. 2, §4.1 · p. 4, §6 · p. 9
  • Results

    Warm-started trees match neural-network policies in LK, IP, and LL, come close in LL-H, and use up to 80% fewer parameters.

    Source: §1 · p. 2, §6 · p. 9
  • Takeaways & Limitations

    Action chunks can reveal short-term plans that may increase user confidence and help engineers or researchers verify robustness and debug performance.

    Source: §1 · p. 2

    Training action-chunked trees from scratch remains challenging in highly stochastic domains.

    Source: §6 · p. 9

Abstract

from arXiv · show

Interpretability offers a solution to safe autonomy by providing transparency into an agent's underlying decision-making model. Within sequential-decision making tasks, differentiable decision trees (DDTs) are one approach to such interpretability, maintaining automatic-differentiable policies while providing humans with a discrete tree-based visualization. Nonetheless, current implementations of DDTs are not well-suited for sequential-decision making domains, as there exists an inherent mismatch between a tree's single-timestep behavior and a human's multi-timestep planning. Our work thus introduces time as a new dimension of interpretability, coined as temporal interpretability, and demonstrates how temporal abstractions via action chunking improve it. We achieve this by first introducing two novel policy gradient algorithms that incorporate action chunking. Additionally, to maintain parameter-efficient trees, we develop an information-theoretic tree restructuring algorithm that modifies the tree during training. Across four simulation environments, we find that warm-starting action chunked DDTs from a distilled action chunked policy is the most effective way to obtain temporally interpretable trees: they match neural network policies in three of the four domains while using up to 80$\%$ fewer parameters. Our code is available at https://github.com/ei5uke/temp-interp.

1. Introduction

Differentiable decision trees offer transparent policies, but their single-timestep actions do not convey the multi-step plans people use to understand sequential decisions. The paper introduces action-chunked policy gradients and tree restructuring; warm-started trees perform competitively while using fewer parameters.

1. Introduction

  • Current DDTs’ single-timestep actions do not provide the plan-level information humans use to interpret sequential decisions.Large trees also require extensive, potentially confusing review of the decision pipeline.Source: §1 · p. 1
  • Warm-started action-chunked DDTs from distilled action-chunked policies match state-of-the-art baselines competitively while using up to 80% fewer parameters.The authors report this as their most effective approach for obtaining temporally interpretable trees.Source: §1 · p. 2
  • Two novel policy-gradient methods incorporate action chunking so DDTs can show users their short-term plans at each timestep.The authors frame this as increasing behavior transparency.Source: §1 · pp. 1–2
  • ITTR dynamically deepens or prunes trees according to policy complexity, reducing tree size by up to 50%.The restructuring aims to counter the parameter growth introduced by action chunking.Source: §1 · p. 2

2. Related Works

The related work covers interpretable AI and action chunking, setting up the paper’s focus on making sequential policies’ short-term plans understandable.

2.1. Interpretable AI

  • Interpretable AI favors transparent logical models such as trees, whereas explainable-AI methods may offer post-hoc explanations without resolving black-box model opacity.The paper emphasizes model transparency and clear logical traces rather than post-hoc explanations.Source: §2.1 · p. 2
  • Temporal interpretability extends pre-deployment model transparency to online understanding, supporting debugging and better-informed mid-trajectory intervention.The paper identifies robotics control as one sequential-decision setting where this can matter.Source: §2.1 · p. 2

2.2. Action Chunking

  • Action chunking extends a policy to output H actions; Temporal Ensemble queries it each timestep and averages recent chunks, though enforced consistency may disadvantage stochastic environments.Open-loop chunking instead queries every H timesteps and relies on each output chunk.Source: §2.2 · pp. 2–3
  • The paper studies whether action-chunked policy gradients help agents depict short-term plans and whether Temporal Ensemble offers performance or interpretability benefits.Source: §2.2 · p. 3

3. Preliminaries

The preliminaries define the RL objective, explain how differentiable trees represent decisions, and introduce action-chunked policy gradients and their interpretability trade-off.

3.1. Problem Formulation

  • The problem is formulated as an MDP, with RL learning a policy that maximizes expected discounted return over trajectories.The MDP is specified by state and action spaces, transition and reward functions, a discount factor, and an initial-state distribution.Source: §3.1 · p. 3

3.2. Differentiable Decision Trees

  • DDTs replace binary decision points with fuzzy sigmoid splits, making tree policies differentiable for learning.Each split outputs the probability that a decision node evaluates TRUE from state features, node weights, and a splitting criterion.Source: §3.2 · p. 3
  • ICCT addresses the loss of tree quality during crispification by directly optimizing over a crisp tree with the straight-through trick.DDTs are generally converted to decision trees through crispification.Source: §3.2 · p. 3

3.3. Policy Gradient

  • Action-chunked policy gradients extend the policy to output H actions while leaving the value function unchanged, but querying only every H timesteps limits continuous plan visibility.Querying at every timestep instead may reduce performance and destabilize trajectories because training accounts for transitions at chunk intervals.Source: §3.3 · p. 3

4. Method

The method trains differentiable decision trees to expose short-horizon plans through action chunking, while restructuring trees during training to keep them compact. It introduces temporal ensemble and temporal prediction policy gradients, alongside Information-Theoretic Tree Restructuring.

4.1. Temporal Ensemble Policy Gradient

  • Temporal Ensemble Policy Gradient queries the chunking policy each timestep and averages actions assigned to the current timestep across the latest H chunks.This temporal ensemble keeps users supplied with the agent’s short-term plan; its weighted log-probability objective is biased when η < 1.Source: §4.1 · pp. 4–5
  • Temporal ensemble can hurt performance in highly stochastic environments, while a low EMA coefficient can make future actions contribute little to the displayed plan.Low EMA coefficients may also cause future chunks to output gibberish or noise.Source: §4.2 · p. 5

4.2. Temporal Prediction Policy Gradient

  • Temporal Prediction Policy Gradient uses the chunk’s first action for control and trains later actions to predict future actions for display.The combined objective targets task performance and accurate future-action predictions, providing short-term plans without enforcing temporal consistency.Source: §4.2 · pp. 5–6

4.3. Information-Theoretic Tree Restructuring

  • ITTR deepens trees when rewards are low and policy complexity is stalled, and prunes leaves with few visitations to keep representations compact.Deepening splits a highly used, low-variation leaf; pruning replaces its parent decision node with the sibling subtree.Source: §4.3 · p. 6

5. Experiments and Results

Warm-started action-chunked DDTs achieve the best or statistically indistinguishable returns among action-chunked trees while using fewer parameters than chunked MLPs; ITTR varies tree size during training.

5.1. Simulation Environments and Mutual Information

  • The four test environments are LK, IP, LL, and a more stochastic LL-H variant with added wind, turbulence, and action noise.Because the observation and action spaces are continuous, ITTR estimates mutual information with InfoNCELoss.Source: §5.1 · p. 6

5.2. Algorithms and Baselines

  • The comparison includes MLPs, CART trees, DDTs, and warm-started DDTs, each tested without chunking or with Temporal Ensemble or Temporal Prediction.The study uses ICCT-Complete trees, RPO, and a chunk horizon of H = 10.Source: §5.2 · pp. 6–7

5.3. Evaluation Results

  • Temporal Ensemble and Temporal Prediction do not significantly modify neural-network policy behavior, though action chunking increases parameter counts.Source: §5.3 · p. 7
  • Training chunked DDTs from scratch is difficult: DDT-PREDICTION can be seed-sensitive, while DDT-ENSEMBLE falls to 160.8 ± 13.1 in LL-H.In LL-H, chunked trees grow from about 75 parameters to 440–506 despite ITTR pruning them from four leaves.Source: §5.3 · p. 7
  • WARM-PREDICTION achieves best or statistically indistinguishable returns among action-chunked trees in every domain, matching neural-network policies with 51–79% fewer parameters than same-chunking MLPs.In LL-H, it reaches 258.4 ± 1.9 and 246.2 ± 3.3; the remaining gap to neural networks is 12–27 reward.Source: §5.3 · p. 7

5.4. Deepening / Pruning Analysis

  • Tree size changes differently by domain: IP remains at two leaves, while LK and LL initially deepen and later prune; LL-H varies more.The authors interpret the deepen-then-prune pattern as first improving representations and then compressing the tree; Figure 2 shows leaf-count trends across environments.Source: §5.4 · p. 7, Figure 2 · p. 8

5.5. Temporal Robustness Verification

  • Future action chunks support online virtual-rollout checks for unwanted termination; the LL-H example reports 67.7% accuracy and predicts crashes about seven timesteps ahead.The example’s tree path visualizes the decision flow, and Table 2 presents evaluation over 100 test runs.Source: §5.5 · pp. 7, 9, Figure 3 · p. 9, Table 2 · p. 9

6. Conclusion

The paper introduces temporal interpretability for differentiable decision trees through action-chunked policy-gradient methods and information-theoretic tree restructuring. Warm-started trees match neural-network policies in three domains and come close in LL-H, while using up to 80% fewer parameters.

6. Conclusion

  • Warm-started action-chunked trees match neural-network policies in LK, IP, and LL, come close in LL-H, and use up to 80% fewer parameters.The authors identify warm-starting from a distilled action-chunked policy as the most effective approach.Source: §6 · p. 9
  • Extending the value function to state chunks and investigating the domains where it applies are proposed as future work.The proposed extension is Vϕ(st:t+H).Source: §6 · p. 9
  • The policy gradients’ non-Markovian features likely increase training difficulty, and training action-chunked trees from scratch remains challenging in highly stochastic domains.The authors also identify ITTR parameter choices and a general tree-restructuring algorithm as further research directions.Source: §6 · p. 9

8. Impact Statement

The impact statement connects temporally interpretable white-box models to the need to understand model behavior in sequential decision-making. The appendix outline lists the paper’s supplementary derivations, algorithm details, visualizations, and ablations.

8. Impact Statement

  • Temporally interpretable white-box models are particularly relevant where understanding a model’s behavior is necessary for trust, viability, and responsibility.The impact statement places this relevance in sequential-decision-making domains.Source: §8 · p. 10

A. Outline of Appendices

  • The appendices cover policy-gradient derivations, Temporal Ensemble details, ITTR visualizations and design, ITTR analysis, additional ablations, and miscellaneous details.The outline locates these materials in sections B–H.Source: A. Outline of Appendices · p. 13

B. Derivation of: Policy Gradient with Action Chunking

The derivation motivates scaling action probabilities by returns at the action-chunk level rather than by total return or timestep-level returns. The Temporal Ensemble explanation describes combining log probabilities across chunks in the policy-gradient objective.

B. Derivation of: Policy Gradient with Action Chunking

  • Action probabilities are scaled by returns at the time-chunk level because total return and timestep-level returns can involve rewards from unrelated timesteps.The derivation relates this chunk-level return to an advantage at t + H, with the discount factor exponentiated because H rewards arrive between t and t + H.Source: B. Derivation of: Policy Gradient with… · pp. 13–14

C. Further explanation: Temporal Ensemble Policy Gradient

  • The policy samples action chunks from a multivariate Gaussian, whose independent variables allow the log probability of each action to be computed separately.These per-action log probabilities feed the Temporal Ensemble objective.Source: C. Further explanation: Temporal… · p. 14
  • Temporal Ensemble policy gradients use a linear combination of action log probabilities, each evaluated under the chunk distribution that produced it.The ensemble log probability is the stated difference from traditional policy gradients.Source: C. Further explanation: Temporal… · p. 14
  • Early-episode ensembles use the chunks available in the history and renormalize their weights; zero vectors and masks handle buffer entries from incomplete ensembles or previous episodes.The procedure addresses ensembles that lack enough chunks for a complete history.Source: C. Further explanation: Temporal… · p. 14

D. Visualization

The LL-H visualization compares a tree at the start and end of training, with each node testing a state feature against a threshold. After restructuring, the illustrated tree has unique comparisons rather than redundant decisions.

D. Visualization

  • The visualizations show crisp trees whose comparator nodes test one feature against a threshold, rather than combining all features.Leaf contents are omitted because action chunking requires parameters that would add visual clutter.Source: D. Visualization · p. 14
  • The LL-H tree’s redundant starting comparisons give way to unique decisions after pruning, avoiding multiple paths that could cover the same behaviors.The caption identifies the compared trees as the beginning and end of a run; each node compares a state feature with a threshold.Source: D. Visualization · p. 14, Figure 4 · p. 14
  • The final tree’s comparisons use angular velocity, x-position, and x-velocity, which the authors identify as the most significant LL-H success features.These are labeled x5, x0, and x2, respectively.Source: D. Visualization · p. 14

E. Information-Theoretic Tree Restructuring

Information-Theoretic Tree Restructuring (ITTR) changes tree structure during training through deepening and pruning operations. Its algorithm uses policy-complexity estimates and leaf visitations to guide those changes.

E. Information-Theoretic Tree Restructuring

  • ITTR checks whether to restructure after specified training steps, once enough policy-complexity estimates have been collected and the global-step threshold is met.The algorithm then tests a complexity-and-reward condition for deepening or a visitation condition for pruning.Source: E. Information-Theoretic Tree… · p. 15
  • When deepening, ITTR replaces a selected leaf with a new decision node and two copies of that leaf.The new node’s weights, bias, and α are sampled from their respective distributions.Source: E. Information-Theoretic Tree… · p. 15
  • For pruning, ITTR removes a leaf with fewer than k visitations when the tree has more than two leaves, then promotes its sibling subtree.The algorithm selects the least-visited leaf and removes its parent decision node.Source: E. Information-Theoretic Tree… · p. 15
  • For warm-started trees in LL-H, ITTR disables deepening and checks pruning every 1e6 steps with k = 4e6.For LL-H, the pruning condition is checked before the deepening condition; the passage also specifies reward thresholds for low rewards.Source: E. Information-Theoretic Tree… · p. 15

F. Further explanation: Number of parameters after deepening / pruning

The section examines how deepening and pruning affect DDT parameter counts, distinguishing tree logic-flow parameters from leaf parameters used for action chunking.

F. Further explanation: Number of parameters after deepening / pruning

G. Ablations

The ablations vary the EMA coefficient to examine performance trends, alongside hyperparameter tables for shared settings and environment-specific ITTR settings.

G. Ablations

  • The five EMA sweeps use η values 0.5, 0.625, 0.75, 0.875, and 1.0.A higher coefficient gives ensemble actions more equal weight, promoting smoother actions but potentially worsening performance.Source: G. Ablations · p. 16
  • In LL-H, return falls from 231.2 ± 1.9 at η = 0.5 to 143.6 ± 19.6 at η = 1.0; performance decreases with higher EMA coefficients in all domains except IP.Figure 5 plots performance against the EMA coefficient in each simulation environment; its black curve regresses results across seeds.Source: G. Ablations · p. 16, Figure 5 · p. 16
  • The EMA coefficient trades off performance against how strongly the agent commits to its displayed plan.Lower coefficients give later planned actions less weight; the unbiased η = 1.0 variant does not perform best, which the authors attribute to enforced consistency rather than estimator bias.Source: G. Ablations · p. 16

H.1. Hyperparameters

  • The hyperparameter tables distinguish settings shared across environments from environment-specific ITTR settings.The supplied captions identify these as Tables 4 and 5, respectively.Source: Table 4 · p. 17, Table 5 · p. 17, §H.1 · p. 17
Loading 2610.10367v1…