Source-linked AI summary
Towards Interpretable Reinforcement Learning Using Attention Augmented Agents
Alex Mott, Daniel Zoran, Mike Chrzanowski, Daan Wierstra, Danilo J. Rezende
TL;DR
The paper tackles the difficulty of interpreting reinforcement-learning agents by introducing a soft, top-down attention mechanism that sequentially queries visual inputs. The resulting attention maps expose task-relevant information, while the agent remains competitive with state-of-the-art Atari agents and supports “what” versus “where” analysis.
Problem
Traditional RL agents do not easily reveal what information drives their decisions or why mistakes occur.
Method
The model uses differentiable top-down attention to query visual inputs, compress the selected information, and feed query-dependent summaries into a recurrent agent.
Results
The agent achieves performance competitive with state-of-the-art agents across a broad range of Atari levels while its attention maps visualize attended input regions.
Takeaways & Limitations
Attention maps support more comprehensive analysis of policy information, including repeated strategies, key features, look-ahead trajectories, and behavioral tripwires.
Takeaways & Limitations
The convolutional vision core has limited ability to produce absolute spatial-position information, motivating the spatial basis.
Abstract
from arXiv · showhide
Inspired by recent work in attention models for image captioning and question answering, we present a soft attention model for the reinforcement learning domain. This model uses a soft, top-down attention mechanism to create a bottleneck in the agent, forcing it to focus on task-relevant information by sequentially querying its view of the environment. The output of the attention mechanism allows direct observation of the information used by the agent to select its actions, enabling easier interpretation of this model than of traditional models. We analyze different strategies that the agents learn and show that a handful of strategies arise repeatedly across different games. We also show that the model learns to query separately about space and content (`where' vs. `what'). We demonstrate that an agent using this mechanism can achieve performance competitive with state-of-the-art models on ATARI tasks while still being interpretable.
1 Introduction
The paper addresses the limited interpretability of traditional reinforcement-learning agents by introducing attention maps that expose task-relevant information used for decisions. Its soft, top-down attention focuses on important visual components and supports analysis of the agent’s decision process.
- Traditional RL architectures are effective but make it difficult to understand decision inputs and errors.
- The proposed agent sequentially queries visual inputs to obtain compressed, query-dependent summaries for generating outputs.
- Attention maps expose the information selected by the agent and help analyze how it solves tasks.
- The model attends to task-relevant components, including regions ahead of the player, enemies, and important moving objects.
- The agent localizes attention using both frame content and absolute spatial position, including in unseen environment configurations.
2 Model
The model combines a recurrent vision core, top-down query generation, soft spatial attention, and compressed answer vectors in a differentiable architecture. A spatial basis preserves location information through the bottleneck and supports separate “what” and “where” queries.
- An attention head queries a large input tensor and returns a low-dimensional summary vector used to produce the model’s output.Multiple attention heads are implemented by producing multiple queries and receiving multiple answers.
- The vision core uses convolutional and recurrent processing, then splits its output into keys and values augmented with a fixed spatial basis.
- The top LSTM generates query vectors from its previous state, and inner products with spatial key locations produce attention logits.
- A spatial softmax normalizes the logits into attention maps, which weight the values tensor before spatial summation into answer vectors.
- The answer vectors and queries feed the LSTM, which produces the next state and output.
- Top-down queries depend on the LSTM state rather than the observation, while the spatial sum creates a severe information bottleneck.
- The spatial basis: The spatial basis preserves positional information after summation and enables mixed queries about object content and location.
- The spatial basis: The convolutional vision core has limited ability to represent absolute spatial position, motivating the added spatial basis.
3 Related work
Prior recurrent attention models span many modalities and tasks and differ in whether attention is hard or soft and windowed or weighting-based. This work adapts key-query-value attention to reinforcement learning with top-down queries and compressed outputs.
- Recurrent attention mechanisms have been applied to question answering, translation, video, image tasks, text classification, generative modeling, tracking, and reinforcement learning.
- Existing mechanisms vary by hard versus soft attention and explicit attention windows versus weighting mechanisms.
- The model uses soft key-query-value attention but generates queries top-down rather than directly from the input.
- Its attention output is highly compressed and lacks spatial structure except for information encoded through the spatial basis.
- Unlike related systems designed for question answering or image captioning, this model generates its own queries for reinforcement-learning decisions.
4 Analysis and Results
The attention agent performs competitively on Atari while exposing task-relevant visual strategies through interpretable attention maps. Across games, it generalizes to novel configurations, plans over trajectories, uses trip-wires, and separates content- from location-based queries.
- Agent and evaluation: The agent uses a three-layer convolutional vision core, convolutional LSTM, four attention queries, and actor-critic training with VTRACE.A fully connected LSTM produces the policy and baseline from query-answer vectors, previous reward, and previous action.
- Agent and evaluation: The attention agent is competitive with state-of-the-art baselines on Atari benchmark results.The comparison uses feedforward and LSTM baselines without attentional bottlenecks.
- Attention strategies: Attention maps concentrate on task-relevant entities, including the player, enemies, power-ups, and score.Unattended regions receive weights close to zero during the spatial summation.
- Reaction to novel states: The agent reacts appropriately to injected enemies in unseen configurations rather than relying only on fixed game patterns.In Seaquest, it attends to the injected fish, moves toward it, fires, and fires again when the first shot does not destroy it.
- Forward planning/scanning: In planning games, the model scans possible future paths and selects trajectories that are safe or rewarding.The paths follow map structure in Ms Pacman and Alien, and the model produces or executes another path when it detects a ghost.
- The role of top-down influence: 541.1% median human normalized score was achieved by the attention agent, versus 274.7% for the fixed-query agent and 274.5% for the L2-Norm Key Agent.Mean scores were 975.5%, 615.2%, and 561.0%, respectively, across the seven games.
- Trip wires: Trip-wires place attention at strategic spatial points, triggering actions when moving objects cross them.Breakout uses a two-stage focus around the ball, while Space Invaders uses trip-wires as shields against projectiles.
- The role of top-down influence: Different attention heads combine content and spatial queries, with some tracking opponent cars, some scanning locations, and some following the player and score.The query’s content component interacts with frame-dependent keys, while its spatial component interacts with a fixed spatial basis.
5 Conclusion
The attention agent achieves competitive ATARI performance while producing interpretable maps of attended information. These maps reveal recurring strategies, including task-focused attention, trajectory scanning, and behavioral tripwires.
- The agent achieves performance competitive with state-of-the-art agents across a broad range of ATARI levels.
- Agents learn to focus on key features, look ahead along short trajectories, and place tripwires to trigger behaviors.
- The model combines “what” and “where” queries to select regions and objects depending on the task.
- Top-down attention supports more informative visualizations of the information used to inform the agent’s policy.
A.1 Agent Description
The attention agent processes RGB ATARI frames through convolutional and recurrent vision components, then uses attention answers and recurrent policy processing to produce policy logits and value estimates.
- The agent processes RGB frames through a two-layer ConvNet and ConvLSTM, producing a 27 × 20 × 128 vision output.
- The vision output is split into keys and values, with a shared spatial basis appended to both tensors.
- A three-layer MLP converts the previous LSTM state into four attention-head queries.
- An answer processor combines queries, answers, the previous action, and the previous reward before policy-core processing.
- The policy core and subsequent MLPs produce policy logits and value estimates.
A.2 Multi-Level Agents
The multi-level agent trains across all ATARI levels using distinct actors feeding a shared learner, with population-based training evolving selected optimization parameters.
- Distinct actors act on different ATARI levels while feeding trajectories to the same learner.
- Population-based training uses a population size of 16.
- The training process evolves the learning rate, entropy cost, RMSProp ϵ, and gradient clipping threshold.
A.3 Agent Performance
Performance is evaluated with expert agents trained separately on ATARI levels and a multi-level agent trained across levels. Freeway and Venture are omitted because every tested agent obtained zero return there.
- The multi-level agent achieves lower scores on almost all levels than the expert agents.
- 55 ATARI-level expert training curves are shown, with Freeway and Venture omitted from the curves.
- Freeway and Venture are omitted because no tested agent achieved a non-zero return on either game.
A.4 Top-Down versus Bottom-Up
The comparison tests whether top-down attention improves Atari performance relative to bottom-up alternatives. In 6 of 7 tested games, removing top-down attention significantly reduces performance.
- Top-Down versus Bottom-Up: In 6 of 7 tested games, agents without top-down attention perform significantly worse than agents with top-down attention.The comparison uses the Fixed Query Agent and L2 Norm Keys agent, trained on single levels for 2e9 frames.
- Top-Down versus Bottom-Up: Table 4 reports the final scores for each agent across all 7 tested levels.
- Top-Down versus Bottom-Up: Figure 9 compares the performance of individual experts on selected ATARI games.
A.5 What-Where Analysis
The analysis separates attention into spatial and content contributions and tests whether the agent relies mainly on highly attended regions. The results show interpretable, sometimes local attention and little performance loss after suppressing low-weight regions.
- What-Where Analysis: What-where maps measure each location’s relative content and spatial contribution to a query’s total attention weight.Colors interpolate between red, white, and blue; blue or red indicates one contribution exceeds the other by more than 10 times.
- Attention-Weight Ablation: Performance does not degrade for hard-attention thresholds t ≤0.1 on Breakout, Ms. Pacman, Seaquest, and Space Invaders.The modified agent suppresses areas with small attention weights before evaluation.
- Attention-Weight Distribution: Some attention heads are highly peaked while others are diffuse in both Ms. Pacman and Space Invaders.The plotted x- and y-axis weight sums range from 0 to 1, and peakiness controls how local the received information is.
- Attention-Weight Ablation: Agents achieve broadly similar scores across games when the t = 0.1 score is compared with the t = 0.0 score.Table 5 reports this ratio from 15 runs at each threshold, with statistical uncertainty for the ratio of mean scores.
- Score Comparisons: Table 3 lists human-normalized scores for agents on all Atari levels, while Table 4 compares attention agents with two bottom-up experiments.