Why the Agent Made that Decision: Contrastive Explanation Learning for Reinforcement Learning

1 Introduction

Deep Reinforcement Learning (DRL) is a powerful technology in machine intelligence, widely used for many applications [Sutton et al., 1999]. However, understanding a DRL agent's decision-making process is challenging, due to the inherent lack of explainability in the high-dimensional, non-linear structure of its underlying Deep Neural Network (DNN) [Heuillet et al., 2021; Hickling et al., 2023]. The lack of transparency undermines users' trust, driving the development of Explainable AI (xAI).

Various methods in computer vision have been proposed to enhance the transparency of AI systems [Ribeiro et al., 2016; Selvaraju et al., 2017; Bach et al., 2015; Shrikumar et al., 2017; Lundberg and Lee, 2017]. At the core, they share a common foundation: attributing the classifier's outputs to more interpretable features and using a saliency map to visualize these attributions. Their only differences are how these attributions are calculated. A high-quality attributionbased explanation should meet several key criteria. First, it should demonstrate faithfulness, meaning that including features with high attribution should lead the model to the target output, and excluding them should prevent it. Second, it should exhibit specificity, ensuring that only critical features receive high attribution. Sometime this is also referred to as sparseness. Finally, it should be robust, meaning the explanation should remain consistent and not change significantly with minor variations in the input.

Attribution-based explanation has been studied for DRL models. [Greydanus et al., 2018] and [Puri et al., 2020] utilized policy distributional shifts as the basis for attribution in RL. Specifically, they calculate attribution of a feature as the difference in Q/V values or action distributions between the original and perturbed states. For example, given agent policy π, the attribution of a feature is proportional to Es′ (|π(s) -π(s′ )|2) where s stands for the original state and s′ represents perturbed states generated for this feature. By calculating the attributions for all features, a saliency map m can be created. However, the perturbation-based explanations lack faithfulness. Since each perturbation focuses only on local features while ignoring the joint impact of feature combinations, overlaying the saliency map with the original state, (m⊙ s), does not result in a feature combination that leads the agent to the target action distribution π(s).

A better approach to enhance faithfulness is to learn a model to predict the saliency map m that minimizes the difference between π(m⊙s) and π(s). Explainer [Stalder et al., 2022] leveraged this idea by training an explanation model for an image classifier. However, Explainer categorizes class labels into target and non-target for each training sample and focus on learning saliency map (or mask) only for the target label while treat all non-target labels as a single group. Unlike a (well trained) image classifier, where predictions for non-target labels are typically close to 0, DRL agent in many scenarios, does not exhibit a clear preference for the actions. Non-target actions may sometimes have probabilities only slightly lower than those of target actions. Analyzing how masking the feature may affect the non-target action probability provides additional information that can be used to train the explanation model more effectively.

The above analysis motivates us to design a trainable saliency map generator for attribution-based explanations and train it using two channels of contrastive information:(i) Action-wise contrast: We believe that environment states contain features that motivate the DRL agent to select both target action and non-target actions. However, the target action is ultimately chosen because it corresponds to higher reward or has a stronger presence. For each action, a saliency map can be generated as an explanation. Choosing features according to the saliency map for a non-target action should push the agent away from the target action, and vice versa. This inspired us to treat the saliency map of the non-target action ma′ and the target action ma as a negative pair, which can be leveraged for contrastive learning [Chopra et al., 2005; Schroff et al., 2015; Gutmann and Hyv¨arinen, 2010]. (ii) Feature-wise contrast: To exclude irrelevant features (e.g. background) from the saliency map, explanations also need to be discriminative in filtering out such information. When only irrelevant features are accessible to the agent, the resulting action distribution should be as uniform as possible. Therefore, the target action's saliency map (m) and its inverted counterpart ( ˜m = 1 -m) form another negative pair for contrastive learning.

In this work, we present VisionMask as an RL explainer that is contrastively trained to generate saliency maps to explain agent's actions. We specifically focus on agents that maps images to actions and consider each pixel value as the interpretable input feature, although the similar technique could be extended to other type of features. We carefully design the objective function to enable self-supervised contrastive learning of explanations from both action-wise and feature-wise perspectives, fostering the generation of more faithful explanations. We conduct evaluation on six RL environments with five baselines based on faithfulness, robustness, and sparseness. Quantitatively, VisionMask outperforms the baselines in terms of faithfulness, while exhibiting strong robustness and high sparseness. Qualitatively, we compared VisionMask with the baselines in two settings: visual comparison and human studies. In the visual comparison, VisionMask provides sharper explanations that align more closely with human interpretations, as demonstrated by counterfactual examples. In the human studies, VisionMask's explanations help users better understand the agent's decisions and calibrate appropriate trust.

The contribution of our work can be summarized as the following:

• We present VisionMask, a novel attribution method for explainable visual reinforcement learning that generates action-specific saliency maps. To the best of our knowledge, this is the first work that learns saliency map in a self-supervised contrastive learning manner, yielding highly faithful explanations across various RL environments.
• VisionMask is model-agnostic in that it works with any vision-based DRL agent and requires only the agent's input state and output action during inference time. VisionMask does not make any modification on the agent, however, each VisionMask model is trained specifically
• for a given agent by including the fixed agent in the back-propagation path for gradient descent.
• We conducted extensive experiments, both qualitatively and quantitatively, to demonstrate the improvements of the saliency map in terms of faithfulness, robustness, and sparseness, as well as a comprehensive ablation study for each component.

2 Background and RelatedWork

The nested structure and non-linear operation of DNN make it challenging for humans to understand how the outputs are derived from inputs. Some existing works address this challenge by explaining model outputs through input attributions [Petsiuk et al., 2018; Ribeiro et al., 2016; Fong and Vedaldi, 2017]. For example, Randomized Input Sampling for Explanation of Black-box Models (RISE) [Petsiuk et al., 2018] is a perturbation-based approach that explains neural image classifiers by applying randomly generated masks to the input image and assessing their impact. Local Interpretable Model-agnostic Explanations (LIME)[Ribeiro et al., 2016] tries to explain local instances by approximating them within a nearby vicinity using a linear decision model, where the explainable elements are superpixels (i.e., small image regions of similar pixels).

Built on top of the DNNs, the lack of transparency of the DRL agents undermines humans trust and hinders their adoption. Several studies have addressed this challenging problems. According to [Qing et al., 2022], existing efforts can be categorized into four main approaches: model explaining, reward explaining, task explaining, and state explaining.

Model explaining relies on inherently interpretable model architectures or auxiliary reasoning mechanisms to generate explanations [Topin et al., 2021]. However, these selfexplanatory or symbolic models often suffer from decreased performance compared to state-of-the-art neural network based RL policies and may lack the representational power needed to learn more complex policies. Reward explanation typically involves using explicitly designed explainable reward functions to generate explanations [Ashwood et al., 2022]. Task explanation considers a policy as a composition of sub-task and action sequences, explaining the behavior in terms of relationships among sub-tasks [Shu et al., 2017]. These approaches often assume the existence of a reward decomposition scheme or a predefined sub-task partition mechanism, which may not hold true in all RL environments.

State explanations are based on agent observations from the environment. This approach determines the significance of explainable elements within a state observation in relation to action selection or reward determination. An explainable element in the state could be a small region in the visual input or semantic features in the environment state. The proposed VisionMask falls into this category.

State explaining methods can further be divided into three categories: attention-based, Shapley value-based and perturbation-based mechanisms, which are detailed below. Attention is a common approach used for explainable RL in several existing works [Annasamy and Sycara, 2019]. However, similar to i-DQN [Annasamy and Sycara, 2019], these methods cannot explain a given DRL model because their attention model must be trained alongside the agent model. Shapley value [Shapley and others, 1953], a concept from game theory, has also been introduced into XRL. SVERL [Beechey et al., 2023] provides a theoretical analysis on how Shapley values can be applied to explain value and policy network in DRL. However, this approach is still in its early stages and assumes prior knowledge of the transition model of the environment, which is not realistic for realistic applications.

Perturbation-based methods [Greydanus et al., 2018; Iyer et al., 2018; Puri et al., 2020] compute saliency maps of the input features by comparing the action probability or value function before and after perturbation. [Greydanus et al., 2018] perturbed the state with Gaussian blur at fixed intervals, while [Iyer et al., 2018] identified objects and perturbed themwith a constant gray value. [Puri et al., 2020] improved on [Greydanus et al., 2018] by removing the effect of perturbations from other irrelevant actions. As discussed in section 1, the performance of these techniques is constrained by predefined perturbation rules and the lack of knowledge accumulation. Furthermore, there is no guarantee that the perturbed input remains physically meaningful.

3 VisionMask

In this section, we present our VisionMask architecture. The primary goal is to generate action-wise saliency maps that attribute the most relevant features in the state to each action. For agents that map images to actions, the features and states correspond to pixels and images.

3.1 ProblemFormulation

Formally, we define the environment as a Markov Decision Process (MDP) {S, A, P, R, γ}, where S represents the state space; A denotes the action space with |A| = K; the state transition function P : S × A → ∆(S) depicts the transition between states based on actions, where ∆(S) represents the set of probability distributions over states; the reward function R : S ×A → R provides the immediate reward for state-action pairs; γ ∈ [0, 1] is discount factor and π : S → ∆(A) represent policy. Return G is defined as G = ∑∞k=0 γ kRk+1, and the expected cumulative reward of a policy π is Eπ[G] = Eπ[∑∞k=0 γ kRk+1], where the expectation is taken with respect to the initial state distribution, transition probabilities, and action probabilities determined by π. VisionMask operates on a given trained expert policy πE such that πE ≈ π∗ = argmaxπ Eπ[G], where π∗ is the optimal policy. We can obtain a dataset of expert demonstrations DE = {(si, ai ∼ πE(si))}Ni=1 consisting of N stateaction pairs, from trajectories sampled while executing πE in the environment. Our goal is to learn an explainer fθ∗ that minimizes the loss θ∗ = argminθ ∑(s,a)∈DE L(a, s, θ) where L is the training loss function to be discussed in section 3.3. The explainer function fθ : S → [0, 1]K×ds , where ds represents the feature size of the state s ∈ S and K denotes the number of candidate actions, predicts the attributions of each action to each feature in the state s. The value of the output is bounded within the range [0, 1], with the (i, j)th element indicating the jth feature's attribution to the ith action.

In the case of visual DRL, s is the visual input of the agent and ds is the number of pixels in s, i.e., ds = W ×H, with W and H representing the width and height of the visual input. The output of the explainer can be viewed as K saliency maps, M = {m0,m1, ....,mK-1}, each corresponds to an action. A saliency value mi[x, y] near 1 indicates that the pixel (x, y) has significant contributions to action i, whereas a value near 0 denotes irrelevance. Overlaying the ith saliency map with the state s will highlight the input features that lead the agent to the ith action.

3.2 Architecture

As shown in Figure 1, we first collect the expert dataset DE using the expert policy πE. From this dataset DE, state-action pairs (si, ai) are sampled and fed to VisionMas fθ to generate the set of saliency maps M.

Generating the saliency map from given visual input is a dense prediction task that shares similarities with image segmentation, where each pixel is assigned a value to indicate whether it belongs to an object or background. Hence, we structure the explainer fθ akin to the widely used image segmentation model, DeepLabv3 [Chen et al., 2017], however, retrain it using self-supervised contrastive learning. To make sure that the output saliency value are bounded to the range [0, 1], a sigmoid function is applied at the output of fθ. For each mi ∈ M, we also calculate a complement map ˜mi = 1-mi highlighting the irrelevant regions for the action i. Then the masks mi and ˜mi are overlaid onto the original state s to generate two masked states si and ˜ si using the following overlay function:

where ⊙ is Hadamard Product and r is a reference value. Numerous options exist for the reference value, such as setting the pixel to zero, assigning a constant value, blurring the pixel, or cropping it. Empirical study shows that setting the reference to the background gives the best results. More details can be found in the ablation study 4.4. See Appendix F for the background values for different environments.

To generate self-supervised contrastive loss to train the model, we query the agent to obtain the corresponding logits zi = z(si) and ˜ zi = z(˜ si), where zi, ˜ zi ∈ RK , and the action probability distributions pi = softmax(zi) and ˜ pi = softmax(˜ zi), where p, pi ∈ [0, 1]K , 0 ≤ i < K. By concatenating each pi and ˜ pi, we have the the action probability distributions of each mask p, ˜p ∈ RK×K ,

3.3 Training Loss

To enable the agent contrastively learn the saliency map ma, we carefully designed the training loss function L as follows:

where La(s, a) is action-wise contrastive action loss, Lne(s, a) is feature-wise loss, Larea(s, a, n) is the area size loss and λne, λarea are regularization hyper-parameters.

Action-wise contrastive loss La. Let a denote the target action chosen by the agent. Our primary goal is to learn explanations faithful to the a, making (a, pa) the only positive pair. Furthermore, the explanation must be discriminative, meaning it should clearly distinguish the target action a from all other possible actions. As a result, every other pair of (a, pi̸=a) is treated as a negative pair. Then, we compute the cross-entropy loss between these pairs,

where zij is the logit for action j when querying the policy with the masked state si. The contrastive action loss La = Lpos + Lneg. Here �k = a� denotes the indicator function which returns 1 if k is the same as label a, and 0 otherwise. Note that we do not compute the loss with pk,i̸=a, as ensuring the faithfulness of the target action a is our primary objective here.

Feature-wise loss Lne. To ensure that the visual input regions selected by mi is necessary for the agent to make decisions, we also need to make sure that the unselected region, i.e., s ˜ mi , does not provide useful information for action selection, hence the action distribution ˜p should follow a uniform distribution. Motivated by this rationale, we define negative entropy loss regarding ˜m as the following:

Area size loss Larea. A low effort way to minimize La and Lne is to include all pixels in the mask, mi, and no pixel in the complement mask, ˜mi, which obviously is not a valid solution. We need to ensure that each importance mask only consists a small number of crucial pixels. Thus, we define Larea using L1 norm as follows:

where Z it the number of pixels in state.

4 Experiments

In this section, we begin by outlining the experimental setup. We then present quantitative and qualitative analyses to evaluate our approach. Additionally, we provide counterfactual explanations to demo VisionMask's faithfulness and sensitivity. Finally, we perform an ablation study to assess the contribution of each component.

4.1 Experimental Setup

Environment Selection. We conduct experiments across three types of environments: Super Mario Bros (SMB) [Kauten, 2018], Enduro, Seaquest and MsPacman [Bellemare et al., 2013; Machado et al., 2018] for 2D game, Highwayenv [Leurent, 2018] for autonomous driving simulation and VizDoom [Wydmuch et al., 2018] for 3D game.

Baseline Selection. We mainly compare our model with perturbation-based baselines for black-box RL such as Greydanus [Greydanus et al., 2018] and SARFA [Puri et al., 2020]. In addition, we also compared with three techniques originally designed to explain image classifiers, including a learning-based method, Explainer [Stalder et al., 2022], and two perturbation-based methods, LIME [Ribeiro et al., 2016] and RISE [Petsiuk et al., 2018]. Although these methods focus on image classification, their main concept is similar to ours: to attribute a given label (action) to a subset of visual features. We use the public implementation from torchray [Fong et al., 2019] for RISE and the original published implementations for other baselines. Among all the baselines, the Explainer is the most similar to VisionMask, as both are learning-based approaches. However, VisionMask employs action-contrastive learning and is trained using different regularization. In the experimental results, we demonstrate that these differences significantly enhances VisionMask's performance.

Expert Policy. We use the open-source pre-trained PPO [Schulman et al., 2017] agents from [Nguyen, 2020] for SMB environment and from stable-baselines3 [Raffin et al., 2021] for Enduro, Seaquest and MsPacman environments. We train the DQN [Mnih et al., 2013] agents from scratch for the Viz- Doom and Highway-env environments. See Appendix A for agents details.

Dataset. We collect state-action pairs for each environment and split the data into 80% for training, 10% for validation, and 10% for testing. All results in this section are reported using the test split. We make this dataset publicly available; see Appendix B for details.

4.2 Quantitative Analysis

Metrics. Since there is no ground truth explanations [Adebayo et al., 2020], it is crucial to select appropriate metrics to evaluate the trustworthiness of explanations. Designing robust and consistent metrics for XAI remains an unresolved challenge [Hedstr¨om et al., 2023]. To alleviate the inconsistency between metrics and avoid selection bias, we closely follow the LATEC [Klein et al., 2024] benchmark to evaluate performance across three dimensions: Faithfulness, Robustness, and Complexity.

Faithfulness measures the correlation between the agent's action and the masked visual input. In this context, we consider three metrics: Accuracy, Deletion and Insertion performance [Petsiuk et al., 2018]. The accuracy gives the percentage of time that masked input and the original input lead to the same action according to the expert agent. It is defined as the following: Accuracy = 1 |Dtest | ∑|Dtest | i=1 �πE(si) = πE(s′i )�, where s′i = si ⊙ mi + r ⊙ ˜mi, and �·� again denotes the indicator function.

Insertion and deletion measure the impact on the target action probability by progressively adding or removing pixels from the original input based on the descending order of their attribution scores, with the most important pixels being inserted/deleted first. The inserted/deleted input image will be processed by the expert agent again to obtain the probability of the target action. By plotting the probability against the number of pixels added or removed, we obtain an insertion and a deletion curve. The insertion/deletion performance is measured by the area under the curve (AUC) of the insertion/deletion curves. A larger (smaller) AUC for the insertion

(deletion) curve means that including (removing) the important pixels identified by the explainer can effectively increase (reduce) the probability of the target action. Hence a larger (smaller) AUC of insertion (deletion) curve indicates more accurate attribution prediction. The pseudo-code of the detailed information of deletion and insertion could be found in Appedix E.

For robustness, we report the Local Lipschitz Estimate (LLE) scores [Alvarez Melis and Jaakkola, 2018], which quantify the local smoothness of explanations by estimating the Lipschitz constant within a specific neighborhood. The Lipschitz constant measures the maximum rate of change of the function, ensuring explanation does not vary too rapidly within the state's neighborhood. Given state si and neighborhood size ϵ, LLE defined as

We evaluate complexity with Sparseness [Chalasani et al., 2020], which uses the Gini index on the vector of absolute attribution values sorted in non-descending order. Sparseness ensures that features genuinely influencing the output have substantial contributions, while insignificant or only slightly relevant features should have minimal contributions. A higher Sparseness indicates more contrastive attribution values and hence more understandable explanations [Chalasani et al., 2020].

Results. In Table 1, our VisionMask achieves the best performance in terms of faithfulness (i.e., Acc, Del, Ins) in all testing environment except SMB, where its deletion score is slightly lower than RISE. However, RISE has significantly lower insertion score and accuracy in this environment. Hence, our explanations are more aligned with the agent's decision-making process. Moreover, compared to, Explainer [Stalder et al., 2022], VisionMask exhibits much higher faithfulness, which suggests the effectiveness of action contrastive learning. Overall, learning-based model performs better compared to perturbation based approach.

In terms of robustness, RISE achieves the best performance in terms of LLE score across all settings. This is because, as a perturbation-based method, RISE uses the same type of perturbation that we used to generate the neighbor image for LLE score calculation. Hence the perturbation almost has no impact to it. On the other hand, RISE has the lowest sparsity, which means the attribution predicted by RISE are evenly distributed and contains little information. For Complexity, benefiting from the binary mask of LIME which is much more sparse compared to other baselines, LIME has the highest sparseness score. However, it has the worst performance in terms of LLE score. Overall, VisionMask achieves the best balance between the robustness and sparseness. See Appendix for the Radar map.

4.3 Qualitative Analysis

In this section, we conduct two qualitative analyses across two settings, visual comparision and human studies.

Visual Comparision. We present example explanations from three environments, SMB, Enduro, and Seaquest, in Figure 2, along with some counterfactual analysis generated from the explanations provided by VisionMask. The examples show that both LIME and RISE fail to generate interpretable explanations. LIME's superpixels are too large to capture the specific regions, while RISE's explanations include almost all pixels. Explainer generates more accurate and interpretable explanations compared to LIME and RISE. This suggests that purely perturbation-based approaches may fail in RL due to the high dynamics of environments and the lack of learning ability. [Greydanus et al., 2018] and SARFA generate better explanations than LIME and RISE

but often focus on irrelavent objects or background. In contrast, VisionMask accurately highlights the relevant regions, providing sharp explanations that are both accurate and interpretable. By removing some of the regions highlighted by the VisionMask, counterfactual analysis could be performed to answer questions like: "Why this action instead of that one?"

Human Studies. To assess whether the saliency maps improve humans' understanding of the agent's decisions, we conduct human studies for SMB, Enduro and Seaquest. We present 63 participants with 9 state-action pairs accompanied by the saliency maps generated by VisionMask. Participants are then asked whether the saliency maps help them better understand the decision-making process of the RL agents. We record the total number of responses as well as the number of positive responses, as shown in Figure 3. Over 89% of participants find the saliency maps generated by VisionMask helpful in understanding the agents' decisions. Details about

• (d) Reference value. Using background produce better faithfulness values.
• (e) Component. Negative Entropy (NE) with L1 area regularization performs better in terms of faithufulness. Contrastive-only (C.), Stalder Area Loss (S.), and Total Variation (TV.)

Table 2: VisionMask ablation experiments on SMB. Default settings are marked in gray .

this study can be found in the Appendix D.

4.4 Ablation Studies

Table 2 shows the results for ablation studies. In all tables, the column marked in gray represents the default setting of VisionMask used to generate results in previous sections.

Area size regularization. In this experiment, we replace the default L1 area regularization with the Min-Max Area Loss from Explainer [Stalder et al., 2022]. Without this constraint, Explainer tends to produce masks where all values are one in order to achieve the best performance. Explainer first vectorizes and sorts the mask elements, and then applies penalties to pixels falling outside the range [a, b], where the hyperparameters a and b specify the minimum and maximum allowable area sizes. Instead, VisionMask uses L1 area regularization to constrain the maximal area, as we do not wish to impose constraints on the minimum area size. To implement action contrastive learning, we need to generate masks for actions that are not selected, and these masks should be allowed to be all zeros if necessary. As shown in Table 2a, using L1 area regularization improves the faithfulness of the explanation and increases training speed.

Max area size. We further vary the maximum permitted area size amax to evaluate its impact. As shown in Table 2c, reducing amax significantly increases the sparseness as it constrains the model to focus on the most discriminative regions without compromising other performance metrics. The 0.1 setting achieves the best results, as expected.

fθ model. In Table 2b, we compare the performance of three different versions of VisionMask with three different segmentation models: our default DeepLabv3 [Chen et al., 2017], Fully Convolutional Network (FCN)[Long et al.,

2015], and UNet[Ronneberger et al., 2015]. The results show that DeepLabv3 achieves a better performance.

Reference Value. We evaluate the impact of hyperparameter r in Equation 1 by setting the reference value to background, black, mean, and blur. The "black" reference sets r to 0, the "mean" reference sets r to the RGB mean value of the entire dataset, and the "blur" reference applies Gaussian blur (kernel is 39, σ = 15) to the image and use the blurred pixel value as r. Table 2d compares the performance of VisionMask with these four different reference values. We choose the background as reference value as it provides a more balanced performance across all metrics. Details about the background reference value are provided in Appendix F.

Loss Function Components. In Table 2e, we evaluate the impact of different components in the loss function on the performance of VisionMask. The results show that, although adding negative entropy and L1 area regularization slightly reduces accuracy, it improves the insertion and deletion scores and significantly boosts the Sp./LLE ratio. In contrast, adding total variation to produce smoother masks has minimal impact on overall performance.

5 Conclusion

We presented VisionMask, an agent-agnostic DRL explanation model trained in self-supervised contrastive learning. VisionMask generates explanations with higher fidelity and better effectiveness compared to existing attribute-based methods. It is our future plan to extend this approach to multimodality input and couple the visual explanation generated by the VisionMask with other information such as agent's longterm goals and future rewards.

Acknowledgments

This research is partially supported by the Air Force Office of Scientific Research (AFOSR), under contract FA9550-24-10078, and NSF award CNS-2148253. The paper was received and approved for public release by Air Force Research Laboratory (AFRL) on May 28th 2024, case number AFRL-20242908 . Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of AFRL or its contractors.