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 RL often has poor sample efficiency and limited generalization, motivating architectures that can represent relations explicitly. The paper uses structured perception and iterative self-attention to guide a model-free policy, achieving strong results in Box-World and StarCraft II while leaving some generalization effects unresolved.
Problem
Deep RL models face low sample efficiency and a propensity not to generalize to seemingly minor task changes.
Method
The paper equips deep RL agents with structured perception and iterative self-attention to reason over entities and their relations.
Results
The relational agents solve more than 98% of Box-World levels, achieve state-of-the-art performance on six StarCraft II mini-games, and surpass human grandmaster performance in four.
Takeaways & Limitations
Structured perception and relational reasoning produce interpretable representations and improve sample complexity, generalization, and overall performance relative to baseline agents.
Takeaways & Limitations
Generalization effects in StarCraft II show high variability, diminish with larger models, and require further investigation.
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 models often struggle with sample efficiency and generalization because they may overfit training data. The paper proposes relational representations and evaluates them in Box-World and StarCraft II.
- Deep RL remains limited by low sample efficiency and weak generalization to minor task changes.These limitations are linked to difficulty learning abstract, interpretable, and generalizable problem representations.
- The approach combines deep learning with Relational RL by learning reusable entity- and relation-centric functions.These functions implicitly reason over relational representations while retaining deep learning’s representational capacity.
- The paper introduces Box-World, a navigation and planning task explicitly designed to test relational reasoning.It also applies the agent to StarCraft II mini-games, where relational reasoning is relevant to interactions among entities.
- The proposed relational agents achieve state-of-the-art performance on six StarCraft II mini-games.The introduction identifies this benchmark result as a principal contribution.
2 Relational reinforcement learning
Relational reinforcement learning represents states, actions, and policies relationally rather than propositionally. This supports generalization and allows background knowledge to guide policy search.
- Relational RL combines reinforcement learning with relational learning or Inductive Logic Programming.States, actions, and policies are represented using a first-order or relational language.
- Relational representations facilitate generalization across goals, states, and actions by reusing knowledge learned earlier.They also provide a way to incorporate background knowledge relevant to the learning problem.
- Background knowledge can be expressed as logical facts and rules, such as predicates describing spatial relations between blocks.These predicates can be reused for different objects during learning.
- The paper translates relational-learning ideas into architectural inductive biases for neural agents operating on sets of scene entities.The agents use iterated, message-passing-like processing to learn pairwise and higher-order interactions.
3 Architecture
The architecture extracts spatially located entities from visual inputs and applies non-local, iterative self-attention to compute relational interactions. Multiple attention blocks and a policy-learning setup support the Box-World and StarCraft II agents.
- Architecture: The architecture uses non-local computations with a shared function and iterative computation as its two guiding principles.Pairwise interactions are computed independently of spatial proximity, while iteration supports higher-order interactions.
- Computing non-local interactions using a shared function: Multi-head dot-product self-attention projects each entity into query, key, and value vectors, then mixes value vectors using attention weights.The weighted interactions are computed across all entities rather than only nearby entities.
- Computing non-local interactions using a shared function: Multiple independent attention heads operate in parallel and may learn different relational semantics during training.Their outputs are concatenated, processed by an MLP, combined with residual connections, and layer-normalized.
- Computing non-local interactions using a shared function: An attention block performs non-local pairwise relational computation, and multiple blocks can approximate higher-order relations through repeated processing.Blocks may use shared recurrent parameters or unshared deep parameters.
- Extracting entities: A convolutional front-end converts pixel inputs into spatially tagged feature vectors compiled as an entity matrix E.The resulting n^2 pixel-feature vectors remain agnostic to what constitutes an entity for a particular task.
- Agent training and task adaptations: The agents use an actor-critic training setup with distributed actors generating experience for a learner that updates the policy and value baseline.For Box-World, the relational module is compared with residual-convolution control agents; StarCraft II adds recurrent processing for recent history.
4 Experiments and results
Experiments evaluate the relational agent in Box-World and StarCraft II, showing strong task performance, interpretable attention patterns, and transfer to some unseen Box-World conditions. StarCraft II results are strong across the mini-games, while generalization findings vary with model size and seed.
- Box-World task: Box-World requires relational reasoning and planning because a unique box sequence reaches the gem, while distractor branches and random spatial locations complicate action selection.Keys are single-use, and opening a distractor box makes the level unsolvable.
- Box-World results: More than 98% of Box-World levels were solved by relational agents, compared with less than 75% for convolutional and fully connected control agents.Short distractor branches required one attention block for top performance, whereas long branches required more blocks.
- Interpretability: Attention analysis found heads linking keys with unlockable locks and relating objects to the agent, suggesting interpretable relational computations for object interaction and navigation.Using one-hot inputs, the analysis also found locks attending to usable keys and agent-gem mutual attention.
- Generalization: Without further training, relational agents solved more than 88% of levels requiring solution lengths of 6, 8, or 10, demonstrating zero-shot transfer to longer Box-World sequences.The tested conditions also included key-lock combinations previously seen only on distractor paths.
- StarCraft II results: The relational StarCraft II agent achieved state-of-the-art performance in six mini-games and surpassed the human grandmaster in four.Scores were measured as means over 30 episodes using the full action set.
- StarCraft II comparison: Across all StarCraft II mini-games, the relational implementation achieved results equal to or better than the control implementation.Both implementations also improved substantially over the previous best, with several possible contributors beyond relational inductive biases.
- Limitations: StarCraft II generalization results were highly variable, with effects diminishing for larger models and further work needed to understand the relational agent’s generalization.The best seed of a medium-sized relational agent achieved better generalization scores in the test scenario.
5 Conclusion
The agents combine structured perception and relational reasoning to learn interpretable, generalizable representations and outperform baseline agents. The authors identify future work in richer perceptual reasoning, hierarchical planning, and understanding learned representation semantics.
- Conclusion: Structured perception and relational reasoning produced interpretable representations that improved sample complexity, generalization, and overall performance over baseline agents.Behavioral analysis linked the representations to better generalization, while attention analysis found interpretable task-relevant entity interactions.
- Conclusion: The attention mechanism’s internal computations were interpretable and consistent with computing task-relevant relations.
- Future work: Future work should investigate more sophisticated structured perceptual reasoning and hierarchical reinforcement learning or planning.These directions aim to translate structured representations and reasoning more fully into structured behaviors.
- Future work: Further study is needed to explore the semantics of learned representations through comparisons with knowledge hard-coded in traditional relational reinforcement learning.
- Conclusion: The relational architectural prior may give a model-free agent access to abstract planning-like behavior without constituting a model in the technical sense.The authors present this as a speculative possibility rather than a demonstrated equivalence.
Task
Box-World levels are procedurally generated from random trees that specify solution paths and distractor branches. This generation creates many distinct levels and supports training-test splits that withhold particular level structures during training.
- Task: Each Box-World level is procedurally generated from a random tree defining the correct box-opening path and multiple distractor branches.Keys, locks, agents, and boxes are positioned randomly, with 20 keys and 20 locks sampled for each level.
- Task: The level generator assigns rewards of +10 for collecting the gem, +1 for opening a solution-path box, and −1 for opening a distractor box.A level ends immediately when the gem is collected or a distractor box is opened.
- Task: The large number of possible generated trees makes repeated exposure to the same level extremely unlikely.
- Task: Procedural generation enables training-test splits that withhold levels conforming to particular cases during training and present them at test time.
Agent architecture
The agent uses an attention-based relational module followed by fully connected layers to produce policy logits and a value baseline. Training uses RMSprop with a tuned learning rate and fixed entropy, discount, and rollout settings.
- Agent architecture: The relational module used 2 to 4 attention heads with embedding size d = 64, followed by feature-wise max pooling and four ReLU fully connected layers.
- Agent architecture: The network produced policy logits π of size 4 and a baseline value V of size 1 through linear projections.The policy logits were normalized for multinomial action selection.
- Training: Training used entropy cost 0.005, discount γ = 0.99, and an unroll length of 40 steps.
- Training: RMSprop training used momentum 0, ϵ = 0.1, decay 0.99, and a tuned learning rate between 1e−5 and 2e−4.The authors also report replicating the results with A3C, although training took longer.
Control agent architecture
The control agent preserves the relational agent’s overall architecture but replaces its relational module with residual convolutional processing. This provides a non-relational baseline with adjustable convolutional depth.
- Control agent architecture: The control agent used the same overall architecture as the relational agent while replacing the relational module with 3 to 6 residual-convolutional blocks.
- Control agent architecture: Each residual block contained two convolutional layers with 3 × 3 kernels, stride 1, and 26 output channels.
- Control agent architecture: The control architecture therefore varied convolutional depth while keeping the comparison centered on the relational module replacement.
B StarCraft II mini-games
The StarCraft II mini-game experiments use extensive training and include generalization evaluations in Collect Mineral Shards and Box-World branching comparisons.
- StarCraft II agents were trained for 10 billion steps using batches of 32 trajectories, each unrolled for 80 steps.
- Figure 6 contrasts forward and backward branching in Box-World, where backward branching permits a more reactive policy.With backward branching, each key opens only one box, although a key type can appear in multiple boxes.
- Figure 7 evaluates Collect Mineral Shards agents trained with 2 marines and tested with 1, 2, 3, 4, 5, or 10 marines.Colored bars show mean scores of the ten best seeds, with standard error bars.
- Approximately 100 experiments were run for each mini-game, following the stated hyperparameter settings and three seeds.
Relational Agent architecture
The relational agent processes spatial observations, recurrent history, and relational representations before producing policy, value, and action-argument outputs.
- The StarCraft II relational architecture follows the Box-World agent while adapting its components to StarCraft II constraints.
- Input-preprocessing: Inputs include minimap, screen, player, and last-action information, with numerical features logarithmically rescaled and categorical features embedded in 10 dimensions.
- State encoding: Independent residual convolutional blocks process spatial inputs and downsample them to [8 × 8 × #channels1] outputs.
- Memory processing: A Conv2DLSTM combines inputs2D with its previous state to produce outputs2D representing aggregated observation history.
- Relational processing: Flattened outputs2D pass through stacked MHDPA blocks, producing relational-spatial and relational-nonspatial pathways.The relational-nonspatial pathway uses feature-wise max-pooling and a 2-layer MLP.
- Output processing: Shared features generate masked policy logits and baseline values through separate 2-layer MLPs.
- Output processing: Sampled actions are embedded in 16 dimensions to condition non-spatial and spatial argument logits.
Control agent architecture
The control agent replaces relational processing with a deep residual model while retaining the remaining architecture, alongside fixed and swept hyperparameter settings.
- The baseline control agent differs from the relational agent only in its relational-processing component.
- Hyperparameters: Table 3 specifies fixed MHDPA settings for the StarCraft II mini-games.
- Relational processing: A 12-layer residual model forms the control agent’s relational-spatial outputs.It contains four blocks of three convolutional layers with skip-connections and ReLU activations.
- Relational processing: The control agent also flattens these tensors and uses a 2-layer MLP to produce relational-nonspatial outputs.
- Hyperparameters: Table 4 lists hyperparameters swept across mini-games.