SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zelei Cheng*
Evanston, IL 60208
zelei.cheng@northwestern.edu
Xian Wu*
xianwu2024@u.northwestern.edu
Jiahao Yu*
jiahao.yu@northwestern.edu
Wenhai Sun
sun841@purdue.edu
Wenbo Guo
henrygwb@purdue.edu
Xinyu Xing
xinyu.xing@northwestern.edu
Despite the promising performance of deep reinforcement learning (DRL) agents in many challenging scenarios, the black-box nature of these agents greatly limits their applications in critical domains. Prior research has proposed several explanation techniques to understand the deep learning-based policies in RL. Most existing methods explain why an agent takes individual actions rather than pinpointing the critical steps to its final reward. To fill this gap, we propose StateMask, a novel method to identify the states most critical to the agent's final reward. The high-level idea of StateMask is to learn a mask net that blinds a target agent and forces it to take random actions at some steps without compromising the agent's performance. Through careful design, we can theoretically ensure that the masked agent performs similarly to the original agent. We evaluate StateMask in various popular RL environments and show its superiority over existing explainers in explanation fidelity. We also show that StateMask has better utilities, such as launching adversarial attacks and patching policy errors.
Deep reinforcement learning has shown promising performance in a variety of challenging scenarios, such as allowing a robot to learn a control policy directly from the camera or sensor inputs [1, 2, 3, 4, 5] or training an agent to master Go games [6, 7, 8]. However, DRL policies typically lack explainability, which is intrinsically difficult for humans to comprehend and establish trust. To broaden the adoption of DRL-based applications in critical fields, there is a pressing need to increase the transparency of DRL agents through explanation.
Prior research has proposed various methods to derive an explanation for DRL. For example, researchers utilized the saliency map (e.g., [9, 10, 11, 12, 13]) and model approximation methods
(e.g., [14, 15, 16, 17]) to identify which regions in an agent's observation are most critical to the agent's current action. Considering the nature of DRL, researchers also explored the relationship between the agent's states and its final reward. Take the work [18] as an example. Researchers demonstrated that the state-reward relationship could also be used as an explanation for DRL agent developers. Besides, they also showed that pinpointing the critical time steps towards the final reward could help agent developers better understand agent behaviors and even troubleshoot policy defects.
In this work, our research focuses on the DRL explanation methods that explore the state-reward relationships. We evaluate the existing DRL explanation approaches in this category on a range of tasks from simple simulated environments to sophisticated real-world applications, and from normal-form games to extensive-form games. We discovered that existing state-reward-based DRL explanation methods provide relatively low fidelity in state-reward relationship identification for various reasons. For example, the method proposed in [18] suffers from low explanation fidelity because of the limited capacity of their proposed approximation model and the error introduced in the final reward approximation. The methods introduced by [19] and [20] experience low fidelity because an agent's value function does not always reflect the state importance to its final reward.
To this end, we propose StateMask², a novel method to pinpoint the time steps or, in other words, states most critical to a DRL agent's expected total reward. At a high level, StateMask works as follows. For a target agent, we train a mask net that blinds the agent to take random actions at certain time steps. If a time step is critical, the mask net does not blind the target agent. The agent takes action as its policy suggests. In contrast, if the time step is less critical, the mask net blinds the agent. The agent takes random actions. The rationale behind this design is that a DRL agent's final reward is not equally important for each time step. Randomly taking some actions at non-critical time steps has no impact on the agent's final reward. Technically, to learn the mask net, we carefully design an objective function that minimizes the performance difference between the original agent and the masked agent. We demonstrate, both theoretically and empirically, that performance differences could monotonically decrease during training.
While other explanation methods could track down the time steps most or least critical to an agent's final reward, this is the first work that could identify an agent's critical time steps with high fidelity across various types of RL environments, especially sophisticated normal-form games and extensive-form games. Evaluating StateMask in various game environments, we show that the explanation derived from our method demonstrates significantly higher fidelity than those produced by existing explanation techniques. Compared with the existing methods, StateMask allows humans to understand the agent's behaviors better and calibrate appropriate trust. Agent users can easily pinpoint an agent's policy defects and patch the identified defeat under the guidance of explanations. Last but not least, security professionals can launch adversarial attacks using the explanation from our method.
In summary, this paper makes the following contributions.
- We propose a new explanation method - StateMask and provide a theoretical analysis. We show that an agent armed with explainability is equivalent to the original agent in terms of agent performance.
- We evaluate StateMask on 10 different tasks ranging from simple normal-form games (e.g., Pong) to sophisticated normal-form games (e.g., StarCraft II) and extensive-form games (e.g., Connect 4). We show that StateMask can significantly outperform existing explainers.
- We demonstrate that our explanation can be highly helpful for understanding agent behavior, launching adversarial attacks, and even remediating the agent's errors.
DRL Explanation. Existing DRL explanation techniques can be mainly categorized into two classes based on their different explanation objectives: (1) the observation-action-based method pinpointing which regions in an agent's observation are most important to its action at an individual time step; 2) the state-reward-based method tracking down which time steps have the most critical influence on the agent's final reward.
2The source code of StateMask can be found in https://github.com/nuwuxian/RL-state_mask.
Recall that a DRL agent's policy network takes as input the agent's current observation and outputs the corresponding action. In this setup, pinpointing regions in an observation most critical to the corresponding action is the same as identifying the most important regions in a deep neural network's (DNN) input. As such, existing works along this direction mainly utilize or extend the explanation methods developed for DNN classifiers to derive such explanations. Technically, they mainly leverage three lines of methods: post-hoc white-box explanation methods [21, 22, 23, 24], post-hoc black-box explanation methods [15, 25, 26, 27], and self-explainable methods [28, 29, 30, 31]. In addition, another line of work identifies important regions by decomposing the agent's instant reward to regions in its current observation and selecting the regions with high rewards [32, 33, 34, 35].
Regarding the second type of explanation methods, most of the techniques identify the critical steps based on the value function [19, 20, 36]. Given that the value function measures the goodness of the current state towards the agent's expected total reward, it is natural to use it for critical time step selection. Recently, Guo et al. [18] proposed EDGE that establishes state-reward relationship by collecting a set of trajectories and then approximating an explanation model offline. However, as discussed in [18], EDGE has two limitations. First, the approximation method is based on a statistical model not effective in deriving explanations for agents that usually generate long trajectories. Second, for each new trajectory, EDGE needs to build an independent approximation model, which significantly limits its efficiency and scalability.
Our proposed technique falls into the second type of explanation, which follows a different design path to overcome the limitations above. As we will discuss later, our method is superior to existing techniques in multiple aspects: higher explanation fidelity (Section 4), better utilities (Section 5), as well as better generalizability (to extensive-form games) and practicability (in the cases where the agent's value function/network is not available).
Perturbation-based Explanation. Perturbation techniques have been widely utilized within the field of computer vision to gain a deeper understanding of black-box image classifier [37, 38, 39, 40]. These approaches involve introducing minor modifications to the target image and observing how the classifier's outputs change, which enables us to identify the regions of greatest importance within the original image. Our explanation method shares a similar idea with these perturbation methods. Rather than explaining supervised learning (e.g., image classifier), our focus is on interpreting a DRL agent by introducing perturbations in non-critical time steps. This temporal perspective provides unique insights into the DRL agent, which has not been explored in prior research.
3 Key Technique
In the following, we describe the rationales behind our design and then discuss the key technical challenges, followed by the technical details of tackling the challenges.
Assumptions. To ensure the practicability of our method, we do not assume access to the agent's value function or policy network. Instead, we only assume having access to the agent's observations and the ability to query the agent's policy and thereby obtain corresponding actions. Note that our explanation method does not change the target agent's policy. For a multi-agent environment, we explain each agent individually.
Design Rationale. To tackle the limitations confronted by existing state-reward-based explanation methods, our idea is to parameterize the importance of the agent's current time step as a neural network model (denoted as state mask). The state mask net takes as input the current state and outputs its importance to the agent's final reward. The key challenge of this solution is to propose a proper criterion or metric to decide the importance of the current time step and thus design an objective function to train this state mask. In our work, we draw insights from existing counterfactual-based explanation methods designed for supervised classifiers [37, 41, 42, 43]. At a high level, these methods identify important features in an input sample by nullifying or perturbing the features and observing the corresponding changes in the classification output. They deem the features triggering the largest prediction changes as the most important ones. Similarly, in our problem, we can decide whether a time step is important by randomizing the agent's action at a certain time step and observing changes in its final reward. A small variation in the final reward indicates the reward is not sensitive
(a) The DRL agent just launched the ball.
(b) The DRL agent attends to catch the ball.
Figure 1: Illustration of the Pong game. The DRL agent controls the yellow (right) paddle and a non-RL rule-based opponent controls the white (left) paddle.
to the action change, further implying the low importance of the state. In contrast, a large reward change infers the high importance of the corresponding state.
To elaborate on this idea, we use apong game as an example. As shown in Figure 1(a), after the DRL agent launches the ball and when the ball is moving towards its opponent, the agent has a few time steps of freedom to choose random actions. This is because the agent's actions at these steps will not change the movement of the ball and will be less likely to influence the opponent's actions as well (given the opponent's focus is on the ball). As a result, forcing the DRL agent to take random actions at these steps will likely have no influence on the game's final result and thus trigger only minor changes to the agent's final reward. However, if we randomly perturb the DRL agent's action at the time steps when the agent is about to touch the ball (Figure 1(b)), it will likely cause the agent to fail and thus trigger a large variation in the agent's final reward.
State Mask Modeling. Guided by this idea, we design the state mask that takes the current state as input and outputs the probability of whether to randomize the agent's action at the current state. The objective of this state mask is to blind a target agent and thus vary its actions at certain states in a trajectory without changing the agent's final reward. This probability can be taken as the importance of the current state (i.e., the higher the probability is, the lower the importance is). This design introduces a challenge. That is how to decide whether varying the action at the current state will influence the agent's future reward.
To address this challenge, we model it as an RL problem, where we fix the target agent's policy and treat it as part of the environment. Then, we introduce another agent to decide whether to randomize the target agent's action at each step or not. Its policy can be modeled as the state mask above. This agent's goal is to preserve the target agent's total reward while randomizing actions at more states.
Figure 2: The StateMask framework.
Following this idea, we propose a novel explanation framework, StateMask. Figure 2 illustrates the proposed approach.
trates this explanation system. As mentioned above, we add a state mask (agent) to the environment. We fit the current state  $s_t$  into both the state mask  $\tilde{\pi}$  and the target agent  $\bar{\pi}$ . For a given state  $s_t$ , the state mask outputs a binary action  $a_t^e$  of either "zero" or "one", and the target agent will sample the action  $a_t$  from its policy. The final action is determined by the following equation
$$
a _ {t} \odot a _ {t} ^ {e} = \left\{ \begin{array}{l l} a _ {t}, & \text {i f} a _ {t} ^ {e} = 0, \\ a _ {\text {r a n d o m}} & \text {i f} a _ {t} ^ {e} = 1, \end{array} \right. \tag {1}
$$
where  $a_{t}^{e} = 0$  represents keeping the target agent's current action. On the contrary,  $a_{t}^{e} = 1$  refers to replacing the target agent's current action with a random action  $a_{\mathrm{random}}$  uniformly sampled from its action space. Note that  $a_{\mathrm{random}}$  is different from  $a_{t}$ .
The goal of our state mask is to find the non-important time steps and randomize their actions without changing the expected total reward of the target agent. Formally, it can be expressed as
$$
J (\theta) = \operatorname {m i n i m i z e} | \eta (\pi) - \eta (\bar {\pi}) |, \tag {2}
$$
where  $\bar{\pi}$  denotes the policy of the target agent and  $\pi$  denotes the policy of the perturbed agent (i.e., integration of the state mask  $\tilde{\pi}$  and target agent  $\bar{\pi}$ ).  $\eta(\cdot)$  is the expected total reward of an agent by following a certain policy. By fixing  $\bar{\pi}$ ,  $\eta(\bar{\pi})$  is also fixed and  $\eta(\pi)$  depends only on the parameters
<sup>3</sup>We measure the performance gap by the  $L_{p}$  distance, and without loss of generality, we use the  $L_{1}$  norm.
of the state mask. As a result, by solving Eqn. (2), we can obtain a state mask. By applying this resolved mask to each state, we will be able to assess the state importance at any time step.
Key Optimization Challenges and Solutions. Obtaining a non-trivial solution of Eqn. (2) presents the following challenge. To ensure the optimization function is differentiable, we need to decompose and remove the absolute value operation. A straightforward solution is to compare the value of  $\eta(\pi)$  and  $\eta(\bar{\pi})$  during each training iteration. If  $\eta(\pi)$  is larger than  $\eta(\bar{\pi})$ , we update the state mask by minimizing  $\eta(\pi)$ . Otherwise, we maximize  $\eta(\pi)$ . With simple reformulation, we can then leverage the state-of-art policy gradient method - Proximal Policy Optimization (PPO) [44] to update the state mask in each iteration. However, the key limitation of this solution is the lack of convergence guarantee. Suppose the policy of the perturbed agent  $\pi$  is near a local optimal and  $\eta(\pi)$  is less than  $\eta(\bar{\pi})$ . Then, updating the perturbed policy by maximizing  $\eta(\pi)$  may lead to an overly large  $\eta(\pi)$ , which, in turn, requires minimizing  $\eta(\pi)$ . When this situation occurs, the perturbed policy  $\pi$  may oscillate around the local optimal without actually converging. To address this problem, we borrow the idea of TRPO technique [45], carefully designing a surrogate objective function for Eqn. (2). As we will elaborate below, through the surrogate objective function, we can guarantee the difference between  $\eta(\pi)$  and  $\eta(\bar{\pi})$  (i.e.,  $|\eta(\pi) - \eta(\bar{\pi})|$ ) monotonically decreases during the training process, and our optimization process will at least converge to a local optimal rather than oscillating around it.
3.2 Technical Details
In this section, we formally introduce our proposed StateMask framework, including modeling the state mask and designing the surrogate objective function to ensure monotonicity.
Defining the State Mask. An infinite-horizon Markov Decision Process (MDP) is defined as  $\{\mathcal{S},\mathcal{A},P,R,\gamma \}$ , where  $\mathcal{S}$  and  $\mathcal{A}$  are the state and action space.  $P:\mathcal{S}\times \mathcal{A}\to \Delta (\mathcal{S})$  is the state transition function.  $R:S\times \mathcal{A}\rightarrow \mathbb{R}$  is the reward function and  $\gamma \in (0,1)$  is the discount factor. A policy  $\pi$  is a mapping from  $\mathcal{S}$  to a probability distribution over  $\mathcal{A}$ . The expected total reward of a policy  $\pi$  is defined as  $\eta (\pi) = \mathbb{E}_{s_0,a_0,\dots}[\sum_{t = 0}^{\infty}\gamma^t R(s_t,a_t)]$ .
Based on this definition, we define our explanation as an Explanation Markov Decision Process problem (E-MDP), in which we formalize state mask as an explanation policy  $\tilde{\pi}:\mathcal{S}\to \Delta (\mathcal{A}^e)$  where  $\mathcal{A}^e = \{0,1\}$ . With this definition, the integration of the explanation policy  $\tilde{\pi}$  and the target policy  $\bar{\pi}$  (that we aim to explain) then becomes  $\pi (\cdot |s) = \tilde{\pi} (a^{e} = 0|s)\bar{\pi} (\cdot |s) + \tilde{\pi} (a^{e} = 1|s)a_{\mathrm{random}}$  where  $a_{\mathrm{random}}\neq \bar{\pi} (\cdot |s)$ .
As mentioned above, we fix  $\bar{\pi}$ . In this way, we can transform the E-MDP into a single-player environment, where the state-transition, reward function, and  $\eta(\pi)$  depend only on the policy of the state mask  $\tilde{\pi}$ . As such, we can redefine  $\eta(\pi)$  as  $\eta(\tilde{\pi})$ . The goal here is to train an optimal policy  $\tilde{\pi}$  to minimize the difference between  $\eta(\bar{\pi})$  and  $\eta(\tilde{\pi})$ .
Solving the Explanation Policy. With the definitions above, the objective function in Eqn. (2) could be reformulated as  $J(\theta) = \text{minimize}|\eta(\tilde{\pi}_{\theta}) - \eta(\bar{\pi})|$ , where  $\theta$  refers to the parameters of the explanation policy. As is discussed in Section 3.1, to ensure the monotonic property, we subtly design a surrogate objective function to solve  $\theta$ . To introduce our designed objective function, we first present the following lemma.
Lemma 1. ([45]) The difference between  $L_{\tilde{\pi}_{\theta_{old}}}\left(\tilde{\pi}_{\theta}\right)$  and  $\eta (\tilde{\pi}_{\theta})$  is bounded by
$$
\left| \eta \left(\tilde {\pi} _ {\theta}\right) - L _ {\tilde {\pi} _ {\theta_ {o l d}}} \left(\tilde {\pi} _ {\theta}\right) \right| \leq C \mathbb {K L} ^ {\max } \left(\tilde {\pi} _ {\theta_ {o l d}} (\cdot | s) \| \tilde {\pi} _ {\theta} (\cdot | s)\right). \tag {3}
$$
$\tilde{\pi}_{\theta}$  denotes the new policy, updated from  $\tilde{\pi}_{\theta_{old}}$ $L_{\tilde{\pi}_{\theta_{old}}}(\tilde{\pi}_{\theta}) = \eta (\tilde{\pi}_{\theta_{old}}) + \sum_{s}\rho_{\tilde{\pi}_{\theta_{old}}}(s)\sum_{a^{e}}\tilde{\pi}_{\theta}(a^{e}|$ $s)A_{\tilde{\pi}_{\theta_{old}}}(s,a^{e})$  , where  $\rho_{\tilde{\pi}_{\theta_{old}}}(s) = \sum_t\gamma^t P(s_t = s|\tilde{\pi}_{\theta_{old}})$  and  $A_{\tilde{\pi}_{\theta_{old}}}$  is the advantage function of  $\tilde{\pi}_{\theta_{old}}$ $\mathbb{KL}^{\max}(\tilde{\pi}_{\theta_{old}}(\cdot |s)\| \tilde{\pi}_{\theta}(\cdot |s)) = \max_s\mathbb{KL}(\tilde{\pi}_{\theta_{old}}(\cdot |s)\| \tilde{\pi}_{\theta}(\cdot |s))$  and  $C$  is a constant.
Based on Lemma 1, we can derive the following theorem (The derivation is in Supplement S1).
Theorem 1. The following inequalities are the sufficient conditions for satisfying  $|\eta(\tilde{\pi}_{\theta}) - \eta(\bar{\pi})| \leq |\eta(\tilde{\pi}_{old}) - \eta(\bar{\pi})|$ .
$$
\begin{array}{l} L _ {\tilde {\pi} _ {\theta_ {o l d}}} \left(\tilde {\pi} _ {\theta}\right) \geq \eta \left(\tilde {\pi} _ {\theta_ {o l d}}\right) + C \mathbb {K L} ^ {m a x} \left(\tilde {\pi} _ {\theta_ {o l d}} (\cdot | s) \| \tilde {\pi} _ {\theta} (\cdot | s)\right), \\ L _ {\tilde {\pi} _ {\theta_ {o l d}}} \left(\tilde {\pi} _ {\theta}\right) + C \mathbb {K L} ^ {\max } \left(\tilde {\pi} _ {\theta_ {o l d}} (\cdot | s) \| \tilde {\pi} _ {\theta} (\cdot | s)\right) \leq 2 \eta (\bar {\pi}) - \eta \left(\tilde {\pi} _ {\theta_ {o l d}}\right). \tag {4} \\ \end{array}
$$
Algorithm 1 The learning algorithm for training the mask net.
Input: Target agent's policy  $\bar{\pi}$  and the estimation of its expected total reward  $\eta (\bar{\pi})$
Output: Mask net  $\tilde{\pi}_{\theta}$
Initialization: Initialize the weights  $\theta$  for the mask net  $\tilde{\pi}_{\theta}$ , Lagrange multiplier  $\lambda = 0$
for iteration  $= 1,2,\ldots$  do
Run the current policy  $\tilde{\pi}_{\theta_{old}}$  in the environment for  $T$  steps and get estimation for  $A_{\tilde{\pi}_{\theta_{old}}}$
Solve the objective function with respect to  $\theta$  in Eqn. (7) and update  $\theta_{old} \gets \theta$
Solve the objective function with respect to  $\lambda$  in Eqn. (8) and update  $\lambda_{old} \gets \lambda$
Based on Theorem 1, we propose the following objective function to update  $\tilde{\pi}_{\theta}$  from  $\tilde{\pi}_{\theta_{old}}$
$$
\max  _ {\theta} L _ {\tilde {\pi} _ {\theta_ {o l d}}} \left(\tilde {\pi} _ {\theta}\right)
$$
$$
\text {s . t .} L _ {\tilde {\pi} _ {\theta_ {o l d}}} \left(\tilde {\pi} _ {\theta}\right) \leq 2 \eta (\tilde {\pi}) - \eta \left(\tilde {\pi} _ {\theta_ {o l d}}\right) - \delta , \tag {5}
$$
$$
\mathbb {E} _ {a ^ {e} \sim \tilde {\pi} _ {\theta}} [ a ^ {e} ] \geq c, \quad \mathbb {E} _ {s \sim \rho_ {\tilde {\pi} _ {\theta_ {o l d}}}} [ \mathbb {K L} (\tilde {\pi} _ {\theta_ {o l d}} (\cdot | s) \| \tilde {\pi} _ {\theta} (\cdot | s)) ] \leq \delta .
$$
where  $\mathbb{E}_{a^e\sim \tilde{\pi}_\theta}[a^e ]\geq c$  constrains the lower bound sparsity for the output of  $\tilde{\pi}_{\theta}$  across all states, with  $c\in [0,1]$  as a constant hyper-parameter. It rules out the trivial solution where the state mask outputs 0 at any state. As detailed in Supplement S2, a policy  $\tilde{\pi}_{\theta}$  solved from Eqn. (5) satisfies the conditions in Eqn. (4) and thus enables the desired monotonicity.
Eqn. (5) can be further transformed into the following function based on the Lagrangian duality.
$$
\min  _ {\lambda \geq 0} \max  _ {\theta} L _ {\tilde {\pi} _ {\theta_ {o l d}}} (\tilde {\pi} _ {\theta}) - \lambda (L _ {\tilde {\pi} _ {\theta_ {o l d}}} (\tilde {\pi} _ {\theta}) - 2 \eta (\bar {\pi}) + \eta (\tilde {\pi} _ {\theta_ {o l d}}) + \delta),
$$