Piecewise Linear Parametrization of Policies: Towards Interpretable Deep Reinforcement Learning

We thank the reviewer for their detailed review of our submission and their precise feedback. We address their concerns in detail below.

Relationship of HC with ReLU networks

We agree that HC is indeed extremely linked with MLP-ReLU networks (abbreviated as ReLU networks in the following), as the latter is actually the starting point from our architecture (first paragraph of Sec 3.1). However, we made some strong changes to the ReLU network architecture in order to achieve desirable interpretability properties which we identified at the end of Sec. 2 and beginning of Sec. 3. Overall, the differences between ReLU networks and HC are best exemplified by Fig. 1 (right): ReLU networks induce (1) a complex partition of the input space and (2) a different linear sub-function in each linear region (one color per linear sub-function). HC also induces (1) a complex partition (useful to solve hard tasks from scratch), but (2) only exhibits a small number of linear sub-functions (illustrated by using only 4 colors, hence 4 linear sub-functions). This is in stark contrast with ReLU networks of all but the smallest architecture.

In addition, we have actually compared ReLU networks (including with small architectures) with HC in our submission: In Appendix B, Fig. 6, we show evidence that HC reaches a better performance-interpretability trade-off than ReLU networks for a wide range of architectures. This leads HC to a high performance while using a very small number of linear sub-policies. In Appendix C, Fig. 7, we show how a ReLU network uses essentially a new sub-policy at each timestep, on the contrary to HC.

We now answer your points detailing the differences between HC and ReLU networks, following your points:

1. Explicit formulation of $a$ and $\theta$

We agree that it is possible to recover $a$ and $\theta$ for ReLU networks, as we indicated in the beginning of Sec 3.1. However, these functions do not appear explicitly in the computation of the output $f(x)$, and need to be recovered (in the case of the local linear coefficients, by taking the input Jacobian of $f$). On the contrary, we make $a$ and $\theta$ explicit within the computation made by the neural net, by composing a Gumbel network $a$ and the local linear coefficients $\theta$. The explicit modeling of $a$ within the forward pass of the actor lets us set the exact number of sub-policies as a hyperparameter, unlike ReLU networks (we expand on this last statement in the answer to your second point).

Following your comment, we modified the manuscript (Sec 3.1) to clarify that $a$ and $\theta$ now appear explicitly in the computation (in purple).

Thank you for your reply and all the edits to the paper.

The related work section has improved, which is nice to see. The comparison with ReLU networks in the Appendix B (Figure 6) is also in the right direction.

The issue I still have with the paper is the lack of baselines. For example, Qiu and Zhu's work is very related and it could be used even with no skills in the leaves. Qiu and Zhu used these skills in their experiments, but they are not needed to run their method. I am sorry for suggesting yet another baseline after the first round of reviews, but it just occurred to me that network distillation could also be a good baseline. You start with a large network with many unique regions and distill it into a much smaller network with fewer regions. A comparison with these methods would be informative.

We thank the reviewer for their feedback. We are glad they enjoyed reading the paper. HC indeed does not escape the performance-interpretability trade-off, and increasing the number of sub-policies tends to improve performance but reduce interpretability (Fig. 4). We think one of the interesting aspects of our work is that the stakeholder can decide this trade-off for themselves, by fixing how many sub-policies they are willing to use.

Questions

• In Fig. 3, you are absolutely right, we are referring to the computation of the logits for the dimension 0 of the action space. We have updated the legend accordingly (in red), thanks.
• We chose SAC because it was a strong continuous control, actor-critic baseline, with several good open-source implementations easy to bootstrap from. SAC is however not adapted to the maze experiments, mainly due to the specificities of the environments such as sparse rewards. We therefore looked for a base algorithm that would (1) be adapted to maze experiments, (2) be open-source and (3) have no competing architectural assumptions on the policy. RIS fit all these categories, which is why we chose it.
• The sub-policies are indeed trainable. We learned each sub-policy (jointly with how to choose them) in each of our experiments. We actually tried to diversify the initialization of each sub-policy earlier in the project, but we did not find any benefits in terms of performance at that time. We therefore find the reviewer’s proposition of diversifying the sub-policies themselves (in addition to how to choose them) could be very interesting to revisit in the future, as it might help each sub-policy to specialize even more.

We thank the reviewer for their review of our work and their suggestions. We understand that the reviewer has concerns about (1) the performance of our method and (2) the interpretability that is gained with HC. We detail our answer to these concerns below.

Performance of HC

We agree with the reviewer that performance is an important factor in user trust. One of HC’s main advantages is that it lets the stakeholder tune the performance-interpretability trade-off through the number of sub-policies. In certain safety-critical environments or when the stakeholder is not a domain expert, the stakeholder might be willing to sacrifice some of the performance in order to gain insights about how the policy interacts with its environment.

On the other hand, a domain expert might be comfortable with more complexity, unlocking higher performance. This is illustrated in Fig. 4: HC64 never catastrophically fails, despite the high difficulty of the maze environments and its relative simplicity compared to the MLP baselines. HC8 struggles to solve the harder tasks but leads to simpler visualizations (e.g., Fig. 5). We think this is especially encouraging given the high difficulty of the mazes. We update our manuscript in Sec 3.2 to express this point more clearly, in blue.

Finally, some applications demand first and foremost the highest level of performance. We think our approach is outside of this scope, since a performance cost has to be paid for the gains in interpretability.

Applications of the method

We thank the reviewer for their remark. In this work, our goal was to study the properties that could improve the interpretability of piecewise-linear policies, and to propose a parametrization that let us achieve as many of these properties as possible. Therefore, we studied and evaluated the HyperCombinator as a general-purpose architecture.

The two levels of interpretability of HC relate to its transparency. The linear coefficients inform of the importance of a feature in the computation made by a subpolicy, and the sequence of chosen behaviors gives insights about the structure of the inner decision process of the agent. In turn, a domain expert can use this transparency to better understand the policy that they roll out (Appendix D.7 gives a low-dimensional example of such an analysis).

The sub-policy sequence visualization could also be used to detect issues at run-time, especially if no video of the agent is available. In Fig. 15 (Appendix D.6), we studied the robustness of HC to perturbations, which we could model as an adverse event. The perturbations that the agent is subjected to are translated into an unusual sub-policy sequence between timesteps 100 and 200. Hence, we could imagine a system that detects a change in the subpolicy sequence, which would in turn raise an alarm that an adverse event has been detected.

Overall, HC opens the gate to new analyses thanks to its transparency, which we think is a significant improvement from the status quo.

Other remarks

• We reformulated the “remarkable” mention in p.6 following your suggestion.
• Contrastive explanations: we thank the reviewer for their suggestion. We agree that the architecture of HC opens up new possibilities, such as analyzing why a given sub-policy was chosen rather than another.
• P2: we clarified the meaning of the property in p.3 following your suggestion.
• Title: we believe that both “parameterization” and “parametrization” are valid ways to spell the word.
• Finger spin consists in the very quick repetition of the same precise movement in order to maximize the reward. Policies with more degrees of freedom have an advantage here as they can reach a better precision thanks to their higher expressivity. We find this not too surprising, as complex policies end up performing better than HC given enough samples (Fig. 12, Appendix D.4). We added a remark about this phenomenon in Appendix D.4.

Questions

It is indeed entirely possible to use interpretable features while conditioning the Gumbel network on pixel data. For instance, in DM control, we could define the observations as multimodal: one pixel representation of the environment, and the associated (interpretable) proprioceptive features. The Gumbel network could be fed the pixel rendering of the environment, which would predict a sub-policy to apply. This sub-policy would be linear with respect to the interpretable features. We did not propose results with pixel data for now as we felt this was beyond the scope of the paper, which introduces the architecture with respect to interpretable features only. However, we do agree it would be an interesting follow-up, and intend on pursuing a pixel-based version of HC in future work.

Hi, thanks a lot for your reply, I have read everything and thought it over for a while.

I think I can agree with most of what you say, the only part I feel we part ways is about applications of the technique. I would like to offer potentially useful ideas for you.

You say "In turn, a domain expert can use this transparency to better understand the policy that they roll out (Appendix D.7 gives a low-dimensional example of such an analysis)."

This is a good start, but why exactly do you want to understand the policy? That's the key question. If it is for debugging purposes, then you need a user study showing it does this. If it is to calibrate appropriate reliance, then again you need users in the loop to show it actually does this, otherwise it's all conjecture.

The sub-policy sequence visualization could also be used to detect issues at run-time, especially if no video of the agent is available

Again, I think you need to show how seeing these issues is helpful. I know it sounds obvious that it should be useful, but too often XAI techniques are published with the author's claiming the method is "interpretable", but all the while failing to demonstrate that the method is understandable and useful to intended practitioners of the system.

I like the paper, which is why I accepted it, but I cannot raise my score any higher with this fundamental issue present. Thanks again and good luck!

Dear reviewer,

Thank you again for the time you already dedicated to reviewing our work.

As the end of the discussion period approaches, we are reaching out again to check if our previous reply answered your concerns. If not, please let us know, as we would be happy to have the opportunity to address your remaining concerns.

The authors

We thank the reviewer for reviewing our submission and for their feedback. We understand that the reviewer has concerns about (1) the novelty of our work and (2) the resulting interpretability and the trade-off with performance. We address these concerns below.

Novelty of our approach

There are several components that distinguish our paper from previous works. Notably, several aspects of the Gumbel network, the module which decides which sub-policy to apply to a given input, are key to let HC learn in the different experiments:

• In particular, the MLP in the Gumbel network leads to an intricate partition of the input space (unlike regular tree methods), which enables a good performance of HC in the DeepMind Control benchmark as well as in select mazes.
• Moreover, we specifically use the straight-through estimator (STE) to guarantee the selection of a single sub-policy at each timestep (instead of a mixture of several sub-policies), during training and during evaluation. This is an important property for interpretability (enabling analyses during training) that was missing from previous works.
• Finally, while previous works (such as Akrour 2018) focused mostly on performance, we study in detail how such a parameterization influenced the piecewise-linear functions we could express, within the scope of interpretability. The different visualizations give insights about how the agent is interacting with its environment, which are not possible with regular neural networks.

To summarize, our paper performs an in-depth study of the limits of MLP with ReLU activations as an interpretable parametrization of piecewise-linear policies. We propose an architecture that parametrizes a piecewise-linear policy along the constraints. We think this fills a gap that was present in the literature.

Characterization of the interpretability gained with HC

The HC architecture provides two levels of interpretability:

1. Low-level interpretability through the value of the linear coefficients of each sub-policy,
2. High-level interpretability by examining the sequence of choices of sub-policies over a trajectory.

The heatmap in Fig. 3 is an example of low-level interpretability. It shows the coefficients of each sub-policy (one column per sub-policy) for one dimension of the action space. For instance, subpolicy #0 assigns a positive coefficient to feature f0 as the corresponding cell is red. Therefore, we know that inputs with greater positive values for f0 will lead to greater predicted actions, all else being equal. Moreover, the coefficient quantifies the exact importance of feature f0 in the computation of the action.

In addition, we perform a more extensive analysis of the meaning of the coefficients in Appendix D.7, where we analyze the learnt policies on the cartpole environment. The lower dimensionality of the environment helps getting a better grasp of the computations and therefore of each sub-policy. For instance, the coefficient values let us understand the usage of each sub-policy and how the agent specialized some to get the CartPole up, and some to maintain it straight. We think this is a considerable improvement over the regular non-interpretable neural network, and makes the decision making process significantly more transparent.

The high-level interpretability brought by our agent is complementary to this low-level interpretability, as we get a better understanding of how the agent processes the information to decide which subpolicy to follow. For instance, Fig. 3 (left) illustrates how the agent learnt to chain efficiently the linear sub-policies in a repeated manner to solve Half-Cheetah with a strong performance, and the hierarchical structure in Fig. 5 reflects the structure of the maze.

Question about interpretability-performance trade-off

While the performance of HC64 decreased with the increasing difficulty of the mazes, it never collapsed (Fig. 4). We think this is especially encouraging as these tasks are extremely challenging, and require dedicated algorithms to solve them (Chane-Sane 2021). This leads us to think that HC can handle complex environments without fully losing performance simply by increasing the number of sub-policies available. For instance, in Fig. 4, we had to trade-off some interpretability by going from 8 sub-policies to 64.

Therefore, the answer to this question ultimately lies with the stakeholder that needs to understand the sub-policy. If that person is a domain expert, it is possible for them to increase the number of sub-policies and therefore increase the performance. A person less familiar with the environment might want to reduce the number of sub-policies that HC can use, sacrificing some of the performance to understand how the agent interacts with the environment. We updated the manuscript (Sec 3.2) to reflect this point (in olive).

Dear reviewer,

Thank you again for the time you already dedicated to reviewing our work.

As the end of the discussion period approaches, we are reaching out again to check if our previous reply answered your concerns. If not, please let us know, as we would be happy to have the opportunity to address your remaining concerns.

The authors