Pay Attention to What and Where? Interpretable Feature Extractor in Vision-based Deep Reinforcement Learning

TL;DR

This work tackles the spatial misalignment problem in attention-based interpretability for vision-based DRL by introducing the Interpretable Feature Extractor (IFE), which combines a Human-Understandable Encoding (HUE) that preserves spatial information with an Agent-Friendly Encoding (AFE) that maintains learning efficiency. The single, interpretable attention mask produced by HUE, together with the learned agent-ready features from AFE, enables clear visualization of both what and where the agent focuses. Evaluated on 57 ATARI games within the Rainbow framework and extended to A3C-LSTM, IFE demonstrates superior spatial preservation and interpretability relative to CNN-based attention and the S3TA baseline, while maintaining competitive data efficiency. The results suggest that IFE provides reliable, transferable explanations for vision-based DRL and can serve as a practical, modular component for interpretable reinforcement learning systems.

Abstract

Current approaches in Explainable Deep Reinforcement Learning have limitations in which the attention mask has a displacement with the objects in visual input. This work addresses a spatial problem within traditional Convolutional Neural Networks (CNNs). We propose the Interpretable Feature Extractor (IFE) architecture, aimed at generating an accurate attention mask to illustrate both "what" and "where" the agent concentrates on in the spatial domain. Our design incorporates a Human-Understandable Encoding module to generate a fully interpretable attention mask, followed by an Agent-Friendly Encoding module to enhance the agent's learning efficiency. These two components together form the Interpretable Feature Extractor for vision-based deep reinforcement learning to enable the model's interpretability. The resulting attention mask is consistent, highly understandable by humans, accurate in spatial dimension, and effectively highlights important objects or locations in visual input. The Interpretable Feature Extractor is integrated into the Fast and Data-efficient Rainbow framework, and evaluated on 57 ATARI games to show the effectiveness of the proposed approach on Spatial Preservation, Interpretability, and Data-efficiency. Finally, we showcase the versatility of our approach by incorporating the IFE into the Asynchronous Advantage Actor-Critic Model.

Figures (10)

• Figure 1: Proposed Interpretable Feature Extractor Architecture for Vision-based Deep Reinforcement Learning
• Figure 2: Illustration of overlapping convolutional operations result in one-to-many transformation.
• Figure 3: Network Architectures used in the paper
• Figure 4: Visualization of the attention mask overlayed on visual input
• Figure 5: Example of 5 consecutive inputs of the model with the attention mask overlay. The figures show the comparison between the proposed model with the traditional CNN model in (a) Pong and (b) Enduro. The attention visualization of the proposed model is much clearer and more interpretable, while that of the CNN approach is distorted and blurred by multiple attention masks for the object. We can also see the attention consistent with the movement of the objects such as the ball or the car.