Paper: Can Differentiable Decision Trees Learn Interpretable Reward Functions? (ICLR-2024)

Reviewer comment (Evaluation):
One of my main concerns is the lack of repetition over random seeds. In my understanding, the cartpole environment uses 3 seeds (which is not enough for RL), but the Atari environments seem to use only one seed (and treat the seed as a tunable hyperparameter?). This is a cause for concern regarding the replicability and generalization of the proposed method. In addition, it means it is basically impossible to tell if the results about using a regularizer or soft output are significant. The mean and standard deviation across random seeds for each experiment should be reported.

Author response:
We have increased the number of seeds for cartpole to 10 and will  update soon the results in the paper with means and standard deviations. We agree that more seeds for Atari is important. You are correct that we only showed results for a single seed (seed 0). We are working on generating results for more seeds for Table 3 and will add these to the camera ready draft.

Reviewer comment (Evaluation):
The inability to interpret the decision for breakout is a bit worrying. I think this is an important result that should be discussed in the main paper for clarity.

Author response:
Due to space constraints we put the Breakout results in the appendix, but will work to add more discussion to the main paper. We would like to highlight the fact that the Breakout results are interpretable. In Fig 14,which depicts our trace for unregularized Breakout Reward DDT,  we find at Node A, states that have more number of bricks missing have a higher probability(>0.5) of being routed to left, i.e. to Node B which only uses a single leaf node to give a reward of +1 while the states in Node A that have no or very few bricks missing have a lower routing probability( routing probability <0.5)  and thus are routed right , i.e. to Node C which again uses a single leaf node to give a reward of 0. We can similarly interpret the regularized reward DDT for Breakout.

Reviewer comment (Presentation):
The synthetic trace method could be explained better, as it seems to be a key part of allowing interpretability in complex environments. But it is only introduced in the experiments section. This makes the experiment section a bit hard to follow. I think that introducing this in the methodology section would make more sense.

Author response:
We agree that the synthetic trace is an important contribution to our work. We will introduce it earlier as suggested.

Reviewer comment (Reproducibility):
How would the synthetic trace method be applied to the simple environments (cartpole and mnist)?

Author response:
We have generated synthetic traces for MNIST and include them in Appendix D.5.  To generate a trace for an internal DDT node, we sort the training data in descending order based on the routing probability assigned to each state. We then sample every Tth state to produce a trace of states that range from those routed most highly to the left child to those most highly routed to the right child. Interestingly, we find that these traces can uncover interesting edge cases and anomalies. e find that B images of 1 are routed to left leaf node as their routing probability>0.5 while images of 2 are routed to right leaf node. For Node C, our trace interestingly captures the fact that 0s get routed to left leaf node while some 2's which are visually closer to 0s  get routed to left leaf node while those 2's that are closer to 3s and  actual 3s get routed to right leaf node. For Node A, we find 1s and higher number of  2s are routed to Node B while 0s,3s, and some 2s , which are in terms of pixels, closer in shape to 0s and 3s are routed to Node C.

We will also add cartpole traces to the appendix of the camera ready draft.

Reviewer comment (Experiments):
Could this method be applied to text data?

Author response:
This is a great question. While out of scope for this paper, we think there may be a possible path forward by further extending the idea of a sophisticated internal node in a reward DDT to a transformer-based or other type of internal node better suited for text inputs. Given a prompt and output, we could feed this into a small transformer node to get it to predict routing probabilities of left and right. It may be possible to use self-attention-based interpretability or local explanations based on counterfactuals [1], then we could try and gain understanding of what parts of the sentence are leading to specific reward outputs.

Another idea for future work would be to take a pretrained LLM that has already been trained via RLHF and then try to distill it into a DDT. The internal nodes could still be transformer nodes or another form of node that can process text such as 1-d convolutions on word embeddings. Thus, the main challenges are how to design an internal node architecture for processing text. These nodes can then be organized hierarchically as a DDT to perform reward learning. We think this is a very exciting area of future work!

[1] Ross et al. "Explaining NLP Models via Minimal Contrastive Editing (MiCE)." ACL-IJCNLP 2021.

Reviewer comment (Evaluation):
Why does the neural network baseline perform poorly for Cartpole? This is surprising since cartpole is a pretty easy task.

Author response:
We were also confused about these results. Before submission, we tried a wide variety of architectures, but were unable to improve performance. We have looked into this more. Upon further inspection we found the problem. We were basing our implementation on prior work [2] where the reward function output was fed through a sigmoid. However, this seems to significantly impact performance. With no sigmoid, results are much better. We will update Table 1 soon with these new results.

[2] Brown et al. "Extrapolating beyond suboptimal demonstrations via inverse reinforcement learning from observations." ICML, 2019.

Reviewer comment (Novelty):
Limited algorithmic novelty.

Author response:
Thank you for your encouraging feedback and great questions.

> Limited algorithmic novelty

We would like to clarify the following contributions:
- Derivation of fully differentiable DT training using pairwise preferences for reward learning. We are the first to demonstrate a method for learning rewards using decision trees in a fully differentiable manner. The idea of soft-routing is common across any kind of DDT, but training a reward function uses a different loss function than prior work on classification. Thus our loss function is one of our paper's contributions. We cannot simply use a standard classification loss because we do not assume access to reward labels for every state.
- Characterization and evaluation of different leaf nodes. As discussed in Section 3.2 of our paper, prior work on DDTs has mainly focused on classification tasks and we borrow the multi-class reward leaf formulation from prior work. Our results show that using classification style leaf nodes do not work as well as our proposed min-max interpolation leaf nodes which we find is best suited for reward DDT learning.
- Explanations per leaf node: pixel based attribution as well as synthetic traces. Our focus on explainable rewards via tree structures and on using both synthetic traces and pixel-based attribution to enable reward function debugging.
- First approach to scale tree-based reward learning to visual inputs.
- Characterization and evaluation of RL-time reward predictions: soft vs hard. Interestingly, we find there is a tension between wanting hard routing for interpretability and wanting soft routing to enable better shaped rewards for reinforcement learning.

We would like to emphasize that while we use and extend some of the prior tools developed for DDT training for classification tasks, our experimental results are highly novel as we are the first to study training DDTs for reward learning to enable better explainability and alignment verification. While learning rewards from preferences is common (e.g. RLHF), prior work is not interpretable. Thus, we believe our contributions will be of great interest to the community.

Reviewer comment (Novelty):
In the related work section, it would be beneficial to discuss research on expert-driven reward design techniques that incorporate structural and interpretability constraints [1, 2, 3].

Author response:
> Related work incorporating structural and interpretability constraints [1, 2, 3].
[1] Jiang, et al. Temporal-Logic-Based Reward Shaping for Continuing Reinforcement Learning Tasks. 2021.
[2] Devidze, et al. Explicable Reward Design for Reinforcement Learning Agents. 2021.
[3] Icarte, et al. Reward Machines: Exploiting Reward Function Structure in Reinforcement Learning. 2022.

Thank you for pointing out this related work. We will add these papers to our related work. The referenced papers are similar to our paper in that they also have a focus on structured and interpretable rewards. However, there are also some key differences. 

Jiang et al. [1] develop a novel theory for reward shaping for average reward RL tasks and incorporate temporal logic specifications to help with RL learning speed. The major distinction between this work and ours is that Jian et al. assume access to true reward and seek to add a shaping reward, whereas we do not assume any access to the ground truth reward, nor to a temporal logic shaping reward signal. Instead we must learn the entire reward function purely from pairwise preferences. One exciting area of future research could be to investigate how to use a user-supplied LTL formula to help guide building the DDT. Another area would be to try and distill a DDT into an LTL specification.

Devidze et al. [2] focus on finding rewards that are sparse and yet speed up policy learning. This problem is analyzed from an algorithmic teaching perspective where a teacher is assumed to have access to a ground-truth reward function and optimal policy for that reward function and must distil the reward function into something that is informative (leads to fast RL) but also sparse. By contrast, our approach focuses on learning rewards that are initially unknown and unspecified. Given a known reward that is not sparse, it may be possible to distill the reward into a DDT to obtain a sparser, but still informative reward. We think this is an interesting area of future work.

Icarte et al. [3] focus on helping with reward design, not reward learning from human feedback. They assume that an RL agent is given access to the specification of the reward function in the form of a finite state machine reward. Icarte et al. then contribute is a collection of RL methods that can exploit a reward machine’s internal structure to improve sample efficiency and learn policies in a more
sample efficient manner. This paper presents a nice formalism for thinking about how to design interpretable rewards that are well suited for RL, but, similar to the other papers, assumes the human can pre-specify this reward. We focus on the problem of using sparse human feedback to teach an AI the reward function rather than needing to specify it manually.

Reviewer comment (Experiments):
Regarding Table 1, could the performance of the baseline neural network reward be enhanced by using a richer representation such as a high-capacity deep neural network?

Author response:
> Table 1 neural network results

We were also confused about these results. Before submission, we tried a wide variety of architectures, but were unable to improve performance. We have looked into this more. Upon further inspection we found the problem. We were basing our implementation on prior work [4] where the reward function output was fed through a sigmoid. However, this seems to significantly impact performance. With no sigmoid, results are much better. We will update Table 1 soon with these new results.

[4] Brown et al. "Extrapolating beyond suboptimal demonstrations via inverse reinforcement learning from observations." ICML, 2019.