[Paper Review] ProtoX: Explaining a Reinforcement Learning Agent via Prototyping

Ronilo J. Ragodos, Tong Wang+2 more|arXiv (Cornell University)|Nov 6, 2022

Explainable Artificial Intelligence (XAI)4 citations

TL;DR

ProtoX is a prototype-based post-hoc explainer for black-box reinforcement learning agents that identifies representative behavioral patterns (prototypes) to explain actions via visual and scenario similarity. It uses self-supervised contrastive learning for visual similarity and an isometry layer to capture contextual scenario similarity, achieving high fidelity to expert policies while providing interpretable, human-understandable explanations grounded in behavioral patterns.

ABSTRACT

While deep reinforcement learning has proven to be successful in solving control tasks, the "black-box" nature of an agent has received increasing concerns. We propose a prototype-based post-hoc policy explainer, ProtoX, that explains a blackbox agent by prototyping the agent's behaviors into scenarios, each represented by a prototypical state. When learning prototypes, ProtoX considers both visual similarity and scenario similarity. The latter is unique to the reinforcement learning context, since it explains why the same action is taken in visually different states. To teach ProtoX about visual similarity, we pre-train an encoder using contrastive learning via self-supervised learning to recognize states as similar if they occur close together in time and receive the same action from the black-box agent. We then add an isometry layer to allow ProtoX to adapt scenario similarity to the downstream task. ProtoX is trained via imitation learning using behavior cloning, and thus requires no access to the environment or agent. In addition to explanation fidelity, we design different prototype shaping terms in the objective function to encourage better interpretability. We conduct various experiments to test ProtoX. Results show that ProtoX achieved high fidelity to the original black-box agent while providing meaningful and understandable explanations.

Motivation & Objective

• To address the interpretability challenge of black-box deep reinforcement learning agents in unstructured environments like images and text.
• To develop a post-hoc explanation method that provides human-understandable insights into *why* an agent takes a specific action in a given state.
• To overcome limitations of attention maps and decision trees by using prototypes that represent salient behavioral patterns across visually diverse but contextually similar states.
• To jointly model visual similarity (via self-supervised contrastive learning) and scenario similarity (via isometry layer) for more robust and meaningful explanations.
• To enable explanation fidelity and interpretability without requiring access to the agent’s policy or environment, using only demonstration trajectories.

Proposed method

• Pre-trains a convolutional encoder using a time-contrastive loss on expert demonstration trajectories to learn visual similarity based on temporal proximity and shared actions.
• Uses a siamese network architecture with contrastive learning to cluster visually similar states in the embedding space.
• Introduces an isometry layer during behavior cloning to adapt the feature space to capture scenario similarity, enabling the model to recognize that visually different states may share the same underlying intent.
• Trains ProtoX via imitation learning (behavior cloning) on expert trajectories, using a loss function that includes prototype shaping terms to improve interpretability and reduce redundancy.
• Computes action evidence as a weighted sum of similarities between input states and prototypes, where weights reflect action connections (how strongly a prototype supports a specific action).
• Generates importance maps via patch masking to highlight which visual regions contribute most to similarity, aiding in diagnosing agent behavior.

Experimental results

Research questions

• RQ1 Can a prototype-based method effectively explain black-box DRL agents in unstructured observation spaces like images?
• RQ2 How can visual similarity and scenario similarity be jointly modeled to improve explanation fidelity and interpretability?
• RQ3 To what extent can self-supervised pre-training on expert demonstrations replace environment access for training a policy explainer?
• RQ4 Can ProtoX detect and explain flawed agent behaviors, such as failing to jump to avoid enemies, by identifying misaligned prototypes?
• RQ5 How do prototype representations compare to attention-based or tree-based explanation methods in terms of human interpretability and fidelity?

Key findings

• ProtoX achieved high explanation fidelity, matching or exceeding that of VIPER and GAIfO, two strong baselines with similar network architectures.
• The model successfully explained complex behaviors in Super Mario Bros. and Seaquest by identifying prototypical scenarios such as jumping over pipes, avoiding Goombas, and anticipating enemy spawns.
• ProtoX diagnosed the bad agent’s failure to jump by revealing that it relied on sky color similarity, with importance maps showing only the sky as salient—indicating a flawed reasoning pattern.
• In contrast, ProtoX trained on the good agent correctly identified strategic elements like hitting bricks to kill enemies, with importance maps highlighting Mario’s position and action context.
• The addition of the isometry layer significantly improved scenario similarity modeling, allowing the model to recognize that visually dissimilar states could represent the same behavioral intent.
• Prototype shaping terms in the objective function reduced redundancy and improved interpretability, resulting in a more concise and meaningful set of behavioral prototypes.