EDGE: Explaining Deep Reinforcement Learning Policies

1 Introduction . Deep reinforcement learning has shown great success in automatic policy learning for various sequential decision-making problems , such as training AI agents to defeat professional players in sophisticated games [ 74 , 65 , 24 , 37 ] and controlling robots to accomplish complicated tasks [ 33 , 38 ] . However , existing DRL agents make decisions in an opaque fashion , taking actions without accompanying explanations . This lack of transparency creates key barriers to establishing trust in an agent ’ s policy and scrutinizing policy weakness . This issue significantly limits the applicability of DRL techniques in critical application fields ( e.g. , finance [ 47 ] and self-driving cars [ 11 ] ) . To tackle this limitation , prior research ( e.g. , [ 9 , 13 , 73 ] ) proposes to derive an explanation for a target agent ’ s action at a specific time step . Technically , this explanation can be obtained by pinpointing the features within the agent ’ s observation of a particular state that contribute most to its corresponding action at that state . Despite demonstrating great potential to help users understand a target agent ’ s ∗Equal Contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . individual actions , they lack the capability to extract insights into the overall policy of that agent . In other words , existing methods can not shed light on the general sensitivity of an agent ’ s final reward from a game in regards to the actions/states in that game episode . Consequently , these methods fall short in troubleshooting an agent ’ s policy ’ s weaknesses when it fails its task . We propose a novel explanation method to derive strategy-level interpretations of a DRL agent . As we discuss later in Section 3 , we define such interpretations as the identification of critical time steps contributing to a target agent ’ s final reward from each game episode . At a high level , our method identifies the important time steps by approximating the target agent ’ s decision-making process with a self-explainable model and extracting the explanations from this model . Specifically , given a well-trained DRL agent , our method first collects a set of episodes and the corresponding final rewards of this agent . Then , it fits a self-explainable model to predict final rewards from game episodes . To model the unique correlations in DRL episodes and enable high-fidelity explanations , rather than simply applying off-the-shelf self-explanation techniques , we develop a novel self-explainable model that integrates a series of new designs . First , we augment a Gaussian Process ( GP ) with a customized deep additive kernel to capture not only correlations between time steps but , more importantly , the joint effect across episodes . Second , we combine this deep GP model with our newly designed explainable prediction model to predict the final reward and extract the time step importance . Third , we develop an efficient inference and learning framework for our model by leveraging inducing points and variational inference . We refer to our method as “ Strategy-level Explanation of Drl aGEnts ” ( for short EDGE ) . 2 With extensive experiments on three representative RL games , we demonstrate that EDGE outperforms alternative interpretation methods in terms of explanation fidelity . Additionally , we demonstrate how DRL policy users and developers can benefit from EDGE . Specifically , we first show that EDGE could help understand the agent ’ s behavior and establish trust in its policy . Second , we demonstrate that guided by the insights revealed from our explanations , an attacker could launch efficient adversarial attacks to cause a target agent to fail . Third , we demonstrate how , with EDGE ’ s capability , a model developer could explain why a target agent makes mistakes . This allows the developer to explore a remediation policy following the explanations and using it to enhance the agent ’ s original policy . Finally , we illustrate that EDGE could help develop a defense mechanism against a newly emerging adversarial attack on DRL agents . To the best of our knowledge , this is the first work that interprets a DRL agent ’ s policy by identifying the most critical time steps to the agent ’ s final reward in each episode . This is also the first work that demonstrates how to use an explanation to understand agent behavior , discover policy vulnerabilities , patch policy errors , and robustify victim policies . 2 Related Work . Past research on DRL explanation primarily focuses on associating an agent ’ s action with its observation at a particular time step ( i.e. , pinpointing the features most critical to the agent ’ s action at a specific time ) . Technically , these methods can be summarized in the following categories . • Post-training explanation is a method that utilizes and extends post-training interpretation approaches ( e.g. , [ 56 , 28 , 36 , 35 ] ) to derive explanation from a DRL agent ’ s policy/value network and thus treat it as the interpretation for that DRL agent ( e.g. , [ 9 , 44 , 32 , 68 , 20 , 72 ] ) . • Model approximation is an approach that employs a self-interpretable model to mimic the target agent ’ s policy networks and then derives explanation from the self-interpretable model for the target DRL agent ( e.g. , [ 13 , 22 , 55 , 14 , 59 , 58 , 87 , 85 ] ) . • Self-interpretable modeling is an approach different from the model approximation techniques above . Instead of mimicking the target agent ’ s policy network , self-interpretable modeling builds a self-explainable model to replace the policy network . Since the new model is interpretable , one can easily derive an explanation for the target agent ( e.g. , [ 92 , 62 , 82 , 42 ] ) . • Reward decomposition is a method that re-engineers a DRL agent ’ s reward function to make the reward gained at each time step more meaningful and explainable . With the more meaningful reward in hand , at each time step , one could use the instant reward gain to interpret the agent ’ s action ( e.g. , [ 73 , 46 , 54 ] ) . 2The source code of EDGE can be found in https : //github.com/Henrygwb/edge . From the objective perspective , our work is fundamentally different from the above DRL explanation research . Rather than pinpointing the features – in an observation – critical for an agent ’ s action , our work identifies the critical time steps contributing to the agent ’ s final reward . Using our explanation , one can better understand the agent ’ s policy , unveil policy weakness , and patch policy errors ( as shown in Section 5 ) . In Supplement S7 , we further conduct a user study to demonstrate the superiority of our method against the above explanation approaches in pinpointing good policies and performing policy forensics . Note that there are two other methods that also understand a DRL policy through the agent ’ s previous memories [ 49 , 23 ] . These works are fundamentally different from ours in two perspectives . First , both methods have a different explanation goals from our work . Specifically , Koul et al . [ 49 ] focuses on identifying whether the action at each time step depends more on the current observation or the previous states . The method proposed in [ 23 ] pinpoints the important steps w.r.t . the subsequent transitions in the FSM extracted from the target agent rather than the final result of an episode . Second , both methods can be applied only to white-box RNN policies , whereas our method is applicable to DRL policies with arbitrary network structures . 3 Key Technique . 3.1 Problem Setup . Consider a DRL game with an agent trained with Q-learning [ 86 , 60 ] or policy gradient [ 48 , 69 , 70 ] . Our work aims to explain this agent ’ s policy by identifying the important steps contributing most to a game episode ’ s final reward/result . To ensure practicability , we allow access only to the environment states , agent ’ s actions , and rewards . We assume the availability of neither the value/Q function nor the policy network . Formally , given N episodes T = { X ( i ) , yi } i=1 : N of the target agent , X ( i ) = { s ( i ) t , a ( i ) t } t=1 : T is the i−th episode with the length T , where s ( i ) t ∈ Rds and a ( i ) t ∈ Rda are the state and action at the time step t. yi is the final reward of this episode.3 Our goal is to highlight the top-K important time steps within each episode X ( i ) . Possible Solutions and Limitations . The most straightforward approach of identifying important time steps is to use the output of the value/Q network as indicators and pinpoint the time steps with the top K highest value/Q function ’ s outputs as the top K critical steps . However , since we do not assume the availability of these networks , this method is not applicable to our problem . A more realistic method is to fit a seq2one model ( i.e. , RNN ) that takes as input the state-action pairs in an episode and predicts the final reward of that episode . With this prediction model , one could utilize a post-training explanation method to derive the time step importance . However , existing post-training explanation techniques usually require approximating the target DNN with more transparent models , which inevitably introduces errors . Additionally , many post-training methods are vulnerable to adversarial attacks [ 61 , 30 , 94 ] or generate model-independent explanations [ 1 , 64 , 2 , 78 ] . As we will show later in Section 4 and Supplement S3 & S5 , these limitations jeopardize the post-training explanations ’ fidelity . A more promising direction is to fit a self-explainable model to predict the final reward . Existing research has proposed a variety of self-explanation methods . Most of them do not apply to our problem because they either can not derive feature importance as explanations [ 5 , 52 , 50 , 18 ] , can not be applied to sequential data [ 21 , 27 , 12 ] , or require explanation ground truth [ 16 , 66 ] . In this work , we consider two self-explainable models that are designed to fit and explain sequential data – an RNN augmented with attention [ 10 , 39 , 34 ] and rationale net [ 51 ] . Technically , both models have the form of g ( θ ( x ) ⊙ x ) , where θ ( · ) is a weight RNN or an attention layer and g ( · ) is the prediction RNN . The output of θ ( · ) can be used to identify the important steps in the input sequence . Despite extracting meaningful explanations , recent research [ 45 , 89 , 17 ] reveals that the explanations given by θ ( · ) can not faithfully reflect the associations ( i.e. , feature importance ) learned by the subsequent prediction model g ( · ) , leading to an even lower fidelity than the post-training explanations in some applications . Additionally , these models are not designed to explain an RL agent and can not fully capture the dependencies within the episodes of that agent . Specifically , the episodes collected from the same agent tend to exhibit two types of dependencies : dependency between the time steps within an episode and the dependency across different episodes . Although they consider the dependency 3For the games with delayed rewards , such as MuJoCo [ 84 ] and Atari Pong [ 8 ] , where a non-zero reward rT is assigned only to the last step of a game , we use rT as yi . For the games with instant rewards ( e.g. , OpenAI CartPole [ 15 ] ) , we compute an episode ’ s total reward as yi , i.e. , yi = ∑ t rt . within each input sequence , these methods can not capture the correlations between different inputs . As is shown in Section 4 and Supplement S3 & S5 , this also jeopardizes their explanation fidelity . 3.2 Explanation Model Design of EDGE In this work , we design a novel self-explainable model by adopting a very different design than existing methods . First , to better capture the associations ( i.e. , feature importance ) learned by the prediction model , we add the explainable module to the final layer rather than the input layer of the prediction model . Formally , our model can be written as g ( f ( x ) ) , where f ( · ) is a feature extractor and g ( · ) is an explainable prediction model . Second , we design a deep Gaussian Process as the feature extractor to capture the correlations between time steps and those across different episodes , which are often exhibited in a set of episodes collected from the same DRL agent . In addition to capturing different levels of correlations , another advantage of deep GPs over typical DNNs is that GPs model the joint distribution of the output signals , enabling access to the output signals ’ uncertainty . Finally , we design an interpretable Bayesian prediction model to infer the distribution of final rewards and deliver time step importance . Below , we first give an overview of our proposed model . Then , we describe how to adapt the traditional GP model to our problem , followed by the design of the final prediction model . Overview . As shown in Fig 1 , given an episode of the target agent X ( i ) , EDGE first inputs it into a RNN encoder , which outputs the embedding of each time step in this episode { h ( i ) t } t=1 : T . EDGE also passes the last step ’ s embedding through a shallow MLP to obtain an episode embedding e ( i ) . Then , EDGE adopts our proposed additive GP framework to process { h ( i ) t } t=1 : T and e ( i ) and obtains a latent representation of the whole episode f ( i ) 1 : T . As introduced later , this representation is able to capture the correlations between time steps and those across episodes . Finally , EDGE inputs f ( i ) 1 : T into our prediction model f ( i ) 1 : T and get the predicted final reward of the input episode . As detailed later , our prediction model is designed based on a linear regression , whose regression coefficient can be used to identify important time steps within in the input episode . Additive GP with Deep Recurrent Kernels . Gaussian Process defines a distribution over an infinite collection of random variables , such that any finite subset of variables follows a multivariate Gaussian distribution [ 63 ] . In Statistical modeling , GP defines the prior of a non-parametric function f : X → R. Formally , if f has a GP prior , i.e. , f ∼ GP ( 0 , kγ ) , where kγ ( · , · ) is a positive semidefinite kernel function parameterized by γ , any finite collections of f ∈ RN follows a multivariate Gaussian distribution ( f |X ) ∼ N ( 0 , KXX ) . Here , KXX ∈ RN×N is the covariance matrix , with ( KXX ) ij = kγ ( xi , xj ) . In our model , we adopt the widely applied square exponential ( SE ) kernel function : kγ ( xi , xj ) = exp ( − 12 ( xi − xj ) TΘk ( xi − xj ) ) , with γ = Θk . Traditional GP with SE kernel [ 63 ] assumes the input space is Euclidean , which is usually invalid for real-world data with high-dimensional inputs [ 3 ] . To tackle this challenge , recent research [ 91 , 53 ] proposes to conduct dimensional reduction via a DNN and then apply a GP to the DNN ’ s latent space . They show that the resultant deep kernel models achieve similar performance to DNNs on complicated datasets . In our model we capture the sequential dependency within an episode by using an RNN as the deep net inside the kernel function . Specifically , given an episode X ( i ) , we first concatenate the state and action ( i.e. , x ( i ) t = [ s ( i ) t , a ( i ) t ] ) , input them into an RNN hϕ , and obtain the latent representation of this episode : { h ( i ) t } t=1 : T , where h ( i ) t ∈ Rq is the state-action embedding at the time t. We also compute an episode embedding by passing the last step ’ s hidden representation through a shallow MLP eϕ1 : h ( i ) T → e ( i ) ∈ Rq . After obtaining { h ( i ) t } t=1 : T and e ( i ) , we then adopt the additive GP framework to capture the correlations between time steps and those across episodes . Formally , an additive GP is the weighted sum of J independent GPs , i.e. , f = ∑ J αjfj . Here , fj ∼ GP ( 0 , kj ) is the j-th GP component , in which the covariance function kj is typically applied to a subset of input features . By assigning every component a GP prior , one can ensure that the mixed-signal f also follows a GP prior [ 25 ] . Following this framework , we construct our deep GP model as the sum of two components ft and fe . Specifically , ft ∼ GP ( 0 , kγt ) models the correlations between time-steps , where the covariance between the t-th steps in episode i and the k-th steps in episode j can be computed by kγt ( h ( i ) t , h ( j ) k ) . Going beyond modeling the correlations between individual steps , fe ∼ GP ( 0 , kγe ) captures a higher level cluster structures within the collected episodes , i.e. , the similarity between episodes . Formally , the episode-level covariance between any pair of time steps in episode i and j is given by kγe ( e ( i ) , e ( j ) ) . Finally , our deep additive GP model can be expressed as : f = αtft + αefe , where αt and αe are the component weights . Given a set of collected episodes represented by T ∈ RN×T× ( ds+da ) , f ∈ RNT is given by : f |X ∼ N ( 0 , k = α2tkγt +α2ekγe ) , where X ∈ RNT× ( ds+da ) is the flattened matrix of T. Prediction Model . To ensure explanability , we use a linear regression as the base of our prediction model , where the regression coefficients reflect the importance of each input entity . Specifically , we first convert the flattened response f back to the matrix form F ∈ RN×T , where the i-th row F ( i ) ∈ RT is the i-th episode ’ s encoding given by our GP model . Then , we define the conditional likelihood for the discrete and continuous final reward , respectively . When yi is continuous , we follow the typical GP regression model [ 63 ] and define the yi = F ( i ) wT + ϵ1 , where w ∈ R1×T is the mixing weight and ϵ1 ∼ N ( 0 , σ2 ) is the observation noise . The conditional likelihood distribution is yi|F ( i ) ∼ N ( F ( i ) wT , σ2 ) . For the discrete final reward with a finite number of possible values , we use the softmax prediction model to perform classification . Formally , we define yi|F ( i ) follows a categorical distribution with p ( yi = k|F ( i ) ) = exp ( ( F ( i ) WT ) k ) ∑ k exp ( ( F ( i ) WT ) k ) . W ∈ RK×T is the mixing weight , where K is the total number of classes . Finally , we combine all the components together and write our explanation model as ( A illustration of our proposed model can be found in Fig . 1 . ) : 4 f |X ∼ N ( 0 , k = α2tkγt + α 2 ekγe ) , yi|F ( i ) ∼ { Cal ( softmax ( F ( i ) WT ) ) , If conducting classification N ( F ( i ) wT , σ2 ) , otherwise , ( 1 ) where the mixing weight is constant . This indicates the time step importance derived from the mixing weight is a global explanation . 5 According to the insight that time steps with a high correlation tend to have a joint effect ( similar importance ) on the game result , we could combine the global explanation with the time step correlations in Kt ( X , X ) to gain a fine-grained understanding of each game . Specifically , given an episode and the top important steps indicated by the mixing weight , we can identify the time steps that are highly correlated to these globally important steps and treat them together as the local explanation of that episode . Supplement S1 introduces another way of deriving episode-specific explanations by replacing the constant mixing weight with a weight obtained by a simple DNN . Note that the episode correlations in Ke ( X , X ) reveal the cluster structure within a set of episodes , which helps categorize the explanations of similar episodes .

Summary: This paper proposes a new pipeline for understanding agent behavior + actions, and identifying which timesteps contributed to the final reward. The key component of this method is the pipeline, which consists of the following:  * An RNN which takes in state-action pairs, and "compresses" the information via its RNN-state at every timestep. * All RNN-states over all time steps are taken and inserted into a zero-mean Gaussian Process (GP), which consists of two independent GPs, one for correlations between time-steps, and another for correlations between episodes. * A feature vector $\mathbf{f}$ is then sampled from this GP, and used to classify the reward.   Because a GP is involved, variational inference is needed for training.  The paper then experimentally verifies this result on Atari Pong + Adversarial Mujoco tasks by showing competitive performance on fidelity (i.e. metric based on replacing an intelligent agent with random actions) and visualizing which parts of an episode contribute most to the fidelity score. Finally, the method is also applied to successfully perform adversarial attacks which also beat previous baselines.