Source-linked AI summary
Relational Deep Reinforcement Learning
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, Murray Shanahan, Victoria Langston, Razvan Pascanu, Matthew Botvinick, Oriol Vinyals, Peter Battaglia
TL;DR
Deep reinforcement learning often has low sample efficiency and struggles to generalize across minor task changes. This paper introduces learned entity- and relation-centric functions with attention-based relational reasoning. Its agents improve baseline performance across Box-World and StarCraft II, achieving state-of-the-art results in six mini-games.
Problem
Deep RL still has low sample efficiency and often fails to generalize across seemingly minor task changes.
Source: §1 · p. 1Method
The approach uses learned, reusable entity- and relation-centric functions with attention blocks that perform non-local pairwise relational computations.
Source: §1 · p. 1, Computing non-local interactions using… · p. 3Results
Takeaways & Limitations
The work opens new directions for RL through a principled hybrid of flexible statistical learning and structured approaches.
Source: §5 · p. 9StarCraft II generalization effects remain uncertain: results were highly variable and diminished with larger models.
Source: Generalization capability · p. 9
Abstract
from arXiv · showhide
We introduce an approach for deep reinforcement learning (RL) that improves upon the efficiency, generalization capacity, and interpretability of conventional approaches through structured perception and relational reasoning. It uses self-attention to iteratively reason about the relations between entities in a scene and to guide a model-free policy. Our results show that in a novel navigation and planning task called Box-World, our agent finds interpretable solutions that improve upon baselines in terms of sample complexity, ability to generalize to more complex scenes than experienced during training, and overall performance. In the StarCraft II Learning Environment, our agent achieves state-of-the-art performance on six mini-games -- surpassing human grandmaster performance on four. By considering architectural inductive biases, our work opens new directions for overcoming important, but stubborn, challenges in deep RL.
1 Introduction
Deep RL can suffer from low sample efficiency and poor generalization, motivating relational representations that support abstract, interpretable understanding. The paper proposes entity- and relation-centric reasoning and evaluates it on Box-World and StarCraft II mini-games.
Introduction
- Deep RL models can be sample-inefficient and fail to generalize to minor task changes, suggesting they overfit rather than learn abstract, interpretable problem representations.The passage situates these limitations amid recent advances in deep RL [3] and cites prior work on generalization limitations [4] [5] [6] [7].Source: §1 · p. 1
- The approach combines relational learning’s generalization power with deep learning by using learned, reusable entity- and relation-centric functions to reason over relational representations.It builds on Relational RL, which advocated relational state, action, and policy representations.Source: §1 · p. 1
- The agent achieves state-of-the-art performance on six StarCraft II mini-games and demonstrates generalization behaviors in Box-World.The authors create and analyze Box-World to target relational reasoning, comparing agents with attention-based relational representations against agents without them.Source: §1 · p. 1
2 Relational reinforcement learning
Relational reinforcement learning represents states, actions, and policies relationally to support generalization and the use of background knowledge. The paper brings relational inductive biases into deep RL through structured scene representations and iterative relational reasoning.
2 Relational reinforcement learning
- Relational RL represents states, actions, and policies in a first-order language, supporting generalization across goals, states, and actions and the use of relevant background knowledge.Background knowledge can be supplied as logical facts and rules.Source: §2 · p. 2
- The representational language, background knowledge, and assumptions form an inductive bias that guides and restricts the search for good policies.The declarative bias determines how concepts can be represented.Source: §2 · p. 2
- The approach translates relational inductive biases into deep RL architectures that process scene entities with iterative, message-passing-like relational reasoning.Entities correspond to local image regions; the agent learns to attend to key objects and compute pairwise and higher-order interactions.Source: §2 · p. 2
3 Architecture
The architecture uses shared, non-local pairwise interactions and iterative computation to model relations among entities, with entities extracted from visual inputs. Its Box-World and StarCraft II implementations adapt this relational module to each task’s inputs and outputs.
3 Architecture
- The design combines shared, non-local pairwise computations with iteration, aiming to learn relations beyond spatially local interactions and capture higher-order interactions.The shared function computes interactions independently of entity proximity; multiple computation steps can represent higher-order relations.Source: §3 · p. 2
Computing non-local interactions using a shared function
- Each attention block uses multi-head dot-product self-attention to mix entity values, then applies an MLP, a residual connection, and layer normalization; stacked blocks support higher-order relations.Blocks can use shared recurrent or unshared deep parameters.Source: Computing non-local interactions using… · p. 3
Extracting entities
- A CNN converts pixels into feature maps, adds each feature vector’s x-y position, and treats the resulting n^2 vectors as entities without prespecifying what an entity is.The entity matrix has dimensions n^2 × k, where k is the number of CNN output channels.Source: Extracting entities · p. 4
Agent architecture for Box-World
- For Box-World, two convolutional layers produce position-tagged features for stacked shared-weight attention blocks, followed by spatial max-pooling and an MLP for policy logits and value.The convolutional layers use 12 and 24 kernels, each 2 × 2 with stride 1.Source: Computing non-local interactions using… · p. 3, Agent architecture for Box-World · p. 4
Agent architecture for StarCraft II
- For StarCraft II, the agent adds residual convolutional blocks and a 2D-ConvLSTM to retain recent action history, and produces spatial and non-spatial action arguments.The ConvLSTM addresses cases where later observations omit details of prior actions, such as a marine’s chosen path.Source: Computing non-local interactions using… · p. 3, Agent architecture for StarCraft II · p. 4
4 Experiments and results
The experiments show that relational reasoning supports strong Box-World performance and transfer, with attention patterns reflecting key-lock and object-agent relations. In StarCraft II, the relational agent achieves state-of-the-art results in six mini-games and exceeds human grandmaster performance in four.
Task description
- Box-World requires unlocking a unique sequence of boxes to collect a gem, while avoiding distractor boxes that make the level unsolvable.Its 12 × 12-pixel rooms place keys and boxes randomly; keys can be used only once, so success depends on relational reasoning rather than object locations.Source: Task description · p. 5
Training results
- More than 98% of Box-World levels were solved with the relational module, versus less than 75% by controls across both distractor variants.One attention block sufficed for short distractor branches, while long branches required more; forward branching requires planning to avoid wrong locks, unlike backward branching's reactive matching strategy.Source: Training results · p. 6
Generalization capability: testing on withheld environments
- More than 88% of levels were solved zero-shot across solution lengths 6, 8, and 10; the relational agent also solved 97% of levels with unseen key-lock combinations.Without the relational module, performance fell to 5% on six-box sequences and 0% on eight- and ten-box sequences; it reached only 13% on levels with new key-lock combinations.Source: Figure 5 · p. 7, Generalization capability: testing on… · p. 7
Visualization of attention weights
- Attention patterns linked keys to matching locks and related objects to the agent, including when one-hot inputs removed shared key-lock color cues.One-hot analysis also found locks attending to keys that could unlock them, and the agent and gem attending to each other and themselves.Source: Visualization of attention weights · pp. 6–7, Figure 4 · p. 6
Training results
- Six mini-games achieved state-of-the-art results, and performance exceeded human grandmaster performance in four.The agent used the full SC2LE action set, with mean scores measured over 30 episodes per mini-game; relational performance matched or exceeded the control across all mini-games.Source: Table 1 · p. 8, Training results · p. 8
Generalization capability
- For five-marine StarCraft II tests, the best-seed relational agent generalized better, but results varied and the effect diminished with larger models.Agents trained with two marines were tested without further training on modified levels with five; the authors note that larger models may be more prone to overfitting and that further study is needed.Source: Generalization capability · p. 9
5 Conclusion
The paper concludes that structured perception and relational reasoning help deep RL agents learn interpretable representations and outperform baselines, with behavioral and attention-mechanism analyses supporting generalization and interpretable computations.
5 Conclusion
- The agents exceed baselines in sample complexity, generalization, and overall performance, while behavioral and attention analyses show interpretable representations and task-relevant internal computations.The behavioral analysis found better generalization; the attention analysis found computations congruent with task-relevant relations.Source: §5 · p. 9
- Future work could improve structured perception with computer-vision mechanisms [34], connect structured reasoning to behavior through hierarchical RL and planning [35] [36], and further investigate learned-representation semantics.The paper proposes examining those semantics through the lens of what might be hard-coded in traditional RRL.Source: §5 · p. 9
- The authors speculate that relational architectural priors may give model-free agents access to general environmental knowledge in a way analogous to imagination-based forward models [37] [38] [39].They describe this as a possible hybrid of flexible statistical learning and structured approaches, while noting the agent is not a model in the technical sense.Source: §5 · p. 9
Appendix · A Box-world · Task
Box-world procedurally generates levels from trees that define solution paths and distractor branches, with randomized object placement and keys and locks. Its large variety supports withholding particular level types for testing.
Task
- Each level uses a randomly generated tree to define the boxes on the path to the gem and distractor branches that lead to dead ends.The agent, keys, and boxes are placed randomly, with enough room to navigate; 20 keys and 20 locks are sampled for each level.Source: Task · p. 12
- Collecting the gem rewards +10, opening a solution-path box rewards +1, and opening a distractor box incurs −1; either collecting the gem or opening a distractor ends the level.Source: Task · p. 12
- Procedural generation produces many possible trees, making it extremely unlikely for the agent to encounter the same level twice.Source: Task · p. 12
- Training-test splits can withhold levels matching a particular case during training and present them to the agent at test time.Source: Task · p. 12
Agent architecture
The agent used multi-head attention followed by max pooling and a four-layer ReLU network to produce policy logits and a baseline. Its training setup specified entropy, discount, unroll, and optimizer settings.
- The model used 2 to 4 attention heads with embedding size 64, alongside entropy cost 0.005, discount γ = 0.99, and unroll length 40 steps.Source: Agent architecture · p. 12
- Feature-wise max pooling fed four fully connected ReLU layers, whose linear projection produced policy logits of size 4 and a baseline of size 1.The normalized policy logits defined a multinomial distribution used to sample actions.Source: Agent architecture · p. 12
- Training used RMSprop with momentum 0, ϵ = 0.1, decay 0.99, and a tuned learning rate from 1e−5 to 2e−4.The authors noted they could replicate the results with A3C, though training took longer.Source: Agent architecture · p. 12
Control agent architecture
The baseline control agent kept the relational agent’s architecture but replaced its relational module with residual-convolutional blocks.
- The baseline used 3 to 6 residual blocks, each comprising two 3 × 3 convolutional layers with stride 1 and 26 output channels.Source: Control agent architecture · p. 12
B StarCraft II mini-games
The supplied passages describe the StarCraft II mini-game training setup and a generalization evaluation for Collect Mineral Shards; they do not report performance results.
B StarCraft II mini-games
- Agents trained for 10 billion steps with Adam, using batches of 32 trajectories unrolled for 80 steps and a linearly decayed learning rate.Entropy was also linearly decayed, but the passage gives no further details.Source: B StarCraft II mini-games · p. 12
- Collect Mineral Shards agents trained on levels with 2 marines were tested with 1, 2, 3, 4, 5, or 10 marines.Figure 7 reports mean scores for the ten best seeds, with standard error bars; the supplied caption gives no score values.Source: Figure 7 · p. 13
- Approximately 100 experiments were run for each mini-game, following Table 4 hyperparameter settings and using 3 seeds.The passage refers to Table 2 for loss-scaling details but does not provide them.Source: B StarCraft II mini-games · p. 13
Relational Agent architecture
The StarCraft II agent closely follows the Box-World architecture, with changes to accommodate SC2 constraints. Its pipeline combines spatial and non-spatial inputs with recurrent memory and relational processing to produce policies, value estimates, and action arguments.
Relational Agent architecture
- SC2 inputs are rescaled or embedded, then spatial and non-spatial sources are encoded separately before recurrent processing.Numerical features receive a logarithmic transformation, categorical features 10-dimensional embeddings; minimap and screen use residual convolutional blocks, while player and last-action inputs use a two-layer MLP.Source: Relational Agent architecture · p. 13
- MHDPA processes recurrently aggregated observations through separate relational-spatial and relational-nonspatial pathways.The spatial pathway restores an 8 × 8 shape; the non-spatial pathway uses feature-wise max pooling followed by a two-layer, 512-unit MLP.Source: Relational Agent architecture · p. 13
- Shared features produce masked policy logits and baseline values through separate two-layer MLP heads.Sampled actions are embedded in 16 dimensions to condition non-spatial argument logits and, after deconvolution, spatial argument logits.Source: Relational Agent architecture · p. 13
Control agent architecture
- The control agent replaces relational processing with a 12-layer residual convolutional model and a separate MLP for non-spatial outputs.The residual model has four blocks of three convolutions and follows [40]; the remaining architecture matches the relational agent.Source: Relational Agent architecture · pp. 13–14
Related papers
This paper cites
- Attention Is All You Need2017
- Semi-Supervised Classification with Graph Convolutional Networks2016
- Relational inductive biases, deep learning, and graph networks2018
- Relation Networks for Object Detection2017
- Neural Combinatorial Optimization with Reinforcement Learning2016
- Attention, Learn to Solve Routing Problems!2018