Policy Distillation with Selective Input Gradient Regularization for Efficient Interpretability

Reinforcement learning (RL) systems have achieved impressive performance in a wide range of simulated domains such as games (Mnih et al., 2015; Silver et al., 2016; Vinyals et al., 2019) and robotics (Lillicrap et al., 2015; Fujimoto et al., 2018; Haarnoja et al., 2018). However, the interpretability of an agent’s decision making and robustness to attacks need to be addressed when applying RL to real-world problems. For instance, in a self-driving scenario, real-time interpretability could explain how an RL agent produces a decision in response to its observed states and enable a safer deployment under real-world conditions and adversarial attacks (Ferdowsi et al., 2018).

Saliency maps in deep learning is a technique used to interpret input features that are believed to be important for the neural network output (Simonyan et al., 2013; Selvaraju et al., 2017; Fong & Vedaldi, 2017; Smilkov et al., 2017; Sundararajan et al., 2017; Zhang et al., 2018). As the issue of interpretability in RL gets more attention, a number of methods have been proposed to generate saliency maps to explain the decision making of RL agents. Existing saliency map methods in RL either use gradients to estimate the influence of input features on the output (Wang et al., 2016) (gradient-based methods) or compute the saliency of an input feature by perturbing it and observing the change in output (Greydanus et al., 2018; Iyer et al., 2018; Puri et al., 2020) (perturbation-based methods). Gradient-based methods can compute saliency maps efficiently with backpropagation. However, the quality of gradient-based saliency maps is generally poor (Rosynski et al., 2020). Perturbation-based methods are effective in highlighting the important features of the input, but at a significant computationally cost, which can make them ineffective when deployed on systems with real-time constraints. As a result, existing RL agents cannot provide high interpretability in a computation-efficient manner.

Different from previous work proposing new saliency calculation methods, we focus on improving the natural interpretability of RL policies. Given a RL policy, we propose an approach of Distillation with selective Input Gradient Regularization (DIGR) that uses policy distillation and input gradient regularization to retrain a new policy. In our approach, input gradient regularization selectively regularizes gradient-based saliency maps of the policy to imitate its interpretable perturbation-based saliency maps. This allows the new RL policy to generate high-quality saliency maps with gradient-based methods and thus achieve both high interpretability and computational efficiency. At the same time, to ensure that input gradient regularization does not cause task performance degradation, we use policy distillation (Czarnecki et al., 2019) to constrain the output of the new RL policy to remain close to the original RL policy.

We evaluate our method in three different tasks, which include an object fetching task from MiniGrid (Chevalier-Boisvert et al., 2018), Breakout from Atari games and CARLA Autonomoud Driving (Dosovitskiy et al., 2017). The results show that RL policies trained with our approach are able to achieve efficient interpretability while maintaining good task performance. Selective input gradient regularization also improves the robustness of RL policies to adversarial attacks. These two desired properties allow the RL policy to better adapt to real-world scenarios.

To summarize, we demonstrate a novel approach to improve the efficient interpretability and robustness to attacks of RL policies based on the utilization of saliency maps. Our approach increases the applicability of RL to real-world problems.

Our approach to achieve both computational efficiency and high interpretability in RL is to produce a policy whose gradient-based saliency maps are comparable to those of perturbation-based methods. To achieve this, given a trained RL policy, we set its perturbation-based saliency maps as supervisory signals and update the weights of the policy so that its gradient-based saliency maps match the perturbation-based saliency maps. Since the computations involved in gradient-based saliency maps are differentiable, we can use stochastic gradient descent to conduct the training. The idea of optimizing gradient-based saliency maps has a close connection with input gradient regularization which imposes constraints on how input gradients behaves. For example, Ross & Doshi-Velez (2018) penalizes input gradients based on expert annotation to prevent the network from “attending” to certain parts of the input in an image classification task. Inspired by this, the training of the gradient-based saliency map in our approach is conducted by selectively penalizing the gradients of input features that have low perturbation-based saliency.

One challenge of selective input gradient regularization is that optimizing gradient-based saliency maps may also affect the policy output and thus degrade the task performance. To avoid this, we conduct policy distillation to ensure that the new policy maintains the same task performance. We give a more formal introduction of our method below.

Given a RL policy $\pi$ and input $s$, we define the function $g$ as the method used in generating gradient-based saliency map $M_{g}$ and function $f$ as the method used in generating perturbation-based saliency map $M_{p}$. Both $M_{g}$ and $M_{p}$ have the same size as input $s$. Each element in the saliency map, ${M_{g}}_{i}$ and ${M_{p}}_{i}$, are computed as

where $g(s,i,\pi)$ and $f(s,i,\pi)$ compute the gradient-based and perturbation-based saliency values of input feature $s_{i}$ given policy $\pi$. These saliency values are then normalized between 0 and 1 to form saliency maps which contain $N$ elements in each map. In this work, perturbation function $m$ induces a Gaussian blur on the input with the input feature of interest $s_{i}$ as the center (Greydanus et al., 2018). It’s worth mentioning that, besides perturbation-based saliency maps, DIGR could be easily extended to utilize other saliency data (e.g. saliency maps from expert annotation) as supervisory signals. In this work, we focus on using perturbation-based saliency maps for input gradient regularization as they show high interpretability and can be computed as long as we have access to the policy and states.

After introducing the process of generating two types of saliency maps given a RL policy and state input, we introduce how they are used in DIGR. Given a trained RL policy $\pi_{t}$, DIGR aims to produce a new policy $\pi_{\theta}$ with parameters $\theta$ that can generate interpretable saliency maps using gradient-based method. Given a state input $s$, the saliency map could differ based on the generation method (gradient-based vs perturbation-based) and the policy ($\pi_{t}$ vs $\pi_{\theta}$) used to generate them. For clarity, we define these 4 types of saliency maps as ${M_{g}}^{t}$, ${M_{g}}^{\theta}$, ${M_{p}}^{t}$, ${M_{p}}^{\theta}$ where $M_{g}$ and $M_{p}$ represent gradient-based and perturbation-based saliency maps, respectively. The superscripts $t$ and $\theta$ represent whether the saliency map is generated by the given trained policy $\pi_{t}$ or the new policy $\pi_{\theta}$. Then the loss function for input gradient regularization is

where $d_{\pi_{\theta}}$ is the state distribution following policy $\pi_{\theta}$ and $N$ is the number of input features in the saliency map. ${M_{p}}^{t}$ and ${M_{g}}^{\theta}$ have the same size and are both indexed by $i$. Threshold $\lambda$ is used in the indicator function $\mathbbm{1}$ to determine whether one input gradient should be penalized. The indicator function $\mathbbm{1}$ returns 1 if $\lambda-{M_{p}}^{t}_{i}\geq 0$ and 0 otherwise. In other words, if the perturbation-based saliency for an input feature is below threshold $\lambda$, the loss penalizes its gradient-based saliency. This selective penalization allows the model to only keep high saliency on task-relevant features selected by the perturbation-based saliency maps.

The final loss function in our approach is a weighted combination of selective input gradient regularization and policy distillation. In practice, generating perturbation-based saliency maps online for input gradient regularization could be time-consuming and slow down the overall training. To address this, we build an offline perturbation saliency dataset $D$ which contains states sampled from $d_{\pi_{t}}$ and the corresponding perturbation-based saliency maps generated in advance. Because of the policy similarity brought by policy distillation, we use $D$ to approximate $d_{\pi_{\theta}}$ for input gradient regularization. As a result, the loss function for DIGR is

where $\alpha$ is a weighting parameter used to balance the loss of input gradient regularization and policy distillation. We show the complete architecture of our approach in Figure 2.

We conducted experiments on three tasks including Red-Fetch-Green in MiniGrid, Breakout in Atari games and CARLA Autonomous Driving to demonstrate the effectiveness of our approach. In Red-Fetch-Green , the red agent needs to locate and pick up the green object while avoiding picking up other distractors in a room composed of 8x8 grids. In Breakout, the paddle is controlled to move at the bottom to ricochet the ball against the bricks and eliminate them for rewards. Besides these two tasks, we designed a CARLA Autonomous Driving task in which the agent needs to control an autonomous car driving on a highway while avoiding collisions. Since CARLA simulator has its simulation clock and time that can be matched with real time, we use it to demonstrate that both high quality and computation efficiency of our approach in interpreting RL policies are important in real-world scenarios.

We propose an approach called DIGR to improve the efficient interpretability of RL by retraining a policy with selective input gradient regularization and policy distillation. Our approach allows RL policies to generate highly interpretable saliency maps with computationally efficient gradient-based methods. We further show that our approach is able to improve the robustness of RL polices to multiple adversarial attacks. Interpretable decision-making and robustness to attacks are two challenges in deploying RL to real-world systems. We believe our approach could help to build trustworthy agents and benefit the deployment of RL policies in practice.