Mitigating Information Loss in Tree-Based Reinforcement Learning via Direct Optimization

Authors: Sascha Marton, Tim Grams, Florian Vogt, Stefan Lüdtke, Christian Bartelt, Heiner Stuckenschmidt
Primary area: reinforcement learning
Keywords: symbolic reinforcement learning, interpretable reinforcement learning, reinforcement learning, decision trees, policy gradient, proximal policy optimization

Abstract:
Reinforcement learning (RL) has seen significant success across various domains, but its adoption is often limited by the black-box nature of neural network policies, making them difficult to interpret. In contrast, symbolic policies allow representing decision-making strategies in a compact and interpretable way. However, learning symbolic policies directly within on-policy methods remains challenging.
In this paper, we introduce SYMPOL, a novel method for SYMbolic tree-based on-POLicy RL. SYMPOL employs a tree-based model integrated with a policy gradient method, enabling the agent to learn and adapt its actions while maintaining a high level of interpretability.
We evaluate SYMPOL on a set of benchmark RL tasks, demonstrating its superiority over alternative tree-based RL approaches in terms of performance and interpretability. Unlike existing methods, it enables gradient-based, end-to-end learning of interpretable, axis-aligned decision trees within standard on-policy RL algorithms. Therefore, SYMPOL can become the foundation for a new class of interpretable RL based on decision trees. Our implementation is available under: https://github.com/s-marton/sympol

Decision: accept (spotlight)
Reviewer scores: 5, 5, 8, 10, 8

Reviews and discussion:
<thread>
creation date:
1729986595918

thread_type:
official_review

<thread_object>
creator: reviewer_6tah

<content>
summary:
The paper presents SYMPOL, a new approach to learning interpretable policies for reinforcement learning problems by building on recent works facilitating gradient-based training of decision tree architectures. SYMPOL performs much better than prior methods for producing decision tree policies (such as discretizing policies or behavior-cloning a decision tree from a replay buffer), and includes a few tricks for improving the stability of training tree-based models.

SYMPOL is evaluated on several RL problems, largely simpler problems such as mountain car or lunar lander, as well as a MiniGrid environment. Finally, the interpretability of SYMPOL policies is displayed using a case study on policy generalization, in which the authors inspect a learned policy, identify areas for improvement or weaknesses in the policy, and then re-train to fix these weaknesses.

strengths:
* The paper builds on very recent work and adequately cites the prior work, making sure to clearly draw the line between novel contributions and techniques which already exist.
* The SYMPOL framework is clearly explained, and the tricks used to improve training stability are highlighted (as opposed to being buried in the appendix). I feel like I could re-implement this from the paper alone.
* The results are compelling. SYMPOL performs quite well (final performance numbers/returns are superior to baselines) while also generalizing well to new versions of the problems (as shown by the Cohen's D experiments)
* The case study on interpretability is compelling, showing the power of an interpretable RL framework such as SYMPOL. 

Overall, the paper is very well written, the framework is clearly described, makes intuitive sense, and does not require significant hand-tuning, and the results and discussion sections are strong.

weaknesses:
* SYMPOL is mentioned to have 50+ nodes, which could become quite unwieldy or difficult to interpret, but we do not see such examples in the paper. It would be illustrative to see a standard SYMPOL policy in addition to the cleaned up versions in Figures 7/9.
* While the framework is sound and the results are strong, the overall novelty of the framework is limited. Apart from introducing a dynamically growing replay buffer, the technical contributions are largely a combination of existing techniques.
* As the paper relies heavily on the contribution of GradTree, it would be helpful to have an overview of that work here, so that the finished paper reads as a more self-contained work.
* The comparisons across all works are not entirely 1-1, as the number of nodes in each tree is not consistent (however, the authors make note of this in the text).

questions:
* There is no ablation on the automated pruning, but automated pruning is mentioned a few times in the work. What are the effects of this on the resulting policy, and at what point are the node automatically pruned? If they are not pruned during training, do they complicate gradient flow?
* How is misaligned behavior detected, as in the case study shown in Section 6? It seems that, because the network's parameters are 1D feature vectors, you would need to reverse engineer the input states for each node, then plot out the resulting tree? Is there another way, or is this manually done?
* Would SYMPOL extend naturally to continuous action spaces or unstructured input (e.g., images)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730040585098

thread_type:
official_review

<thread_object>
creator: reviewer_7ps3

<content>
summary:
Authors propose SYMPOL. SYMPOL alleviates previous challenges of using RL to learn decision tree policies by using a differentiable tree representation. They also increase training stability by collecting longer rollouts at the end of training and using large batch sizes. The paper is well written. However, with the current experiments, it is unclear to me if the good experimental results are due to the use of GradTree or to the stabilizing of the training or both. But in any case, the contribution is not clear.

strengths:
The related work and preliminary sections are excellent. Misgeneralization experiment and information loss experiments are insightful. Topic of interpretability is crucial.

weaknesses:
There are two main weaknesses to this paper. 

1) Contributions. Is the contribution plugging GradTree into PPO ? 

I do find the ablation study insightful on why the training process needs strong stabilyzing and digging this further could be a nice contribution. Similarly, it seems to me that Silva et al 2020 is very similar to GradTree in the sense that GradTree is a differentiable formulation of DTs and Silva et al is too. I would have liked a more in depth comparison of the two! 

2) Experiments (just the section 5.2). Baselines seem underperforming. 

From looking at Silva et al 2020 and VIPER 2018, I feel like SA-DT (d=5) and D-SDT should be able to solve CartPole contrary to what you show in your table 1. Also why isn't SA-DT called VIPER? 

If I were you I would rewrite the section 5 centered around two questions: why use direct tree optimization rather than post hoc distillation of a neural net like VIPER? After showing that direct optim might be better than post hoc fitting of trees, you should ask which of the two sota direct optim is better and why: Silva et al or SYMPOL?

questions:
Could you remove fig 2 that does not bring anything to the paper and replace it with appendix results (such as figure 10)? Even better, could you summarize the section 4.1 or GradTree in a schematic and replace the current figure 2 with it please? If you decide to put figure 10 in the main paper instead of the schematic at figure 2, please smooth your curves.

Could you please remove figure 3 and replace it with somehting more useful? I think one does not need a plot to understand your rollout schedule.

Could you write a dedicated parargaph that explains the difference between SYMPOL and Silva et. al. 2020 ? I think your work is actually very similar to their.?

Bonus, add CUSTARD (Topin et al 2021) as a baseline in your table 1.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730201795522

thread_type:
official_review

<thread_object>
creator: reviewer_6l9x

<content>
summary:
This work presents SYMPOL, a method for online RL using decision trees as policies. 
The authors include the GradTree into PPO, thus backpropagating from a neural critic to an arithmetic decision tree-based actor.
They evaluate their method on different continuous and discrete environments, and demonstrate that SYMPOL allow to detect and correct misalignments problem that SIMPOL allows to find.

strengths:
**The method is well motivated.**
GBRL, the only interpretable tree-based method that directly learns tree instead of distilling them after training, leads to huge and many trees. 

**The potential impact of online learning of trees is huge for interpretable RL.**
The online learning of transparent policies could allow for detection of misalignment during training, and potentially correction of the misalignment during training. Distillation methods do not have this advantage.

**The structure of the paper makes it easy to follow**. Most of the paper is clear, in terms of structure, particularly the experimental evaluation. The bold points, together with the clear figures and tables, allow the reader to quickly grasp most of the experimental evaluation, even if it took me a bit of time to understand that SA-DT is an implementation of Dagger/VIPER.

weaknesses:
**The difference between SYMPOL and other baselines is not clear.**
The preliminaries section would benefit for a bigger explanation on GradTree, as its integration to PPO is the main contribution and what distinguish this work from many existing interepretable baselines such as NUDGE, INSIGHT, GBRL and SCoBots ([1]).
It also took me quite some time to understand that SA-DT is equivalent to running Dagger or VIPER with their MLP.

**Evaluation seems weird.** It feels like the baselines are underperforming, compared to what is reported in the literature. For example, the D-SDT is underperforming compared to what VIPER reports (with a few node). I don't know why, as it is okay if their baselines would perform worse than distilled ones. Directly learning trees in an online fashion, could allow for e.g. LLM guidance to provide the trees with correct inductive biases, that would avoid e.g. misalignment problems during learning.

**Missing litterature on misgeneralisation (and interpretable RL).**
On misgeneralisation, NNs have been shown to systematically misgeneralize on different Atari environment [2], even when the environment is simplified [3], and DT-based policies have been shown to allow to mitigate it [1].


[1] Delfosse et al. "Interpretable concept bottlenecks to align reinforcement learning agents." Advances in Neural Information Processing Systems 37 (2024). 

[2] Farebrother et al. Generalization and regularization in DQN. arXiv (2018).

[3] Delfosse et al. "HackAtari: Atari Learning Environments for Robust and Continual Reinforcement Learning." arXiv (2024).

questions:
* Why is it that the extracted trees are performing so low in comparison to what is reported by the VIPER's authors ?
* Is there any reason why you are not using the set of Object Centric Atari environments in your evaluation ? It would allow to compare with other deep and interpretable baselines such as NUDGE, INSIGHT, SCoBots, ...

I'm willing to revise my scores if experiments that compare SYMPOL to these interpretable methods on some Atari games are reported. It would allow the reader to compare these methods. Even if SYMPOL slightly underperforms, for interpretability, game scores are not the most important focus.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731119965685

thread_type:
official_review

<thread_object>
creator: reviewer_sgid

<content>
summary:
This paper proposes a differentiable tree-structured model that can be optimized via RL algorithms. The proposed method includes a tree-structure actor and a neural network critic. Experiments are conducted on several RL benchmarks include discrete and continuous action spaces, compared with previous tree-based methods.

strengths:
1. The proposed method can be applied into both discrete and continuous action spaces.

2. The proposed method improves interpretability of RL policies.

weaknesses:
The reviewer is not convinced with the claim of "first method that trains a DT via stochastic policy gradient in an end-to-end manner". Making the learning of DTs differentiable has been investigated for a long time, such as [1]. It is safe to make this claim through a thorough literature review.

There are some other works focusing on learning a differentiable symbolic policy which is based on first-order logic, such as [2,3], which is not compared and discussed in this paper.

Looks like that defining the syntax and semantics of DTs relies on human knowledge, maybe the reviewer is wrong, which can be clarified during rebuttal.

[1] Treeqn and atreec: differentiable tree-structured models for deep reinforcement learning.

[2] Neural logic reinforcement learning. 2019.

[3] GALOIS: boosting deep reinforcement learning via generalizable logic synthesis. 2022.

questions:
The reviewer is confused about the information loss experiments, described in Fig 1 and Table 3. Is the test reward on the same environment as the trained one? If so, why is there a performance drop of (Silva et al. (2020))? Did you mean the trained model has a stochastic probability while the test model is greedy that is, selecting the action with the maximum Q?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731226739151

thread_type:
official_review

<thread_object>
creator: reviewer_dz3q

<content>
summary:
This work employs a tree-based model integrated with on-policy RL, like PPO, to maintain better interpretability with less information loss.
On a set of benchmark RL tasks, it demonstrates its superiority over alternative tree-based RL approaches in terms of performance and interoperability.

strengths:
1. The motivation is good, aiming to improve the interpretability of RL, which is a key factor for deployment due to its adoption of DNNs. 
2. The authors try to improve the stability of DTs as a policy via more implemented skills.

weaknesses:
1. This work evaluates its experiments on a series of toy tasks that could be easily solved. Can it achieve competitive performance in complex tasks, like Mujoco and Atari games?
2. This work is also hard to follow, the contribution and the method should be highlighted and stated.
3. Please abatement the contribution of this work. This work is not the first to use DT to represent RL policy, i.e., the work "MIXRTs: Toward interpretable multi-agent reinforcement learning via mixing recurrent soft decision trees", arXiv:2305.10091, 2022. They utilized SDT in multi-agent RL complex tasks two years ago.
4. There are some typos in References.

questions:
1. line 414- 415: "In RL, especially for smaller policies, the runtime is mainly determined by the actor-environment interaction..." - This sentence is hard to follow. Can the authors unpack this sentence?
2. Did this work simply incorporate GradTree into the PPO framework?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732115581970

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Summary

comment:
We greatly appreciate the reviewers' constructive feedback and acknowledgment of our method's strengths and novelty. In particular, the reviewers `[DZ3Q, sgid, 6L9X, 7Ps3, 6tAh]` highlighted the motivation and relevance of our approach, noting its *“huge potential impact for interpretable RL through online learning of trees.”* `[6L9X]`. Furthermore, the reviewers pointed out our method's advantage over alternative approaches that do not directly optimize the policy `[6L9X, 6tAh]`, as well as the clarity of our presentation `[6L9X, 7Ps3, 6tAh]` and the value of the insightful case study `[7Ps3, 6tAh]`.

We summarize below the main concerns raised by the reviewers and explain how we have addressed each one. We have provided detailed responses to all questions directly to each reviewer. Additionally, we have uploaded a revised version of the manuscript, incorporating adjustments based on the reviewers' feedback. We have also included a version of the revised manuscript highlighting all changes, provided in the supplementary material for transparency.

* **Distinguishing SYMPOL from Differentiable Decision Tree Methods** `[DZ3Q, sgid, 7Ps3]`: We revised the related work section to more clearly distinguish SYMPOL from methods that employ differentiable decision trees in a separate paragraph, addressing the reviewers’ feedback regarding this point (L138-L143; L247-L252). In related work, differentiable decision trees, as for example used by Silva et al. (2020), are soft decision trees (SDTs), achieving differentiability by relaxing discrete decisions for feature and path selection at internal nodes. This approach is fundamentally different from SYMPOL, which does **not** use differentiable decision trees. Instead, SYMPOL leverages GradTree to optimize standard, non-differentiable decision trees through gradient descent. As a result, SYMPOL retains univariate, hard decisions, which are more interpretable than the oblique splits involving multiple features in SDTs and differentiable decision trees. Furthermore, SYMPOL does not require any post-processing to obtain interpretable, standard decision trees, which would typically result in a substantial information loss, as demonstrated in Section 5.2 of our paper.

* **Comparison with Related Work** `[6L9X, 7Ps3]`: In the revised version, we clarified the relationship between SA-DT and VIPER  (L343-L346; Section A.3), which was not explicitly outlined in the initial submission. In general, SA-DT can be considered as a version of DAGGER. Therefore, SA-DT and VIPER are conceptually similar, but VIPER (Q-DAGGER) includes advanced sampling and weighting strategies for training samples — techniques not incorporated in SA-DT as used, for example, by Silva et al. (2020). To address this, we implemented VIPER in our evaluation, enabling a direct comparison. The evaluation results remain consistent with those in the initial submission. SYMPOL achieves superior performance in most cases. This is also in-line with the results reported by `[4]`, which demonstrate that the sampling in VIPER does not result in superior performance compared to DAGGER/SA-DT when learning small, interpretable decision trees which is the focus of our paper.

* **Recentering of contributions** `[7Ps3]`: We clarified our contributions which are now centered around the problem we are solving: Mitigating the information loss inherent in existing approaches through a direct optimization of an interpretable decision tree policy with SYMPOL.

* **SYMPOL with (affine) linear combinations of features** `[7Ps3]`: We extended SYMPOL to include (affine) linear features and observed a significant performance increase in environments such as LunarLander, where (affine) linear feature combinations are relevant (reward improvement from -57 to 105). These results suggest that SYMPOL with (affine) linear features is capable of addressing more complex environments. We will add complete experiments with (affine) linear features to the appendix as soon as our experiments including HPO have finished. In addition, we are currently evaluating the performance of SYMPOL with affine linear feature on (Object-Centric) Atari environments, where preliminary results appear promising.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper presents SYMPOL, integrating tree-based models with policy gradient algorithms for interpretable reinforcement learning policies. The reviewers praised the paper's clear writing and thorough empirical evaluation. Initial concerns about technical novelty and similarity to previous work were addressed through detailed clarifications of SYMPOL's unique integration of GradTree and stability-enhancing modifications. The authors made significant improvements during review, particularly in extending SYMPOL to handle affine linear features and demonstrating its unique approach to maintaining hard decision boundaries versus soft decision trees. 
Moving forward, I suggest the authors consider expanding evaluations to more complex environments and providing clearer differentiation from recent work on differentiable decision trees. Given the paper's valuable contribution to interpretable RL and the authors' strong engagement with reviewer feedback through substantive improvements, I recommend acceptance.