SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Devleena Das
School of Interactive Computing
ddas41@gatech.edu
Sonia Chernova
chernova@gatech.edu
beenkim@google.com
As more non-AI experts use complex AI systems for daily tasks, there has been an increasing effort to develop methods that produce explanations of AI decision making that are understandable by non-AI experts. Towards this effort, leveraging higher-level concepts and producing concept-based explanations have become a popular method. Most concept-based explanations have been developed for classification techniques, and we posit that the few existing methods for sequential decision making are limited in scope. In this work, we first contribute a desiderata for defining "concepts" in sequential decision making settings. Additionally, inspired by the Protégé Effect which states explaining knowledge often reinforces one's self-learning, we explore how concept-based explanations of an RL agent's decision making can in turn improve the agent's learning rate, as well as improve end-user understanding of the agent's decision making. To this end, we contribute a unified framework, State2Explanation (S2E), that involves learning a joint embedding model between state-action pairs and concept-based explanations, and leveraging such learned model to both (1) inform reward shaping during an agent's training, and (2) provide explanations to end-users at deployment for improved task performance. Our experimental validations, in Connect 4 and Lunar Lander, demonstrate the success of S2E in providing a dual-benefit, successfully informing reward shaping and improving agent learning rate, as well as significantly improving end user task performance at deployment time.
Black-box AI systems are increasingly being deployed to help end-users with everyday tasks. Examples include doctors leveraging decision support systems to aid in diagnosis [39], warehouse managers relying on robots for goods transportation [8], and drivers using autonomous vehicles for assisted driving [18]. To increase the transparency of these black-box models, researchers have developed numerous techniques to provide explanations of agent decision making [45, 2, 23, 51, 14, 9, 30, 15].
A popular method towards non-expert friendly explanations has been to attribute higher-level "concepts" to an agent's decision making, and these concepts have primarily been used to explain classification-based AI systems [30, 33, 54, 19]. An example of a concept-based explanation for a classification model includes, "wing color" for a bird class label [19]. In sequential decision making, concept-based explanations have been less widely explored, and existing works leverage preconditions of states and action costs [48] or logical formulas [22] as "concepts".
In this work, we posit that concept-based explanations may not only benefit the end-user but also benefit the AI agent. Our claim is loosely motivated by the Psychological phenomenon, the Protégé Effect, that states explaining material to another student also helps the explainer learn and reinforce material [7, 24, 17]. Thus the objective of our work is to demonstrate the dual-benefit of concept-based explanations to both the end-user for improved understanding, as well as the AI agent for improved learning rate. Note, prior works, typically in Reinforcement Learning (RL), have explored various methods to improve agent learning, a common approach including the use of human feedback and natural language commands as a method of reward shaping to improve RL agent sample efficiency and learning rate [38, 50, 20]. However, these methods have focused on leveraging language as a one-way benefit to the agent, not considering how language provided to the agent during training may also benefit end-user understanding at deployment. Therefore in our work, we explore the following research question: Can concept-based explanations have a two-way benefit to both the user and agent, such as to improve end-user understanding of a deployed agent's decision making, while also improving the agent's learning rate via reward shaping?" As a solution, we contribute a unified framework, State2Explanation (S2E), that learns a single joint embedding model between agent states and associated concept-based explanations and leverages such learned model to (1) inform reward shaping for agent learning, and (2) improve end-user understanding of the agent's decision making at deployment.
Additionally, we contribute a desiderata of what entails a "concept" in a concept-based explanation for sequential decision making systems. Given that sequential decision making agents engage in long-term interaction with their environment, we posit that the scope of concept-based explanations should span beyond representing preconditions and action cost [48], and control logic [22]. Specifically, we believe concept-based explanations in sequential decision making should function at a much higher-level of abstraction, highlighting knowledge that are applicable across multiple states, and most importantly, expressing a positive or negative influence towards the agent's goal. Our claim is motivated by the importance of the agent's goal in a sequential decision making formulation [44, 49].
Contributions. (1) We provide desiderata for what constitutes a "concept" in concept-based explanations of sequential decision making. (2) We introduce a novel framework, State2Explanation (S2E) which learns a joint embedding model between agent states-action pairs and concept-based explanations to provide a benefit to both the agent, by informing reward shaping, as well as the end user by providing explanations that improve user task performance. We perform both model and user evaluations of our S2E framework in two complex RL domains, Connect 4 and Lunar Lander.
Our model evaluations demonstrate that S2E can successfully inform reward shaping, comparable to expert-defined reward functions. Additionally, our user evaluations demonstrate that S2E can successfully provide meaningful concept-based explanations to end users such that exposure to our explanations significantly improve user task performance.
In this work, we explore the utility of concept-based explanations in providing a two way benefit to both the AI agent during training, via reward shaping, as well as end user at deployment time, via improved user task performance. Below, we highlight related prior work.
Reward Shaping using Language Many prior works have focused on improving an agent's learning rate or sample efficiency [35, 21]. Closely related to our work are methods providing reward shaping via language. Specifically, [38] propose ELLA, which leverages a relevance classifier to associate a lower level instruction to a higher level task for providing a subtask shaping reward. In [20], the authors leverage natural language instruction to derive potential-based shaping rewards that encourage an RL agent to more frequently select relevant actions in the game of Montezuma's Revenge. Similarly, [50] learn a joint embedding model between subgoal language commands and agent states to provide shaping rewards for completing subtasks in StarCraft II. Specifically, [50] demonstrate that their model can generalize to unseen, semantically similar language commands, and improve agent learning rate. In our work, we take inspiration from [50], and learn a joint embedding model between agent states and associated natural language commands to inform agent reward shaping. However, our work differs in that our natural language commands are concept-based explanations of the associated state, and not subgoal descriptions required for task completion. By learning a common embedding space between state explanations and state representations, our work
provides a two-way benefit to both the RL agent for learning as well as end-user for improved understanding.
Explanations for non-AI Experts Prior works have established that concept-based explanations for classification tasks provide explanations that are meaningful to end-users [30, 33, 54, 19]. Methods for concept-based explanations have included Concept Activation Vectors [30], Concept Bottleneck Models [33], and Concept Embedding Models [54]. For classification tasks, "concepts" have primarily represented human-interpretable features that a model's prediction is sensitive to (i.e. 'wing color' for image classification.) [33]. Related to sequential decision making, [26] provide concept-based explanations of 3D action recognition CovNets by clustering learned, human-interpretable features (i.e. "grass" or "hand") from segmented videos. In [48], concept-based explanations for sequential decision making are formulated in relation to state preconditions (e.g. "the action move-left failed in the state as the precondition skull-not-on-left was false in the state") and action costs (e.g., "executing the action attack in the presence of the concept skull-on-left, will cost at least  $500$ ). Specifically, in [48], concepts have represented any factual statement a user associates to a given state, such as "skull-on-left" in the previous example. Similarly, in [22], concepts have represented logical formulas for summarizing RL policies. Additionally, [29] propose a prototype wrapper network for training interpretable RL policies with user-defined prototypes. Similar to [48], the authors in [29] define prototypes as factual observations about a state, such as "turn right" or "turn left". Other explanation generation methods have demonstrated that end user understanding is improved when explanations are presented in natural language rather than features or labels [14, 5, 9, 10]. Inspired by these findings, our work contributes the first method for producing natural language concept-based explanations, where we define a general desiderata of concepts that goes beyond action costs and preconditions to capture goal-driven decision making.
3 Concept Based Explanations for Sequential Decision Making
We posit that the objective of a concept-based explanation, in sequential decision making settings, is to communicate higher-level knowledge that aids the agent and/or user in reasoning about the utility of taking a given action from a specific state. In this section, we provide a general definition of concept-based explanations and "concepts", in sequential decision making, that extends beyond prior usage of state preconditions and action-costs [48], and control logic [22]. Specifically, we define a concept-based explanation,  $\mathcal{E}_{C_i}^i$ , as one that presents, in natural language form, the set of concepts,  $C_i = \{c_j\dots c_m\}$ , that can describe a particular state-action pair  $(s_i,a_i)$  in the context of goal  $G$ . This naturally poses the question of: "What is a concept  $c_{j}$ ?" Below are several desiderata for defining a concept in sequential decision making settings:
1. A concept should be grounded in human domain knowledge. A concept,  $c_{j}$ , should be contextualized in the overall task domain. For example, although an agent's goal may be to optimize its Q-function, we posit that "higher Q-value" is not a valid concept. Instead, a concept should represent a higher-level abstraction of what the "higher Q-value" relates to in the domain. For example, in Chess, an agent receiving a "higher Q-value" for a particular state-action pair may be translated by a domain expert to "capturing a queen".
2. A concept should relate to the task goal. A concept,  $c_{j}$ , is an abstracted representation of  $(s_{i}, a_{i})$  that encompasses how  $(s_{i}, a_{i})$  may lead to or inhibit the agent's goal  $G$ . This claim is motivated by the general objective of a sequential decision making agent, that is to perform a sequence of actions  $\langle a_{1}, a_{2}, \dots, a_{n} \rangle$ ,  $a_{i} \in A$ , defined via a plan or policy  $\pi$ , that transforms the agent's current state  $I$  to its goal state  $G$  [44, 49]. For example, we do not consider the description "Knight in A4" as a valid concept, since, by itself, it does not provide relation to the game's objective. However, we would consider "king safety" as a valid concept, since it relates to game's objective of checkmate.
3. A concept should be generalizable. A concept,  $c_{j}$ , should generalize across multiple  $(s_{i}, a_{i})$  pairs. This means that a concept is typically not a unique description of a single  $(s_{i}, a_{i})$ . For example, a valid concept in Chess is "checkmate" since the abstraction of "checkmate" can generalize to many different  $(s_{i}, a_{i})$ . An invalid concept involves detailing the entire game board by piece position.
Note, above we define a concept,  $c_{j} \in C_{i}$ , in the context of describing a single state-action pair,  $(s_{i},a_{i})$ . However, in many domains a meaningful concept  $c_{j}$  can only be defined in terms of a
sequence of state-action pairs. This is especially true in domains in which state-action pairs are sampled at high-frequency, such as robotics and video games [40, 56]. For example, concepts such as "rotate wrist" or "bend arm" can only be defined by a series of state-action pairs, as opposed to a single state-action pair. As a result, we define concept-based explanation in these domains as  $\mathcal{E}_{C_i}^{[i,i + n]}$  which presents the set of concepts  $C_i = \{c_j..c_m\}$ , for a series of  $(s_{[i,i + n]},a_{[i,i + n]})$  pairs.
Finally, we highlight that some  $(s_i,a_i)$  or  $(s_{[i,i + n]},a_{[i,i + n]})$  pairs may not have a concept-based explanation associated with them. In other words,  $C_i = \emptyset$  for some  $(s_i,a_i)$  pair or  $(s_{[i,i + n]},a_{[i,i + n]})$  pairs. This is especially true in long horizon tasks, in which a  $(s_i,a_i)$  may not relate directly to  $G$  but is important for another, future  $(s_{i + n},a_{i + n})$  pair which then leads to  $G$ . For example, "Move the pawn to D5 so that two moves from now you may capture the rook". Our definition of concept-based explanations do not consider these cases, and is an opportunity for future work.
(c)
Figure 1: Our S2E framework involves (a) learning a joint embedding model  $M$  from which a  $\mathcal{E}_{C_i}^i$  is extracted and utilized (b) during agent training to inform reward shaping and benefit agent learning, and (c) at deployment to provide end-users with  $\mathcal{E}_{C_i}^i$  for agent actions.
4 State2Explanation (S2E) Framework
Figure 1 introduces our State2Explanation (S2E) framework which includes a three-step process to provide a benefit to both the RL agent and the end-user. Specifically, S2E learns a joint embedding model  $M$  to align agent state-action pair  $(s_i, a_i)$  to concept-based explanations  $\mathcal{E}_{C_i}^i$  (Fig. 1a). S2E then leverages the learned  $M$  to inform reward shaping during agent training (Fig. 1b). Finally, S2E leverages the learned  $M$  to provide  $\mathcal{E}_{C_i}^i$  about the trained agent's decision making to end-users at deployment (Fig. 1c). Below, we further detail S2E.
4.1 Joint Embedding Model  $M$
A core component of S2E is learning a joint embedding model,  $M$ , to align an agent state-action pair,  $(s_i, a_i)$  with an associated concept-based explanation  $\mathcal{E}_{C_i}^i$ . The motivation for learning  $M$  is to explicitly map an agent's representation,  $s_i$ , into a representation understandable by our non-AI expert users,  $\mathcal{E}_{C_i}^i$ . Joint embedding models have been leveraged in many applications, such as in image captioning [55, 34], knowledge-graph generation [28, 43], and RL [42, 41, 50] to improve task learning. In the context of S2E, we demonstrate that joint embedding  $M$  can provide the dual benefit to i) agent learning, and ii) improved end user understanding of agent actions. Our  $M$  is inspired by [50] who learn an embedding space in the context of hierarchical RL tasks. However, unlike in [50], our  $M$  is not constrained to hierarchical RL tasks, and we consider domains in which multiple states can be associated to multiple, non-unique concepts.
The input to our joint embedding model,  $M$ , is defined by  $\langle [s_i,s_{i - 1}],g_i,\mathcal{E}_{C_i}^i,y_i\rangle$ , which includes an agent's current state  $s_i$ , previous state  $s_{i - 1}$ , other contextual game information (i.e., player in a multiplayer game, whether game is over)  $g_{i}$ , associated concept-based explanation  $\mathcal{E}_{C_i}^i$ , and  $y_{i}$  which defines if  $s_i$  and  $\mathcal{E}_{C_i}^i$  are aligned or misaligned. Note, including  $s_i$  and  $s_{i - 1}$  implicitly encodes the
action  $a_i$  taken to transition from  $s_{i-1}$  to  $s_i$ . Our contrastive loss function (Eq. 1) is adapted from [50], in which  $\theta$  represents the model's parameters, and  $i \in N$  represents the  $i^{th}$  training sample. Specifically,  $M$  minimizes the L2-norm of the difference between the embedding vectors when  $s_{\text{embed}}^i$  and  $\exp_{\text{embed}}^i$  are aligned ( $y = 0$ ) and maximize the L2-norm of the difference when  $s_{\text{embed}}^i$  and  $\exp_{\text{embed}}^i$  are misaligned ( $y = 1$ ).
$$
L (\theta) = \frac {1}{N} \sum_ {i} ^ {N} \left(\left\| s _ {\text {e m b e d}} ^ {i} - \mathcal {E} _ {\text {e m b e d}} ^ {i} \right\| - y\right) ^ {2} \tag {1}
$$
To train  $M$ , we leverage a dataset  $D = \{D^a, D^m\}$ , in which  $D^a$  and  $D^m$  denote the set of aligned and misaligned data samples, respectively. To collect  $D^a$ , we leverage domain knowledge to annotate relevant  $s_i$  with its aligned concept set  $C_i$ . The concept set  $C_i$  is then produced into a natural language explanation,  $\mathcal{E}_{C_i}^i$ , manually via fixed, expert-defined templates. That is, each  $C_i$  has one templated  $\mathcal{E}_{C_i}^i$ . The list of concepts for our evaluation domains are in Section 5.1. To collect  $D^m$ , each  $s_i$  is paired with  $z$  concepts not present in concept set  $C_i$ . Section 5.2 provides more details on  $D$ .
Explanation Retrieval Given  $M$ , for any  $(s_i, a_i)$ , we extract the closest explanation embedding  $exp_{\mathrm{embed}}^i$  for corresponding  $s_{\mathrm{embed}}^i$ , and decode  $exp_{\mathrm{embed}}^i$  to an  $\mathcal{E}_{C_i}^i$  similar to image-to-text retrieval [31, 6]. Specifically, we rank the set of possible explanation embeddings,  $\{exp_{\mathrm{embed}}^i..exp_{\mathrm{embed}}^k\}$ , by their L2-norm distance to  $s_i$ . The decoding of the best ranked  $exp_{\mathrm{embed}}^i$  to  $\mathcal{E}_{C_i}^i$  is then a vocabulary look-up. It is possible to utilize large language models (LLMs), and generate  $\mathcal{E}_{C_i}^i$  from  $exp_{\mathrm{embed}}^i$  which we consider future work. Instead, in this work, we focus on validating the dual utility of embedding-based  $\mathcal{E}_{C_i}^i$  retrieval to both the RL agent as well as end-user.
4.2 Reward Shaping in RL via Joint Embedding Model  $M$
During agent learning, we utilize  $M$  to retrieve the closest  $\mathcal{E}_{C_i}^i$  associated with a  $(s_i, a_i)$  for informing reward shaping. Prior work has shown that reward shaping can improve agent learning rate and sample efficiency [11, 25, 32]. However, designing an effective dense reward function is not trivial [16, 12], often requiring a two-stage trial-and-error procedure to 1) determine when a state-action pair deserves an intermediate reward, and 2) determine what reward value such state-action pair should be attributed. Many methods, such as inverse reinforcement learning, utilize learned, black box reward functions to minimize the trial-and-error process [37, 1, 57]. However, with black-box reward functions, we lose transparency of why a  $(s_i, a_i)$  is rewarded. To balance the design-transparency trade-off of dense reward functions, we leverage  $M$  to inform when a  $(s_i, a_i)$  pair should be rewarded based on the retrieved  $\mathcal{E}_{C_i}^i$ , and as a result, remove the first step of the trial-and-error process. Specifically, if  $M$  returns a meaningful  $\mathcal{E}_{C_i}^i$  for a  $(s_i, a_i)$  (e.g., concept set  $C_i \neq \emptyset$ ), then such  $(s_i, a_i)$  represents concepts that influence the agent's goal, and should receive an intermediate reward.
Note,  $M$  does not remove the trial-and-error process needed to determine what values each reward component should receive. Instead, we perform a hyperparameter sweep to assign shaping values to each concept component when an expert-determined, continuous shaping function is not derivable. Thus, the "Explanation-to-reward-shaping" module in S2E performs a look-up to assign a reward value based on the retrieved  $\mathcal{E}_{C_i}^i$ . In Section 5.3 we validate how leveraging  $M$  to inform reward shaping can improve an agent's learning rate comparable to expert-defined shaping functions.
4.3 Concept-Based Explanations to End Users via Joint Embedding Model  $M$
Once an RL agent is trained and deployed, we leverage  $M$  to provide  $\mathcal{E}_{C_i}^i$  about an agent's decision-making to non AI expert end-users. Recall from Section 3 that  $\mathcal{E}_{C_i}^i$  can include multiple concepts (i.e.  $|C_i| > 1$ ). We posit that in some cases,  $\mathcal{E}_{C_i}^i$  with too many concepts may become ineffective in aiding end-user understanding. In fact, prior works discuss the challenge of presenting information at the "right" level of abstraction to end-users [46, 53]. To the best of our knowledge, there are no explicit guidelines for what information abstractions are beneficial in sequential decision making settings. From analyzing prior end-user friendly explanation techniques [9, 14], we infer two important dimensions to consider for explanation abstraction in sequential decision making: Information Filtering (InF), and Temporal Grouping (TeG).
Information Filtering (InF) InF filters out concepts within  $\mathcal{E}_{C_i}^i$  to only include those that immediately influence the goal. For example, consider a  $\mathcal{E}_{C_i}^i$  in Lunar Lander, "Fire left engine because it brings lander closer to the center, decreases lander velocity to avoid crashing, and decreases tilt". Applying InF may produce a filtered explanation, "Fire left engine because it brings lander closer to the center and decreases tilt" given a  $(s_i, a_i)$  when the lander needs to immediately correct its tilt and
Connect 4 Concepts	Description
BW	Blocks an opponent win
CD	Provides center dominance
3IR	Creates a three-in-a-row
3IR_BL	Creates a three-in-a-row but blocked from a win
W	A four-in-a-row win
NULL (C_i = ∅)	Any neutral state that advances the game
$\varepsilon_{\text{concept}}$ : Play column 4 because it provides center dominance and creates a three-in-a-row.
Lunar Lander Concepts	Description
POS	Brings lander closer to center
VEL	Decrease lander speed to avoid crashing
TILT	Decrease tilt
RLEG	Encourage right leg contact with ground
LLEG	Encourage left leg contact with ground
MF	Conserve main fuel usage
SF	Conserve side fuel usage
L	Land without crashing
$\varepsilon_{\mathrm{concept}}$  : Fire left engine because it moves lander closer to the center, decreases lander speed to avoid crashing, decreases tilt of lander, and conserves main fuel usage.
Figure 2: The tables present concepts for Connect 4, derived from [3], and Lunar Lander, derived from [4], with accompanying example  $\mathcal{E}_{C_i}^i$ .
position to avoid a trajectory path leading to a crash. While the other concepts are important, they are filtered out since they do not effect the immediate chance of meeting  $G$ .