Visual Rationalizations in Deep Reinforcement Learning for Atari Games

TL;DR

This paper introduces a visualization method to explain the decision-making process of deep reinforcement learning agents in Atari games, enhancing transparency and interpretability of their actions.

Contribution

It proposes a novel visualization approach to generate visual rationalizations for black-box deep RL agents, improving understanding of their decisions.

Findings

Visual rationalizations improve interpretability of RL agents.

Method applied successfully to Atari game agents.

Enhances transparency in deep reinforcement learning models.

Abstract

Due to the capability of deep learning to perform well in high dimensional problems, deep reinforcement learning agents perform well in challenging tasks such as Atari 2600 games. However, clearly explaining why a certain action is taken by the agent can be as important as the decision itself. Deep reinforcement learning models, as other deep learning models, tend to be opaque in their decision-making process. In this work, we propose to make deep reinforcement learning more transparent by visualizing the evidence on which the agent bases its decision. In this work, we emphasize the importance of producing a justification for an observed action, which could be applied to a black-box decision agent.

Figures

Equations

$$R_t=\sum _{k=0}^{\infty }{\gamma }^k{r}_{t+k+1}$$

$${\theta }_{t+1}={\theta }_t+{\nabla }_{{\theta }_t}\log \pi (a_t|s_t;{\theta }_t)A_t.$$

$${\theta }_{t+1}={\theta }_t+{\nabla }_{{\theta }_t}\log \pi (a_t|s_t;{\theta }_t)A_t+\beta {\nabla }_{{\theta }_t}H(\pi (s_t;{\theta }_t)),$$

$${\alpha }_k^c=\frac{1}{Z}\sum _i\sum _j\frac{\partial y^c}{\partial L_{ij}^k}.$$

$${\alpha }_k^a=\frac{1}{Z}\sum _i\sum _j\frac{\partial h^a}{\partial L_{ij}^k},$$

$$L_{GradCAM}^a=ELU(\sum _{k=1}^K{\alpha }_k^a{L}^k).$$

Full text

1 Introduction

Due to strong results on challenging benchmarks over the last few years, enabled by the use of deep neural networks as function approximators  [6, 11, 17] deep reinforcement learning has become an increasingly active field of research. While neural networks allow reinforcement learning methods to scale to complex problems with large state spaces, their decision-making is opaque and they can fail in non-obvious ways, for example, if the network fails to generalize well and chooses an action based on the wrong feature. Moreover, recent work [7] has shown that these methods can lack robustness, with large differences in performance when varying hyperparameters, deep learning libraries or even the random seed. Gaining insight into the decision-making process of reinforcement learning (RL) agents can provide new intuitions about why and how they fail. Moreover, agents that can justify with visual elements why a prediction is consistent to a user are more likely to be trusted [18]. Generating such post-hoc explanations, also referred to as rationalizations, does not only increase trust, but also it is a key component for understanding and interacting with them [3]. Motivated by explainability as a means to make the black-box neural networks transparent, we propose to visualize the decision process of a reinforcement learning agent by using Grad-CAM [15].

Grad-CAM creates an activation map that shows prominent spaces of activation given an input image and class, typically in an image classification task. The activation map is calculated through a combination of the convolutional neural network weights and the gradient activations created during a forward pass of the input image and class in the neural network.

Applying this method instead to a reinforcement learning agent, it wil be used to construct action-specific activation maps that highlight the regions of the image that are the most important evidences, for the predicted action of the RL agent. We evaluate these visualizations on three Atari 2600 games using the OpenAI Gym wrapper, created precisely to tackle difficult problems in deep reinforcement learning. The range of games in the wrapper are diverse in difficulty: they have different long-term reward mechanics and a different action space per game. These difficulties are of interest when looking to explain why the agent takes a specific action given a state.

This paper is structured as follows: The next section, 2, discusses related works in both reinforcement learning and explainable AI. Section 3 presents the visual rationalization model and explains how it is adapted to reinforcement learning tasks. Following after that is section 4 which provides the setup required for experiments. This section also provides the results for the rationalization model, including where the model fails. The last section, section 5, provides an conclusion to the experiments.

2 Related Work

In this section, we discuss previous works relevant to reinforcement learning and explainable artificial intelligence.

2.1 Deep Reinforcement Learning

In general, there are two main methods in deep reinforcement learning. The first method uses a neural network to approximate the value function that estimates the value of state, action pairs as to infer a policy. One such value function estimation model is called the Deep-Q Network (DQN), which has had garnered much attention to the field of deep reinforcement learning due to impressive results on challenging benchmarks such as Atari 2600 games [11]. Since the release of this model, a range of modifications have been proposed that have improved this model such as the Deep Recurrent-Q model, the Double DQN model and the Rainbow DQN model [5, 9, 19].

The second method used in deep reinforcement learning approximates the policy directly, by parameterizing the policy and using the gradient of these parameters to calculate an optimal policy. This method is called the policy gradient method, and a much cited example of such a method is known as the REINFORCE line of algorithms [20]. More recent examples of policy gradient methods include Trust Region Policy Optimization and Proximal Policy Optimization [13, 14].

A hybrid that combines value function methods and policy gradient methods is known as the actor-critic method. In this method, the actor is trying to infer a policy using a state, action pair and the critic is assigning a value to the current state of the actor. In this paper we use the Asynchronous Advantage Actor-Critic (A3C) model which has been used to achieve human level performance on a wide range of Atari 2600 games [10].

2.2 Explainable AI

Generating visual or textual explanations of deep network predictions is a research direction that has recently gained much interest [1, 8, 12, 23]. Following the convention described by Park et al. in [12], we focus on post-hoc explanations, namely rationalizations where a deep network is trained to explain a black box decision maker which is useful in increasing trust for the end user.

Textual rationalizations are explored in Hendricks et al. [8] which proposes a loss function based on sampling and reinforcement learning that learns to generate sentences that realize a global sentence property, such as class specificity. Andreas et al. [1] composes collections of jointly-trained neural modules into deep networks for question answering by decomposing questions into their linguistic substructures, and using these structures to dynamically instantiate modular networks with reusable components.

As for visual rationalizations, Zintgraf et al. [23] propose to apply prediction difference analysis to a specific input. [12] utilizes a visual attention module that justifies the predictions of deep networks for visual question answering and activity recognition. In [4] the authors propose to use a perturbation method that selectively blurs regions to calculate the impact on an RL agent’s policy. Although this method demonstrates important regions for the agent’s decision making, the method used in this paper highlights important regions without the need for such a perturbation method.

Grad-CAM [15] uses the gradients of any target concept, i.e. predicted action, flowing into the final convolutional layer to produce a coarse localization map highlighting the important regions in the image for predicting the concept. It has been demonstrated on image classification and captioning. In this work, we adapt it to two reinforcement learning tasks to visually rationalize the predicted action.

3 Visual Rationalization Model

In reinforcement learning, an agent interacting with an environment over a series of discrete time steps observes a state111Here we assume problems where partial observability can be addressed by representing a state as a small number of past observations $s_{t}\in\mathcal{S}$, takes an action $a_{t}\in\mathcal{A}$ and receives a reward $r_{t}$ and observes the next state $s_{t+1}\in\mathcal{S}$. The agent is tasked with finding a policy $\pi:\mathcal{S}\times\mathcal{A}\rightarrow[0,1]$, a function mapping states and actions to probabilities whose goal is to maximize the discounted sum of rewards:

[TABLE]

which is the return with discount factor $\gamma\in[0,1]$.

3.1 Asynchronous Advantage Actor Critic Learning

Gradient based actor-critic methods split the agent in two components: an actor that interacts with the environment using a policy $\pi(a|s;\theta)$, and a critic that assigns values to these actions using the value function $V(s;\theta)$. Both the policy and the value function are directly parameterized by $\theta$. Updating the policy and value function is done through gradient descent

[TABLE]

With $A_{t}=R_{t}-V(s_{t};\theta_{t})$, an estimation of the advantage function [21]. In [10], the policy gradient actor-critic uses a series of asynchronous actors that all send policy-gradient updates to a single model that keeps track of the parameters $\theta$. In our implementation the actor output is a softmax vector of size $|\mathcal{A}|$, the total number of actions the agent can take in the specific environment. Because our visual rationalization model uses the actor output only, the scalar critic output will be ignored for the purposed of this paper. However, in future work, exploring the critic’s explanations could be of interest. To ensure exploration early on an entropy regularization term $H$ is introduced with respect to the policy gradient,

[TABLE]

where $\beta$ is a hyper parameter discounting the entropy regularization.

3.2 Visual Rationalization

Our visual rationalization is based on Grad-CAM [15], and constitutes of computing a class-discriminative localization map $L_{GradCAM}^{s}\in{R}^{u\times v}$ using the gradient of any target class. These gradients are global-average-pooled to obtain the neuron importance weights $a^{c}_{k}$ for class $c$, for activation layer $k$ in the CNN222$k$ is usually chosen to be the last convolutional layer in the CNN.:

[TABLE]

Adapting this method in particular to the A3C actor output, let $h^{a}$ be the score for action $a$ before the softmax, $\alpha_{k}^{a}$ now represents the importance weight for state $a$ in activation layer $k$:

[TABLE]

with $|h|=|\mathcal{A}|$, the total amount of actions the agent can take. The gradient then gets weighted by the forward-pass activations $L^{k}$ and passes an ELU activation333the Exponential Linear Unit has been chosen in favor of the ReLU used in the original Grad-CAM paper due to the dying ReLU effect described in [22]. to produce a weighted class activation map:

[TABLE]

This activation map has values in the range $[0,1]$ with higher weights corresponding to a stronger response to the input state. This can be applied to the critic output in the same fashion. The resulting activation map can bilinearly extrapolated to the size of the input state and can then be overlayed on top of this state to produce a high-quality heatmap that indicate regions that motivate the agent to take action $a$. A visual representation of this process is depicted in figure 1.

4 Experiments

In this section, we first provide the details of our experimentation setup. We then show qualitative examples evaluating how our model performs in three of the Atari games. Throughout this section, red bounding boxes and red arrows indicate important regions of a state.

4.1 Setup

The Atari 2600 game environment is provided through the Arcade Learning Environment wrapper in the OpenAI Gym framework [2]. The framework has multiple version of each game but for the purpose of this paper the NoFrameskip-v4 environment will be used(OpenAI considers NoFrameskip the canonical Atari environment in gym and v4 is the latest version). Each state is represented as a $210\times 160\times 3$ pixel image with a 128-colour palette, and each state is preprocessed to a $84\times 84\times 1$ image as input to the network. A side-effect of this preprocessing is that the visual score will be removed from the state in most games, but the agent still gets the reward per state implicitly through the environment.

In our experiments, we use three Atari games, namely Pong, BeamRider and Seaquest, all depicted in Figure 2. All three games have a different action space (see Table 1), and a different long term-reward system for the agent to learn.

4.1.1 Pong.

Pong has six actions with three of the six being redundant (FIRE is equal to NOOP, LEFT is equal to LEFTFIRE and RIGHT is equal to RIGHTFIRE). The agent is displayed on the right and the enemy on the left and the first player to score 21 goals wins.

4.1.2 BeamRider.

In BeamRider the agent is displayed at the bottom and the agent has to traverse a series of sectors where each sector contains 15 enemies (remaining enemies is displayed at the top-left) and a boss at the end. The agent has three torpedoes that can be used specifically to kill the sector boss, but these can also be used to destroy debris that appear in later sectors. Learning how to use the torpedoes correctly is not necessary to succeed in the game, but it provides for long-term rewards in the form of bonus points.

4.1.3 Seaquest.

In Seaquest the agent is dependent on a limited amount of oxygen, depicted at the bottom of the state. The agent can ascend to the surface which will refill the oxygen bar and it drops off any swimmers that the agent has picked up along the way for a bonus reward. Resurfacing requires learning a long-term reward dependency which is not easily learned [16]. Surfacing is not just used to refill the oxygen bar but also to drop-off any swimmer that the agent has found underwater which results in additional points. A different way to get a positive reward is to kill sharks.

4.2 Learning A Policy

Training an agent to gain human-like or superhuman-like performance in a complex environment can take millions of input frames. In this section we take the same approach as Greydanus et al. [4], in which the authors argue that deep RL agents, during training, discard policies in favor for better ones. Seeing how an agent is reacting to different situations at different times of training might make it clear how an agent is trying to maximize long-term rewards. To demonstrate this, two agents have been trained for a different number of frames. The first model which will be called the Full Agent has been trained using (at least) 40 million frames. The second agent which will be called the Half Agent has been trained using 20 million frames, except for the case of Pong where it has been trained using 500,000 frames, due to the fact that Pong is an easier game to learn. The mean score and variance can be found in Table 2. For both games a sequence of states were manually sampled, after which both agents have evaluated444evaluated in this case means having forwarded each state that has been manually sampled through the model. the state to learn spatial-temporal information. States were manually sampled by having a person (one of the authors) play one episode of each game. The states were sampled manually to make sure the samples were not biased towards one agents’ policy.

4.2.1 Pong.

For Pong, the Full Agent has learned to shoot the ball in such a way that it scores by hitting the ball only once each round. The initial round might differ, but after that all rounds are the same: the Full Agent shoots the ball up high which makes the ball bounce off the wall all the way down over the opponent’s side, at which point the agent retreats to the lower right corner. This would indicate that the Full Agent is not reacting to the ball most of the time, but is waiting to exploit a working strategy that allows it to win each round. In contrast, he Half Agent is actively tracking the ball at each step and could potentially be losing some rounds because of this. The tracking behavior of the Half Agent is demonstrated in Figure 3 at frames 50, 51 and 53 indicated with a red box. In these frames the Half Agent’s attention is focused on the ball and the corresponding action is to go up to match it.

4.2.2 BeamRider.

For BeamRider, both agents have learned to hit enemies but the Full Agent has a higher average return. Looking at figure 4, both agents have a measure of attention on the two white enemy saucers, but the intensity of attention differs; the Full Agent has high attention on the enemies, in comparison with the Half Agent which has low attention on the enemies. The Half Agent is either going right which is essentially a NOOP in that area or it could be shooting at the incoming enemy. More interesting are the last two frames: 175 and 176. The attention of the Full Agent turns from the directly approaching enemy saucer to the enemy saucer on the left of it, and the agent would try to move into its direction (LEFTFIRE). the Full Agent’s attention in frame 176 is placed in a medium degree at the trajectory of its own laser that will hit the enemy saucer in the next frame. This could indicate that the Full Agent knows it will hit the target and is thus moving away from it, to focus on the other remaining enemy.

From the analysis of both agents another interesting result is discovered: the agents do not learn to properly use the torpedoes. At the beginning of each episode/level both agents would fire torpedoes until they are all used up and then continue on as usual. In Figure 5 this phenomena is demonstrated through a manually sampled configuration evaluated by the Full Agent only (the results are the same for the Half Agent). The torpedoes have not been used yet, on purpose, and there are enemies coming towards the agent at different time-steps. Looking at the Grad-CAM attention map, it would appear to be highly focused on the remaining three torpedoes in the upper right corner indicated by a red box. This occurs even when the action chosen by the agent is not of the UP-variety which would trigger firing a torpedo.

4.3 Agent Failing

A different way of looking at how rationalizations aid in understanding the behavior of an agent is by looking at when an agent fails at its task. In the context of BeamRider and Seaquest, this means looking at the last couple of frames before the agent dies.

4.3.1 BeamRider.

In the situation depicted in Figure 6 the agent is approached by a number of different enemies, one of which only appears after sector 7: the green bounce craft, depicted inside a red box in the first four frames. This is an enemy that can only be destroyed by shooting a torpedo at it, and it jumps from beam to beam trying to hit the agent which is what kills the agent eventually in the last frame. In all frames the Grad-CAM model is focused at the nearest three enemies, and the agent is shooting using LEFTFIRE in the direction of the green bounce craft. This could add extra weight to the idea that the agent does not know how to use the torpedoes correctly, but perhaps also that the agent might not be able to distinguish one enemy from another; the piece of green debris to the left of the green bounce craft looks quite similar to it.

4.3.2 Seaquest.

The agent playing Seaquest has a different problem: it has not learned the long term strategy of surfacing for oxygen. An example of a death due to this is depicted in Figure 7. The oxygen bar is highlighted by a red box, and it is noticeable that there is no direct or intense activations produced by the rationalization model on the oxygen bar. This could indicate that the agent has never made a correlation between the oxygen bar depleting and the episode ending. A multitude of factors could lead the agent to not learn this such as not having enough temporal knowledge or a lack of exploratory actions. A solution to this could be the use of Fine-Grained Action Repetition which selects a random action and performs this action for a decaying number of times [16].

4.4 Failure Cases of Our Model

Looking at Figure 7 a prominent activation is depicted in the form of a vertical bar at the top of the state. This vertical bar might seem a bit too ambiguous and even hard to interpret. This type of activation map come in two varieties: activations that highlight only static objects and avoid any non-static objects like agents or enemies and activations that do highlight seemingly at random.

The activations that highlight everything except for non-static objects are noticeable in Figure 8 in the case of Pong and BeamRider. For Pong, the activations are not focused on the ball but on everything except for the ball which could still indicate some pattern for the agent. For BeamRider the activations are highlighting areas directly next to the non-static agent and enemies in the state. This could indicate that the agent is calculating the trajectory of enemies or possible safe locations for it to go to.

The activations that are seemingly at random are depicted in Figure 8 in the last two Seaquest frames. A possible explanation for this could be that the agent is not provided with enough evidence and is indifferent to taking any action, which is reflected in the ambiguous activation map.

5 Conclusion

In this work, we have presented a post-hoc explanation framework that visually rationalizes the output of a deep reinforcement learning agent. Once the agent has made the decision of which action to take, the model propagates the gradients that lead to that action back to the image. Hence, it is able to visualize the activation map of the action as a heatmap. Our experiments on three Atari 2600 games indicate that the visualizations successfully attend to the regions such as the agent and the obstacle that lead to the action. We argue that such visual rationalizations, i.e. post-hoc explanations, are important to enable communication between users and the agents. Future work will include a quantitative evaluation in the form of a user study or developing an automatic evaluation metric for these kind of visual explanations.