Original: D EEP REINFORCEMENT LEARNING WITH RELATIONAL INDUCTIVE BIASES
Revised: D EEP REINFORCEMENT LEARNING WITH RELATIONAL INDUCTIVE BIASES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We introduce an approach for augmenting model-free deep reinforcement learning agents with a mechanism for relational reasoning over structured representations, which improves performance, learning efﬁciency, generalization, and interpretability.
Revised: We introduce an approach for augmenting model-free deep reinforcement learning agents with a mechanism for relational reasoning over structured representations, which improves performance, learning efﬁciency, generalization, and interpretability.

Original: Our architecture encodes an image as a set of vectors, and applies an iterative message-passing procedure to discover and reason about relevant entities and relations in a scene.
Revised: Our architecture encodes an image as a set of vectors, and applies an iterative message-passing procedure to discover and reason about relevant entities and relations in a scene.

Original: In six of seven StarCraft II Learning Environment mini-games, our agent achieved state-of-the-art performance, and surpassed human grandmasterlevel on four.
Revised: In six of seven StarCraft II Learning Environment mini-games, our agent achieved state-of-the-art performance, and surpassed human grandmasterlevel on four.

Original: In a novel navigation and planning task, our agent’s performance and learning efﬁciency far exceeded non-relational baselines, it was able to generalize to more complex scenes than it had experienced during training.
Revised: In a novel navigation and planning task, our agent’s performance and learning efﬁciency far exceeded non-relational baselines, it was able to generalize to more complex scenes than it had experienced during training.

Original: Moreover, when we examined its learned internal representations, they reﬂected important structure about the problem and the agent’s intentions.
Revised: Moreover, when we examined its learned internal representations, they reﬂected important structure about the problem and the agent’s intentions.

Original: The main contribution of this work is to introduce techniques for representing and reasoning about states in model-free deep reinforcement learning agents via relational inductive biases.
Revised: The main contribution of this work is to introduce techniques for representing and reasoning about states in model-free deep reinforcement learning agents via relational inductive biases.

Original: Our experiments show this approach can offer advantages in efﬁciency, generalization, and interpretability, and can scale up to meet some of the most challenging test environments in modern artiﬁcial intelligence.
Revised: Our experiments show this approach can offer advantages in efﬁciency, generalization, and interpretability, and can scale up to meet some of the most challenging test environments in modern artiﬁcial intelligence.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Recent deep reinforcement learning (RL) systems have achieved remarkable performance in very challenging problem domains (Mnih et al., 2015; Silver et al., 2016), in large part because of their ﬂexibility in how they learn and exploit the statistical structure underlying observations and reward signals.
Revised: Recent deep reinforcement learning (RL) systems have achieved remarkable performance in very challenging problem domains (Mnih et al., 2015; Silver et al., 2016), in large part because of their ﬂexibility in how they learn and exploit the statistical structure underlying observations and reward signals.

Original: But the downsides to such ﬂexibility often include low sample efﬁciency and poor transfer beyond the speciﬁcs of the training environment (Zhang et al., 2018;
Revised: But the downsides to such ﬂexibility often include low sample efﬁciency and poor transfer beyond the speciﬁcs of the training environment (Zhang et al., 2018;

Original: Lake et al., 2017; Kansky et al., 2017).
Revised: Lake et al., 2017; Kansky et al., 2017).

Original: Various structured approaches to RL (e.g. Dzeroski et al.
Revised: Various structured approaches to RL (e.g. Dzeroski et al.

Original: ; Driessens & Dzeroski (2004);
Revised: ; Driessens & Dzeroski (2004);

Original: Diuk et al.
Revised: Diuk et al.

Original: ; Garnelo et al.
Revised: ; Garnelo et al.

Original: (2016)) have attempted to overcome these limitations by explicitly incorporating entity-based and symbolic representations, and specialized building blocks for solving the task at hand.
Revised: (2016)) have attempted to overcome these limitations by explicitly incorporating entity-based and symbolic representations, and specialized building blocks for solving the task at hand.

Original: Although these approaches are often highly efﬁcient, they constrain the representations and admissible learning algorithms, they struggle to learn rich representations, and they are therefore conﬁned to relatively simple tasks and data conditions.
Revised: Although these approaches are often highly efﬁcient, they constrain the representations and admissible learning algorithms, they struggle to learn rich representations, and they are therefore conﬁned to relatively simple tasks and data conditions.

Original: To strike favorable tradeoffs between ﬂexibility and efﬁciency, a number of recent approaches have explored using relational inductive biases in deep learning, to reap the beneﬁts of ﬂexible statistical learning and more structured approaches.
Revised: To strike favorable tradeoffs between ﬂexibility and efﬁciency, a number of recent approaches have explored using relational inductive biases in deep learning, to reap the beneﬁts of ﬂexible statistical learning and more structured approaches.

Original: Methods such as “graph networks” (Scarselli et al., 2009;
Revised: Methods such as “graph networks” (Scarselli et al., 2009;

Original: Li et al., 2015; Battaglia et al., 2018) explicitly represent entities and their relations using using sets and graphs, and perform relational reasoning using learned message-passing (Gilmer et al., 2017) and attention (Vaswani et al., 2017; Hoshen, 2017; Velickovic et al., 2017; Wang et al., 2017) schemes.
Revised: Li et al., 2015; Battaglia et al., 2018) explicitly represent entities and their relations using using sets and graphs, and perform relational reasoning using learned message-passing (Gilmer et al., 2017) and attention (Vaswani et al., 2017; Hoshen, 2017; Velickovic et al., 2017; Wang et al., 2017) schemes.

Original: Because they are implemented using deep neural networks, they can learn transformations from input observations into task-relevant entities, as well as functions for computing rich interaction among these entities.
Revised: Because they are implemented using deep neural networks, they can learn transformations from input observations into task-relevant entities, as well as functions for computing rich interaction among these entities.

Original: This provides a powerful capacity for combinatorial generalization, where their learned building blocks can be composed to represent and reason about novel scenarios (Battaglia et al., 2016;
Revised: This provides a powerful capacity for combinatorial generalization, where their learned building blocks can be composed to represent and reason about novel scenarios (Battaglia et al., 2016;

Original: Raposo et al., 2017; Dai et al., 2017; Wang et al., 2018;
Revised: Raposo et al., 2017; Dai et al., 2017; Wang et al., 2018;

Original: Selsam et al., 2018; Hamrick et al., 2018;
Revised: Selsam et al., 2018; Hamrick et al., 2018;

Original: Kipf et al., 2018; Sanchez-Gonzalez et al., 2018).
Revised: Kipf et al., 2018; Sanchez-Gonzalez et al., 2018).

Original: Drawing on several lines of work, we introduce an approach for incorporating relational inductive biases for entity- and relation-centric state representations, and iterated relational reasoning, into a deep RL agent.
Revised: Drawing on several lines of work, we introduce an approach for incorporating relational inductive biases for entity- and relation-centric state representations, and iterated relational reasoning, into a deep RL agent.

Original: In contrast with prior work exploring relational inductive biases in deep RL (e.g.,
Revised: In contrast with prior work exploring relational inductive biases in deep RL (e.g.,

Original: Wang et al., 2018), our approach does not rely on a priori knowledge of the structure of the problem and is agnostic to the particular relations that need to be considered.
Revised: Wang et al., 2018), our approach does not rely on a priori knowledge of the structure of the problem and is agnostic to the particular relations that need to be considered.

Original: To handle raw visual input data, our architecture used a convolutional front-end to compute embeddings of sets of entities, similar to previous work in visual question answering, physical prediction, and video understanding (Santoro et al., 2017; Watters et al., 2017; Wang et al., 2017).
Revised: To handle raw visual input data, our architecture used a convolutional front-end to compute embeddings of sets of entities, similar to previous work in visual question answering, physical prediction, and video understanding (Santoro et al., 2017; Watters et al., 2017; Wang et al., 2017).

Original: To perform relational reasoning, we used a self-attention mechanism (Vaswani et al., 2017; Hoshen, 2017; Velickovic et al., 2017) applied iteratively within each timestep, which can be viewed as learned message-passing (Li et al., 2015;
Revised: To perform relational reasoning, we used a self-attention mechanism (Vaswani et al., 2017; Hoshen, 2017; Velickovic et al., 2017) applied iteratively within each timestep, which can be viewed as learned message-passing (Li et al., 2015;

Original: Gilmer et al., 2017).
Revised: Gilmer et al., 2017).

Original: Our deep RL agent is based on an off-policy advantage actor-critic (A2C) method which is very effective across a range of standard RL environments (Espeholt et al., 2018).
Revised: Our deep RL agent is based on an off-policy advantage actor-critic (A2C) method which is very effective across a range of standard RL environments (Espeholt et al., 2018).

Original: Our results show that this relational deep RL agent scale to very challenging tasks, achieving state-ofthe-art performance on six out of seven StarCraft II mini-games (Vinyals et al., 2017), surpassing grandmaster level on four mini-games.
Revised: Our results show that this relational deep RL agent scale to very challenging tasks, achieving state-ofthe-art performance on six out of seven StarCraft II mini-games (Vinyals et al., 2017), surpassing grandmaster level on four mini-games.

Original: Additionally, we introduce a novel navigation and planning task, called “Box-World”, which stresses the planning and reasoning components of the policy, factoring out other challenges like complex vision or large action spaces.
Revised: Additionally, we introduce a novel navigation and planning task, called “Box-World”, which stresses the planning and reasoning components of the policy, factoring out other challenges like complex vision or large action spaces.

Original: Our agent reaches higher ceiling performance, more efﬁciently, than non-relational baseline, and is able to generalize to solve problems with more complex solutions than it had been trained on within this task.
Revised: Our agent reaches higher ceiling performance, more efﬁciently, than non-relational baseline, and is able to generalize to solve problems with more complex solutions than it had been trained on within this task.

Original: We also found that the intermediate representations involved in the relational computations were interpretable, and suggest that the agent has rich understanding of the underlying structure of the problem.
Revised: We also found that the intermediate representations involved in the relational computations were interpretable, and suggest that the agent has rich understanding of the underlying structure of the problem.

Original:  R ELATIONAL DEEP RL AGENT ARCHITECTURE RL ALGORITHM
Revised:  R ELATIONAL DEEP RL AGENT ARCHITECTURE RL ALGORITHM

Original: We started with a deep RL algorithm based on a distributed advantage actor-critic (A2C) method (Espeholt et al., 2018), which the schematic on the left of Figure 2 summarizes.
Revised: We started with a deep RL algorithm based on a distributed advantage actor-critic (A2C) method (Espeholt et al., 2018), which the schematic on the left of Figure 2 summarizes.

Original: The agent creates an embedded state representation, S from its input observation, which is then used to compute as output a policy, π (the “actor”), and a baseline value, B (the “critic”).
Revised: The agent creates an embedded state representation, S from its input observation, which is then used to compute as output a policy, π (the “actor”), and a baseline value, B (the “critic”).

Original: The π consists of logits over the set of possible actions, from which an action to perform is sampled, and the B is an estimate of the statevalue function at the current state.
Revised: The π consists of logits over the set of possible actions, from which an action to perform is sampled, and the B is an estimate of the statevalue function at the current state.

Original: During learning, the B is used to compute the temporal-difference error, which is used both to optimize π to generate actions with greater returns than B predicts, and to optimize B to more accurately estimate state values.
Revised: During learning, the B is used to compute the temporal-difference error, which is used both to optimize π to generate actions with greater returns than B predicts, and to optimize B to more accurately estimate state values.

Original: For further algorithmic details, including the distributed actor/learner framework and off-policy correction methods, please see Espeholt et al.
Revised: For further algorithmic details, including the distributed actor/learner framework and off-policy correction methods, please see Espeholt et al.

Original: I NPUT MODULE
Revised: I NPUT MODULE

Original: Our agent takes as input an image of a scene, transforms it into an embedded state representation,
Revised: Our agent takes as input an image of a scene, transforms it into an embedded state representation,

Original: S , which is a spatial feature map returned by a convolutional neural network (CNN).
Revised: S , which is a spatial feature map returned by a convolutional neural network (CNN).

Original: The relational module (Figure 2) operated on a set of entities, so our architecture transformed the m × n × f feature map into an N × f (where N = m · n ) set of “entity vectors”, E , by reshaping S so that each of E ’s row, e i , corresponded to a feature vector, s x,y , at a particular x, y location in the embedded scene.
Revised: The relational module (Figure 2) operated on a set of entities, so our architecture transformed the m × n × f feature map into an N × f (where N = m · n ) set of “entity vectors”, E , by reshaping S so that each of E ’s row, e i , corresponded to a feature vector, s x,y , at a particular x, y location in the embedded scene.

Original: Crucially, this allowed for non-local computation between entities (Wang et al., 2017), unconstrained by their coordinates in the spatial feature map.
Revised: Crucially, this allowed for non-local computation between entities (Wang et al., 2017), unconstrained by their coordinates in the spatial feature map.

Original: Because removing the spatial map’s coordinate structure could prevent spatial information from being available for downstream computation,
Revised: Because removing the spatial map’s coordinate structure could prevent spatial information from being available for downstream computation,

Original: we
Revised: we

Original: [Figure]
Revised: [Figure]

Original: concatenated the spatial x, y coordinates of each s x,y onto the corresponding e i to ensure it was preserved.
Revised: concatenated the spatial x, y coordinates of each s x,y onto the corresponding e i to ensure it was preserved.

Original: This feature-to-entity transformation procedure is analogous to how relation networks, visual interaction networks, and the general class of non-local neural networks (Santoro et al., 2017;
Revised: This feature-to-entity transformation procedure is analogous to how relation networks, visual interaction networks, and the general class of non-local neural networks (Santoro et al., 2017;

Original: Watters et al., 2017; Wang et al., 2017) process input images.
Revised: Watters et al., 2017; Wang et al., 2017) process input images.

Original: R ELATIONAL MODULE
Revised: R ELATIONAL MODULE

Original: To perform one step of relational reasoning, our architecture computed pairwise interactions between each entity and all others (including itself), denoted p i,j , and updated each entity by accumulating information about all of its interactions,
Revised: To perform one step of relational reasoning, our architecture computed pairwise interactions between each entity and all others (including itself), denoted p i,j , and updated each entity by accumulating information about all of its interactions,

Original: ˜ e i ← { p i,j } j =1: N .
Revised: ˜ e i ← { p i,j } j =1: N .

Original: This is analogous to how the general class of graph neural networks and their relatives perform entity-wise relational updates (Scarselli et al., 2009; Gilmer et al., 2017).
Revised: This is analogous to how the general class of graph neural networks and their relatives perform entity-wise relational updates (Scarselli et al., 2009; Gilmer et al., 2017).

Original: To perform these one-step relational computations, we used self-attention (Vaswani et al., 2017; Hoshen, 2017; Velickovic et al., 2017; Wang et al., 2017), though various other implementations (e.g., those reviewed in Battaglia et al.
Revised: To perform these one-step relational computations, we used self-attention (Vaswani et al., 2017; Hoshen, 2017; Velickovic et al., 2017; Wang et al., 2017), though various other implementations (e.g., those reviewed in Battaglia et al.

Original: (2018)) should also be viable.
Revised: (2018)) should also be viable.

Original: Our speciﬁc self-attention implementation was based on Vaswani et al.
Revised: Our speciﬁc self-attention implementation was based on Vaswani et al.

Original: (2017)’s multi-head dotproduct attention (MHDPA), which we explain in detail to make clear its role as a mechanism of relational reasoning.
Revised: (2017)’s multi-head dotproduct attention (MHDPA), which we explain in detail to make clear its role as a mechanism of relational reasoning.

Original: MHDPA (see Figure 2’s expanded box) projects the entities, E , into matrices of query, key, and value vectors, Q , K , and V , respectively (which are normalized via “layer normalization” (Ba et al., 2016)).
Revised: MHDPA (see Figure 2’s expanded box) projects the entities, E , into matrices of query, key, and value vectors, Q , K , and V , respectively (which are normalized via “layer normalization” (Ba et al., 2016)).

Original: The similarities between query q i and all keys, k j =1: N are computed by a dot-product, which are normalized into attention weights, w i , via a softmax function, and then used to compute the pairwise interaction terms, p i,j = w i,j v j .
Revised: The similarities between query q i and all keys, k j =1: N are computed by a dot-product, which are normalized into attention weights, w i , via a softmax function, and then used to compute the pairwise interaction terms, p i,j = w i,j v j .

Original: The accumulation of the interactions for each entity is the sum of these pairwise interactions, a i = (cid:80) j
Revised: The accumulation of the interactions for each entity is the sum of these pairwise interactions, a i = (cid:80) j

Original: =1: N p i,j .
Revised: =1: N p i,j .

Original: This can be efﬁciently computed using matrix multiplications, A = softmax( d − 12 QK T ) V , where d is the dimensionality of the query and key vectors (i.e. Q and K have shape N × d ).
Revised: This can be efﬁciently computed using matrix multiplications, A = softmax( d − 12 QK T ) V , where d is the dimensionality of the query and key vectors (i.e. Q and K have shape N × d ).

Original: Like Vaswani et al.
Revised: Like Vaswani et al.

Original: (2017), we also use multiple, independent attention “heads”, applied in parallel (which our attention analyses in Results 3.1 suggests the heads may assume distinct relational semantics through training).
Revised: (2017), we also use multiple, independent attention “heads”, applied in parallel (which our attention analyses in Results 3.1 suggests the heads may assume distinct relational semantics through training).

Original: The updated entities are computed as a function of the accumulated interactions, ˜ e i
Revised: The updated entities are computed as a function of the accumulated interactions, ˜ e i

Original: = g θ (cid:0) a h =1: H i (cid:1) (where h indexes the head) by concatenating all a h
Revised: = g θ (cid:0) a h =1: H i (cid:1) (where h indexes the head) by concatenating all a h

Original: =1: H i together, passing them to a multilayer perceptron (MLP), summing the output with e i (i.e. a residual connection), and applying a ﬁnal stage of layer normalization (Ba et al., 2016).
Revised: =1: H i together, passing them to a multilayer perceptron (MLP), summing the output with e i (i.e. a residual connection), and applying a ﬁnal stage of layer normalization (Ba et al., 2016).

Original: We included this non-linear transformation ( g θ ) to facilitate the computation of more complex relationships between the entities.
Revised: We included this non-linear transformation ( g θ ) to facilitate the computation of more complex relationships between the entities.

Original: This is analogous to what is done in Santoro et al.
Revised: This is analogous to what is done in Santoro et al.

Original: This one-step relational update process, which we term a “block”, can be applied iteratively using shared (recurrent) or unshared (deep) parameters, to compute higher order interactions among entities, analogous to message-passing on graphs (Li et al., 2015; Gilmer et al., 2017).
Revised: This one-step relational update process, which we term a “block”, can be applied iteratively using shared (recurrent) or unshared (deep) parameters, to compute higher order interactions among entities, analogous to message-passing on graphs (Li et al., 2015; Gilmer et al., 2017).

Original: Here we refer to the stack of multiple relational blocks as the relational module .
Revised: Here we refer to the stack of multiple relational blocks as the relational module .

Original: Underlying graph Observation
Revised: Underlying graph Observation

Original: B r an c h l eng t h =
Revised: B r an c h l eng t h =

Original: B r an c h l eng t h =        8 Environment steps 1e 1e 0.
Revised: B r an c h l eng t h =        8 Environment steps 1e 1e 0.

Original: F r a c t i on s o l v ed
Revised: F r a c t i on s o l v ed

Original: Environment steps
Revised: Environment steps

Original: F r a c t i on s o l v ed     10 12 14 0.
Revised: F r a c t i on s o l v ed     10 12 14 0.

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Baseline (
Revised: Baseline (

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: O UTPUT MODULE
Revised: O UTPUT MODULE

Original: The relational module’s ﬁnal output, ˜ E , is then used to compute π and B as follows.
Revised: The relational module’s ﬁnal output, ˜ E , is then used to compute π and B as follows.

Original: The ˜ E matrix, with shape N × f , is reduced to an f -dimensional vector by max-pooling over the entity dimension.
Revised: The ˜ E matrix, with shape N × f , is reduced to an f -dimensional vector by max-pooling over the entity dimension.

Original: This pooled vector is then passed to a small MLP, which returns an ( c + 1) -dimensional vector.
Revised: This pooled vector is then passed to a small MLP, which returns an ( c + 1) -dimensional vector.

Original: This vector is split into a c -dimensional vector of π ’s logits (where c is the number of discrete actions), and a scalar baseline value estimate, B .
Revised: This vector is split into a c -dimensional vector of π ’s logits (where c is the number of discrete actions), and a scalar baseline value estimate, B .

Original: The π logits are normalized using a softmax function, and used as probability parameters for a multinomial distribution, from which an action is randomly sampled.
Revised: The π logits are normalized using a softmax function, and used as probability parameters for a multinomial distribution, from which an action is randomly sampled.

Original:  E XPERIMENTS AND RESULTS 3. B OX -W ORLD T ASK DESCRIPTION
Revised:  E XPERIMENTS AND RESULTS 3. B OX -W ORLD T ASK DESCRIPTION

Original: We ﬁrst introduce a controlled navigation environment, called Box-World 1 , which we designed to be perceptually simple, but combinatorially complex, and require abstract relational reasoning and planning.
Revised: We ﬁrst introduce a controlled navigation environment, called Box-World 1 , which we designed to be perceptually simple, but combinatorially complex, and require abstract relational reasoning and planning.

Original: It consists of a 12 × 12 pixel room with keys and boxes scattered randomly.
Revised: It consists of a 12 × 12 pixel room with keys and boxes scattered randomly.

Original: The room also contains an agent, represented by a single dark gray pixel, which can move in four directions: up , down , left , right (see Figure 1).
Revised: The room also contains an agent, represented by a single dark gray pixel, which can move in four directions: up , down , left , right (see Figure 1).

Original: Keys are represented by a single colored pixel.
Revised: Keys are represented by a single colored pixel.

Original: The agent can pick up a loose key (i.e., one not adjacent to any other colored pixel) by walking the avatar (i.e. the visual depiction of the agent’s location) over it.
Revised: The agent can pick up a loose key (i.e., one not adjacent to any other colored pixel) by walking the avatar (i.e. the visual depiction of the agent’s location) over it.

Original: Boxes are represented by two adjacent colored pixels – the pixel on the right represents the box’s lock and its color indicates which key can be used to open that lock; the pixel on the left indicates the content of the box which is inaccessible while the box is locked.
Revised: Boxes are represented by two adjacent colored pixels – the pixel on the right represents the box’s lock and its color indicates which key can be used to open that lock; the pixel on the left indicates the content of the box which is inaccessible while the box is locked.

Original: To collect the content of a box the agent must ﬁrst collect the key that opens the box (the one that matches the lock’s color), and then walk over the lock, which makes the lock disappear.
Revised: To collect the content of a box the agent must ﬁrst collect the key that opens the box (the one that matches the lock’s color), and then walk over the lock, which makes the lock disappear.

Original: At this point the content of the box becomes accessible and can be picked up by the agent.
Revised: At this point the content of the box becomes accessible and can be picked up by the agent.

Original: Most boxes contain keys that, if made accessible, can be used to open other boxes.
Revised: Most boxes contain keys that, if made accessible, can be used to open other boxes.

Original: One of the boxes contains a gem, represented by a single white pixel.
Revised: One of the boxes contains a gem, represented by a single white pixel.

Original: The agent’s goal is to acquire the gem by unlocking the box that contains it and picking it up by walking over it.
Revised: The agent’s goal is to acquire the gem by unlocking the box that contains it and picking it up by walking over it.

Original: Keys that an agent has in its possession are depicted in the input observation as a pixel in the top-left corner.
Revised: Keys that an agent has in its possession are depicted in the input observation as a pixel in the top-left corner.

Original: In each level there is a unique sequence of boxes that need to be opened in order to reach the gem.
Revised: In each level there is a unique sequence of boxes that need to be opened in order to reach the gem.

Original: Opening one wrong box (a distractor box) consumes the held key and leads to a dead-end, where the gem cannot be reached and the level becomes unsolvable.
Revised: Opening one wrong box (a distractor box) consumes the held key and leads to a dead-end, where the gem cannot be reached and the level becomes unsolvable.

Original: There are three user-controlled parameters that contribute to the difﬁculty of the level: 1) the number of boxes in the path to the goal (solution length); 2) the number of distractor branches; 3) the length of the distractor branches.
Revised: There are three user-controlled parameters that contribute to the difﬁculty of the level: 1) the number of boxes in the path to the goal (solution length); 2) the number of distractor branches; 3) the length of the distractor branches.

Original: In general, the task is computationally difﬁcult for a few reasons.
Revised: In general, the task is computationally difﬁcult for a few reasons.

Original: First, a key can only be used once, so the agent must be able to reason about whether a particular box is along a distractor branch or along the solution path.
Revised: First, a key can only be used once, so the agent must be able to reason about whether a particular box is along a distractor branch or along the solution path.

Original: Second, keys and boxes appear in random locations in the room, demanding a capacity to reason about keys and boxes based on their abstract relations, rather than based on their spatial proximity.
Revised: Second, keys and boxes appear in random locations in the room, demanding a capacity to reason about keys and boxes based on their abstract relations, rather than based on their spatial proximity.

Original: RL A GENTS FOR B OX -W ORLD
Revised: RL A GENTS FOR B OX -W ORLD

Original: Our agent follows the reinforcement learning setup and architecture described in section (Sec. 2).
Revised: Our agent follows the reinforcement learning setup and architecture described in section (Sec. 2).

Original: For a non-relational baseline, we replaced the agent’s relational module with a variable number of residual convolution blocks.
Revised: For a non-relational baseline, we replaced the agent’s relational module with a variable number of residual convolution blocks.

Original: See the Appendix for further details of the input, relational, and output modules, including hyperparameters and training procedures.
Revised: See the Appendix for further details of the input, relational, and output modules, including hyperparameters and training procedures.

Original: R ESULTS
Revised: R ESULTS

Original: The training set-up consisted of Box-World levels with solution lengths of at least 1 and up to 4.
Revised: The training set-up consisted of Box-World levels with solution lengths of at least 1 and up to 4.

Original: This ensured that an untrained agent would have a small probability of reaching the goal by chance, at least on some levels.
Revised: This ensured that an untrained agent would have a small probability of reaching the goal by chance, at least on some levels.

Original: 2 The number of distractor branches was randomly sampled from 0 to 4.
Revised: 2 The number of distractor branches was randomly sampled from 0 to 4.

Original: Training was split into two variants of the task: one with distractor branches of length 1; another one with distractor branches of length 3 (see Figure 3).
Revised: Training was split into two variants of the task: one with distractor branches of length 1; another one with distractor branches of length 3 (see Figure 3).

Original: Agents augmented with our relational module achieved close to optimal performance in the two variants of this task, solving more than 98% of the levels.
Revised: Agents augmented with our relational module achieved close to optimal performance in the two variants of this task, solving more than 98% of the levels.

Original: In the task variant with short distractor branches, an agent with a single relational block was able to achieve top performance.
Revised: In the task variant with short distractor branches, an agent with a single relational block was able to achieve top performance.

Original: In the variant with long distractor branches, a greater number of relational blocks was required, consistent with the conjecture that more blocks allow higher-order relational computations.
Revised: In the variant with long distractor branches, a greater number of relational blocks was required, consistent with the conjecture that more blocks allow higher-order relational computations.

Original: In contrast, our baseline agents, which can only rely on convolutional and fully-connected layers, performed signiﬁcantly worse, solving less than 75% of the levels across the two task variants.
Revised: In contrast, our baseline agents, which can only rely on convolutional and fully-connected layers, performed signiﬁcantly worse, solving less than 75% of the levels across the two task variants.

Original: We observed similar results when we repeated the experiments using two alternative RL algorithms: asynchronous advantage actor-critic (A3C, Mnih et al., 2016) and distributed Q-learning with prioritized experience replay (Horgan et al., 2018).
Revised: We observed similar results when we repeated the experiments using two alternative RL algorithms: asynchronous advantage actor-critic (A3C, Mnih et al., 2016) and distributed Q-learning with prioritized experience replay (Horgan et al., 2018).

Original: We note that with Q-learning training took signiﬁcantly longer (see Figure 7 in
Revised: We note that with Q-learning training took signiﬁcantly longer (see Figure 7 in

Original: Appendix).
Revised: Appendix).

Original: We also repeated these experiments, but with backward branching in the underlying graph used to generate the level.
Revised: We also repeated these experiments, but with backward branching in the underlying graph used to generate the level.

Original: With backward branching the agent does not need to plan far into the future; when it is in possession of a key, a successful strategy is always to open the matching lock.
Revised: With backward branching the agent does not need to plan far into the future; when it is in possession of a key, a successful strategy is always to open the matching lock.

Original: In contrast, with forward branching scenes, the agent can use a key on the wrong lock (i.e. on a lock along a distractor branch).
Revised: In contrast, with forward branching scenes, the agent can use a key on the wrong lock (i.e. on a lock along a distractor branch).

Original: Thus, forward branching demands more complicated forward planning to determine the correct locks to open, which contrasts with backward branching scenarios where an agent can adopt a more reactive policy, always opting to open the lock that matches the key in possession.
Revised: Thus, forward branching demands more complicated forward planning to determine the correct locks to open, which contrasts with backward branching scenarios where an agent can adopt a more reactive policy, always opting to open the lock that matches the key in possession.

Original: Indeed, the baseline agents performed much better in backward (versus forward) branching scenes, which suggests that its weak relational reasoning capacity is to blame for its poor performance in the forward branching condition (see Figure 8 in Appendix).
Revised: Indeed, the baseline agents performed much better in backward (versus forward) branching scenes, which suggests that its weak relational reasoning capacity is to blame for its poor performance in the forward branching condition (see Figure 8 in Appendix).

Original: V ISUALIZATION OF ATTENTION WEIGHTS
Revised: V ISUALIZATION OF ATTENTION WEIGHTS

Original: We next looked at speciﬁc rows of the attention weight matrix, w i , which corresponded to relevant objects in the scene.
Revised: We next looked at speciﬁc rows of the attention weight matrix, w i , which corresponded to relevant objects in the scene.

Original: Figure 4 shows the result of this analysis when the attending entities (source of the attention) are objects along the solution path.
Revised: Figure 4 shows the result of this analysis when the attending entities (source of the attention) are objects along the solution path.

Original: For one of the attention heads, each key attends mostly to the locks that can be unlocked with that key.
Revised: For one of the attention heads, each key attends mostly to the locks that can be unlocked with that key.

Original: In other words, the attention weights reﬂect the options available to the agent once a key is collected.
Revised: In other words, the attention weights reﬂect the options available to the agent once a key is collected.

Original: For another attention head, each key attends mostly to the avatar’s entity vector.
Revised: For another attention head, each key attends mostly to the avatar’s entity vector.

Original: This suggests that object-avatar relations are important, which may help measure of relative position and support navigation.
Revised: This suggests that object-avatar relations are important, which may help measure of relative position and support navigation.

Original: In the case of RGB pixel inputs, the relationship between keys and locks that can be opened with that key is confounded with the fact that keys and the corresponding locks have the same RGB representation.
Revised: In the case of RGB pixel inputs, the relationship between keys and locks that can be opened with that key is confounded with the fact that keys and the corresponding locks have the same RGB representation.

Original: We therefore repeated the analysis, this time using one-hot representation of the input, where the mapping between keys and the corresponding locks is arbitrary.
Revised: We therefore repeated the analysis, this time using one-hot representation of the input, where the mapping between keys and the corresponding locks is arbitrary.

Original: We found evidence that: 1) keys attend to the locks they can unlock
Revised: We found evidence that: 1) keys attend to the locks they can unlock

Original: 2) locks attend to the keys that can be used to unlock them
Revised: 2) locks attend to the keys that can be used to unlock them

Original: 3) all the objects attend to the avatar’s location
Revised: 3) all the objects attend to the avatar’s location

Original: 4) the avatar gem attend to other themselves
Revised: 4) the avatar gem attend to other themselves

Original: G ENERALIZATION CAPABILITY : TESTING ON WITHHELD ENVIRONMENTS
Revised: G ENERALIZATION CAPABILITY : TESTING ON WITHHELD ENVIRONMENTS

Original: As we observed, the attention weights captured a link between a key and its corresponding lock, using a shared computation across entities.
Revised: As we observed, the attention weights captured a link between a key and its corresponding lock, using a shared computation across entities.

Original: If the function used to compute the weights (and hence, used to determine that certain keys and locks are related) has learned to represent some general, abstract notion of what it means to “unlock” – e.g., unlocks(key, lock) – then this function should be able to generalize to key-lock combinations that it has never observed during training.
Revised: If the function used to compute the weights (and hence, used to determine that certain keys and locks are related) has learned to represent some general, abstract notion of what it means to “unlock” – e.g., unlocks(key, lock) – then this function should be able to generalize to key-lock combinations that it has never observed during training.

Original: Similarly, a capacity to understand “unlocking” shouldn’t necessarily be affected by the number of locks that need to be unlocked to reach a solution.
Revised: Similarly, a capacity to understand “unlocking” shouldn’t necessarily be affected by the number of locks that need to be unlocked to reach a solution.

Original: We thus tested the model under two conditions, without further training : 1) on levels that required opening a longer sequence of boxes than it had ever observed (6, 8 and 10), and 2) on levels that required using a key-lock combination that was never required for reaching the gem during training, instead only being placed on distractor paths.
Revised: We thus tested the model under two conditions, without further training : 1) on levels that required opening a longer sequence of boxes than it had ever observed (6, 8 and 10), and 2) on levels that required using a key-lock combination that was never required for reaching the gem during training, instead only being placed on distractor paths.

Original: As shown in Figure 5, in the ﬁrst condition, the agent with the relational module solved more than 88% of the levels, across all three solution length conditions.
Revised: As shown in Figure 5, in the ﬁrst condition, the agent with the relational module solved more than 88% of the levels, across all three solution length conditions.

Original: In contrast, the performance of the agent trained without the relational module collapsed to 5% when tested on sequences of 6 boxes, and to 0% on sequences of 8 and 10.
Revised: In contrast, the performance of the agent trained without the relational module collapsed to 5% when tested on sequences of 6 boxes, and to 0% on sequences of 8 and 10.

Original: On levels with new key-lock combinations, the agent augmented with a relational module solved 97% of the new levels.
Revised: On levels with new key-lock combinations, the agent augmented with a relational module solved 97% of the new levels.

Original: The agent without the relational module performed poorly, reaching only 13%.
Revised: The agent without the relational module performed poorly, reaching only 13%.

Original: Together, these results show that the relational module offers a greater capacity for zero-shot transfer to more complex and previously unseen problems.
Revised: Together, these results show that the relational module offers a greater capacity for zero-shot transfer to more complex and previously unseen problems.

Original: This is likely a consequence of our approach’s object- and relation-centric learning, which is less sensitive to the speciﬁc conjunctions of objects and relations it has experienced during training.
Revised: This is likely a consequence of our approach’s object- and relation-centric learning, which is less sensitive to the speciﬁc conjunctions of objects and relations it has experienced during training.

Original: [Figure]
Revised: [Figure]

Original: F r a c t i on s o l v ed 0.
Revised: F r a c t i on s o l v ed 0.

Original: T r a i n .
Revised: T r a i n .

Original: T e s t
Revised: T e s t

Original: T r a i n .
Revised: T r a i n .

Original: T e s t10 

Revised: T e s t10

Original: F r a c t i on s o l v ed
Revised: F r a c t i on s o l v ed

Original: T r a i n .
Revised: T r a i n .

Original: T r a i n .10 

Revised: T r a i n .10

Original: Test
Revised: Test

Original: Test 0.
Revised: Test 0.

Original: Not required during training
Revised: Not required during training

Original: Withheld key
Revised: Withheld key

Original: [Figure]
Revised: [Figure]

Original: Relational
Revised: Relational

Original: Baseline Relational
Revised: Baseline Relational

Original: Baseline
Revised: Baseline

Original: 3. S TAR C RAFT II MINI - GAMES T ASK DESCRIPTION
Revised: 3. S TAR C RAFT II MINI - GAMES T ASK DESCRIPTION

Original: StarCraft II is a popular video game that presents a difﬁcult challenge for RL agents.
Revised: StarCraft II is a popular video game that presents a difﬁcult challenge for RL agents.

Original: It is a multiagent game where each player controls a large number (hundreds) of units that need to interact and collaborate (see Figure 1).
Revised: It is a multiagent game where each player controls a large number (hundreds) of units that need to interact and collaborate (see Figure 1).

Original: It is partially observable and has a large state and action space, with more than 100 possible actions.
Revised: It is partially observable and has a large state and action space, with more than 100 possible actions.

Original: The consequences of any single action – in particular, early decisions in the game – are typically only observed many frames later, posing difﬁculties in temporal credit assignment and exploration.
Revised: The consequences of any single action – in particular, early decisions in the game – are typically only observed many frames later, posing difﬁculties in temporal credit assignment and exploration.

Original: We trained our agents on the suite of 7 mini-games developed for the StarCraft II Learning Environment (SC2LE, Vinyals et al., 2017).
Revised: We trained our agents on the suite of 7 mini-games developed for the StarCraft II Learning Environment (SC2LE, Vinyals et al., 2017).

Original: These mini-games were proposed as a set of speciﬁc scenarios that are representative of the mechanics of the full game and can be used to test agents in a simpler set up with a better deﬁned reward structure, compared to the full game.
Revised: These mini-games were proposed as a set of speciﬁc scenarios that are representative of the mechanics of the full game and can be used to test agents in a simpler set up with a better deﬁned reward structure, compared to the full game.

Original: The 7 mini-games demand different strategies and skills.
Revised: The 7 mini-games demand different strategies and skills.

Original: Notably, Collect Mineral Shards (CMS), Defeat Roaches (DR) and
Revised: Notably, Collect Mineral Shards (CMS), Defeat Roaches (DR) and

Original: Defeat Zerglings and Banelings (DZB), require controlling multiple units with precision and in a coordinated manner.
Revised: Defeat Zerglings and Banelings (DZB), require controlling multiple units with precision and in a coordinated manner.

Original: To achieve a high score in CMS, the player (or agent) needs to command the two marines independently, deﬁning navigation paths that do not overlap too much in order to efﬁciently cover most of the ground.
Revised: To achieve a high score in CMS, the player (or agent) needs to command the two marines independently, deﬁning navigation paths that do not overlap too much in order to efﬁciently cover most of the ground.

Original: In DR, the player must focus-ﬁre to succeed, which requires selecting the same attack target, one at a time, for all the marines.
Revised: In DR, the player must focus-ﬁre to succeed, which requires selecting the same attack target, one at a time, for all the marines.

Original: In DZB, the enemy units have different abilities, so a good strategy may involve splitting the army into groups with different roles.
Revised: In DZB, the enemy units have different abilities, so a good strategy may involve splitting the army into groups with different roles.

Original: RL A GENTS FOR S TAR C RAFT II
Revised: RL A GENTS FOR S TAR C RAFT II

Original: For StarCraft II (SC2), we start by ﬁrst constructing a strong control baseline to test against our relational agent.
Revised: For StarCraft II (SC2), we start by ﬁrst constructing a strong control baseline to test against our relational agent.

Original: The architecture is similar to that of the Box-World agent, with changes to accommodate speciﬁc features of the SC2 environment.
Revised: The architecture is similar to that of the Box-World agent, with changes to accommodate speciﬁc features of the SC2 environment.

Original: We increased the model capacity by using 2 residual convolutional blocks in the input module, each consisting of 3 convolutional layers.
Revised: We increased the model capacity by using 2 residual convolutional blocks in the input module, each consisting of 3 convolutional layers.

Original: We added a 2D-ConvLSTM immediately downstream of the residual blocks, so the subsequent computations were sensitive to the recent history of observations, rather than only the current observation, thus accounting for partial observability.
Revised: We added a 2D-ConvLSTM immediately downstream of the residual blocks, so the subsequent computations were sensitive to the recent history of observations, rather than only the current observation, thus accounting for partial observability.

Original: Including this memory mechanism was important because the actions a StarCraft agent issues to units are carried out over multiple timesteps (e.g., MOVE UNIT
Revised: Including this memory mechanism was important because the actions a StarCraft agent issues to units are carried out over multiple timesteps (e.g., MOVE UNIT

Original: TO COORDINATE (1, 5)), and the agent needs a way of not “forgetting” what actions it has issued previously.
Revised: TO COORDINATE (1, 5)), and the agent needs a way of not “forgetting” what actions it has issued previously.

Original: Otherwise it might keep reissuing the same action to a unit, rather than choosing another action for another unit.
Revised: Otherwise it might keep reissuing the same action to a unit, rather than choosing another action for another unit.

Original: Finally, the relational module depicted in section (Sec. 2) was replaced by a stack of residual convolutional blocks in the control agent.
Revised: Finally, the relational module depicted in section (Sec. 2) was replaced by a stack of residual convolutional blocks in the control agent.

Original: For the agent’s output, alongside action a and value V , the network produces two sets of action-related arguments: non-spatial arguments ( Args ) and spatial arguments ( Args x,y ).
Revised: For the agent’s output, alongside action a and value V , the network produces two sets of action-related arguments: non-spatial arguments ( Args ) and spatial arguments ( Args x,y ).

Original: These arguments are used as modiﬁers of particular actions (see (Vinyals et al., 2017)).
Revised: These arguments are used as modiﬁers of particular actions (see (Vinyals et al., 2017)).

Original: Args are produced from the output of the aggregation function, whereas Args x,y result from upsampling the output of the relational module.
Revised: Args are produced from the output of the aggregation function, whereas Args x,y result from upsampling the output of the relational module.

Original: See the Appendix for further details.
Revised: See the Appendix for further details.

Original: R ESULTS
Revised: R ESULTS

Original: For these results we used the full action set provided by SC2LE and performance was measured as the mean score over 30 episodes  for each mini-game.
Revised: For these results we used the full action set provided by SC2LE and performance was measured as the mean score over 30 episodes of the best run for each mini-game.

Original: Our agent implementations achieved high scores across all the mini-games (Figure 6), consistently outperforming the previous best models.
Revised: Our agent implementations achieved high scores across all the mini-games (Figure 6), consistently outperforming the previous best models.

Original: In particular, the agent augmented with a relational module achieved state-of-the-art results in six mini-games and its performance surpassed that of the human grandmaster in four of them.
Revised: In particular, the agent augmented with a relational module achieved state-of-the-art results in six mini-games and its performance surpassed that of the human grandmaster in four of them.

Original: Head-to-head comparisons between our two implementations show that the agent with the relational component (relational) achieves similar or better results than the one without (control).
Revised: Head-to-head comparisons between our two implementations show that the agent with the relational component (relational) achieves similar or better results than the one without (control).

Original: Concretely, for Collect Mineral Shards, Defeat Roaches and Defeat Zerglings and Banelings, the relational agent improved relative to the control agent by 9 , 8 and 134 points on average, with a standard error of 1.5, 6.4 and 18.2, respectively.
Revised: Concretely, for Collect Mineral Shards, Defeat Roaches and Defeat Zerglings and Banelings, the relational agent improved relative to the control agent by 9 , 8 and 134 points, with a standard error of 1.5, 6.4 and 18.2, respectively.

Original: We note that both models improved substantially over the previous best (Vinyals et al., 2017), and in some mini-games, they approach an empirical ceiling performance.
Revised: We note that both models improved substantially over the previous best (Vinyals et al., 2017), and in some mini-games, they approach an empirical ceiling performance.

Original: This can be attributed to a combination of factors: improved RL algorithm (Espeholt et al., 2018), more robust architecture, better hyperparameter tuning to address issues of credit assignment and exploration, better action selection procedure, and longer training.
Revised: This can be attributed to a combination of factors: improved RL algorithm (Espeholt et al., 2018), more robust architecture, better hyperparameter tuning to address issues of credit assignment and exploration, better action selection procedure, and longer training.

Original: Next, we focus on differences afforded by relational inductive biases and turn to particular generalization tests to determine the behavioural traits of the control and relational agents.
Revised: Next, we focus on differences afforded by relational inductive biases and turn to particular generalization tests to determine the behavioural traits of the control and relational agents.

Original: G ENERALIZATION CAPABILITY
Revised: G ENERALIZATION CAPABILITY

Original: As observed in Box-World, a capacity to better understand underlying relational structure – rather than latch onto superﬁcial statistics – may manifest in better generalization to never-before-seen situations.
Revised: As observed in Box-World, a capacity to better understand underlying relational structure – rather than latch onto superﬁcial statistics – may manifest in better generalization to never-before-seen situations.

Original: To test generalization in SC2 we took agents trained on Collect Mineral Shards, which involved using two marines to collect randomly scattered minerals and tested them, without further training, on modiﬁed levels that allowed the agents to instead control up to ten marines.
Revised: To test generalization in SC2 we took agents trained on Collect Mineral Shards, which involved using two marines to collect randomly scattered minerals and tested them, without further training, on modiﬁed levels that allowed the agents to instead control up to ten marines.

Original: It’s worth highlighting the agents have never been exposed to a single observation present in these generalization experiments.
Revised: It’s worth highlighting the agents have never been exposed to a single observation present in these generalization experiments.

Original: Intuitively, if an agent understands that marines are independent units that can be coordinated, yet controlled independently to collect resources, then increasing the number of marines available should only affect the underlying strategy of unit deployment, and should not catastrophically break model performance.
Revised: Intuitively, if an agent understands that marines are independent units that can be coordinated, yet controlled independently to collect resources, then increasing the number of marines available should only affect the underlying strategy of unit deployment, and should not catastrophically break model performance.

Original: We observed that—at least for medium sized networks—some interesting generalization capabilities emerge, with the best seeds of the relational agent achieving better generalization scores in the test scenario.
Revised: We observed that—at least for medium sized networks—some interesting generalization capabilities emerge, with the best seeds of the relational agent achieving better generalization scores in the test scenario.

Original: However, we noticed high variability in these quantitative results, with the effect diminishing when using larger models (which may be more prone to overﬁtting on the training set).
Revised: However, we noticed high variability in these quantitative results, with the effect diminishing when using larger models (which may be more prone to overﬁtting on the training set).

Original: Qualitative analysis of the policies revealed distinct behaviours for the best performing control and relational agents: while the former adopted a "land sweep strategy", controlling many units as a group to cover the space, the latter managed to independently control several units simultaneously, suggesting a ﬁner grained understanding of the game dynamics.
Revised: Qualitative analysis of the policies revealed distinct behaviours for the best performing control and relational agents: while the former adopted a "land sweep strategy", controlling many units as a group to cover the space, the latter managed to independently control several units simultaneously, suggesting a ﬁner grained understanding of the game dynamics.

Original: Further work is required to draw ﬁrm conclusions about the generalization capabilities of a relational agent in more complex domains such as StarCraft II (see Figure 9 in Appendix).
Revised: Further work is required to draw ﬁrm conclusions about the generalization capabilities of a relational agent in more complex domains such as StarCraft II (see Figure 9 in Appendix).

Original: Given the combinatoric richness and multi-agent aspects of the full StarCraft II game, an agent is frequently exposed to situations in which it might not have been trained on.
Revised: Given the combinatoric richness and multi-agent aspects of the full StarCraft II game, an agent is frequently exposed to situations in which it might not have been trained on.

Original: Thus, an improved capability to generalize to new, unseen situations aided by a better understanding of underlying abstract entities and their relations is fundamental.
Revised: Thus, an improved capability to generalize to new, unseen situations aided by a better understanding of underlying abstract entities and their relations is fundamental.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: By introducing structured perception and relational reasoning into deep RL architectures, our agents can learn interpretable representations, and exceed baseline agents in terms of sample complexity, ability to generalize, and overall performance.
Revised: By introducing structured perception and relational reasoning into deep RL architectures, our agents can learn interpretable representations, and exceed baseline agents in terms of sample complexity, ability to generalize, and overall performance.

Original: Behavioral analyses showed that the learned representations allowed for better generalization, which is characteristic of relational representations.
Revised: Behavioral analyses showed that the learned representations allowed for better generalization, which is characteristic of relational representations.

Original: Analysis of attention weights showed that the model’s internal computations were interpretable, and congruent with the computations we would expect from a model computing task-relevant relations.
Revised: Analysis of attention weights showed that the model’s internal computations were interpretable, and congruent with the computations we would expect from a model computing task-relevant relations.

Original: One important future direction is to explore ways to scale our approach to larger inputs spaces, without suffering, as this and other approaches do (e.g., Wang et al., 2017; Santoro et al., 2017), from the quadratic complexity that results from considering all input pairs.
Revised: One important future direction is to explore ways to scale our approach to larger inputs spaces, without suffering, as this and other approaches do (e.g., Wang et al., 2017; Santoro et al., 2017), from the quadratic complexity that results from considering all input pairs.

Original: Possible avenues involve using a distinct attentional mechanisms that scales linearly with the number of inputs (Hoshen, 2017) or ﬁltering out unimportant relations (Malinowski et al., 2018).
Revised: Possible avenues involve using a distinct attentional mechanisms that scales linearly with the number of inputs (Hoshen, 2017) or ﬁltering out unimportant relations (Malinowski et al., 2018).

Original: Other future directions include exploring perceiving complex scenes via more structured formats, such as scene graphs (Xu et al., 2017;
Revised: Other future directions include exploring perceiving complex scenes via more structured formats, such as scene graphs (Xu et al., 2017;

Original: Chen et al., 2018), which could be powerful additions to our approach’s input module.
Revised: Chen et al., 2018), which could be powerful additions to our approach’s input module.

Original: More complex relational modules could be explored, such as richer graph network implementations (Battaglia et al., 2018), learned approaches for inducing compositional programs (Reed & De Freitas, 2015;
Revised: More complex relational modules could be explored, such as richer graph network implementations (Battaglia et al., 2018), learned approaches for inducing compositional programs (Reed & De Freitas, 2015;

Original: Parisotto et al., 2017; Allamanis et al., 2017; Devlin et al., 2017) and reasoning about structured data (Neelakantan et al., 2015; Liang et al., 2016), or even explicit logical reasoning over structured internal representations (Evans et al., 2018), drawing inspiration from more symbolic approaches in classic AI.
Revised: Parisotto et al., 2017; Allamanis et al., 2017; Devlin et al., 2017) and reasoning about structured data (Neelakantan et al., 2015; Liang et al., 2016), or even explicit logical reasoning over structured internal representations (Evans et al., 2018), drawing inspiration from more symbolic approaches in classic AI.

Original: Our approach may also interface well with approaches for hierarchical RL (Vezhnevets et al., 2017), planning (Guez et al., 2018), and structured behavior representation (Huang et al., 2018), so that the structured internal representations and patterns of reasoning can translate into more structured behaviors.
Revised: Our approach may also interface well with approaches for hierarchical RL (Vezhnevets et al., 2017), planning (Guez et al., 2018), and structured behavior representation (Huang et al., 2018), so that the structured internal representations and patterns of reasoning can translate into more structured behaviors.

Original: More speculatively, this work blurs the line between model-free agents, and those with a capacity for more abstract planning.
Revised: More speculatively, this work blurs the line between model-free agents, and those with a capacity for more abstract planning.

Original: An important feature of model-based approaches is making general knowledge of the environment available for decision-making.
Revised: An important feature of model-based approaches is making general knowledge of the environment available for decision-making.

Original: Here our inductive biases for entityand relation-centric representations and iterated reasoning reﬂect key knowledge about the structure of the world.
Revised: Here our inductive biases for entityand relation-centric representations and iterated reasoning reﬂect key knowledge about the structure of the world.

Original: While not a model in the technical sense, it is possible that the agent learns to exploit this relational architectural prior similarly to how an imagination-based agent’s forward model operates (Hamrick et al., 2017;
Revised: While not a model in the technical sense, it is possible that the agent learns to exploit this relational architectural prior similarly to how an imagination-based agent’s forward model operates (Hamrick et al., 2017;

Original: Pascanu et al., 2017; Weber et al., 2017).
Revised: Pascanu et al., 2017; Weber et al., 2017).

Original: More generally, our work opens new directions for RL via a principled hybrid of ﬂexible statistical learning and more structured approaches.
Revised: More generally, our work opens new directions for RL via a principled hybrid of ﬂexible statistical learning and more structured approaches.

Original: A PPENDIX A B OX - WORLD T ASK
Revised: A PPENDIX A B OX - WORLD T ASK

Original: Each level in Box-world is procedurally generated.
Revised: Each level in Box-world is procedurally generated.

Original: We start by generating a random graph (a tree) that deﬁnes the correct path to the goal – i
Revised: We start by generating a random graph (a tree) that deﬁnes the correct path to the goal – i

Original: [Figure]
Revised: [Figure]

Original: e., the sequence of boxes that need to be opened to reach the gem.
Revised: e., the sequence of boxes that need to be opened to reach the gem.

Original: This graph also deﬁnes multiple distractor branches – boxes that lead to dead-ends.
Revised: This graph also deﬁnes multiple distractor branches – boxes that lead to dead-ends.

Original: The agent, keys and boxes, including the one containing the gem, are positioned randomly in the room, assuring that there is enough space for the agent to navigate between boxes.
Revised: The agent, keys and boxes, including the one containing the gem, are positioned randomly in the room, assuring that there is enough space for the agent to navigate between boxes.

Original: There is a total ofkeys and 20 locks that are randomly sampled to produce the level.
Revised: There is a total ofkeys and 20 locks that are randomly sampled to produce the level.

Original: An agent receives a reward of + for collecting the gem, +1 for opening a box in the solution path and − 1 for opening a distractor box.
Revised: An agent receives a reward of + for collecting the gem, +1 for opening a box in the solution path and − 1 for opening a distractor box.

Original: A level terminates immediately after the gem is collected or a distractor box is opened.
Revised: A level terminates immediately after the gem is collected or a distractor box is opened.

Original: The generation process produces a very large number of possible trees, making it extremely unlikely that the agent will face the same level twice.
Revised: The generation process produces a very large number of possible trees, making it extremely unlikely that the agent will face the same level twice.

Original: The procedural generation of levels also allows us to create different training-test splits by withholding levels that conform to a particular case during training and presenting them to the agent at test time.
Revised: The procedural generation of levels also allows us to create different training-test splits by withholding levels that conform to a particular case during training and presenting them to the agent at test time.

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: v ed
Revised: v ed

Original: F r a c t
Revised: F r a c t

Original: i
Revised: i

Original: [Figure]
Revised: [Figure]

Original: l v ed
Revised: l v ed

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: B RL TRAINING PROCEDURE
Revised: B RL TRAINING PROCEDURE

Original: We used distributed A2C agents with off-policy corrections (Espeholt et al., 2018).
Revised: We used distributed A2C agents with off-policy corrections (Espeholt et al., 2018).

Original: Each agents consisted of 100 actors generating trajectories of experience
Revised: Each agents consisted of 100 actors generating trajectories of experience

Original: [Table]
Revised: [Table]

Original: and a single learner, which learns π and
Revised: and a single learner, which learns π and

Original: B using the actors’ experiences.
Revised: B using the actors’ experiences.

Original: The model updates were performed on GPU using mini-batches of 32 trajectories provided by the actors via a queue.
Revised: The model updates were performed on GPU using mini-batches of 32 trajectories provided by the actors via a queue.

Original: The agents used an entropy cost of 0 .
Revised: The agents used an entropy cost of 0 .

Original: 005 , discount ( γ ) of 0 .
Revised: 005 , discount ( γ ) of 0 .

Original: 99 and unroll length of 40 steps.
Revised: 99 and unroll length of 40 steps.

Original: Training was done using RMSprop optimiser with momentum of 0, (cid:15) of 0 . 1 and a decay term of 0 . 99 .
Revised: Training was done using RMSprop optimiser with momentum of 0, (cid:15) of 0 . 1 and a decay term of 0 . 99 .

Original: The learning rate was tuned, taking values between 1e − 5 and 2e − 4 .
Revised: The learning rate was tuned, taking values between 1e − 5 and 2e − 4 .

Original: A GENT ARCHITECTURE
Revised: A GENT ARCHITECTURE

Original: The input module contained two convolutional layers with 12 and 24 kernels, 2 × 2 kernel sizes and a stride of 1 , followed by a rectiﬁed linear unit (ReLU) activation function.
Revised: The input module contained two convolutional layers with 12 and 24 kernels, 2 × 2 kernel sizes and a stride of 1 , followed by a rectiﬁed linear unit (ReLU) activation function.

Original: The output was tagged with two extra channels indicating the spatial position ( x and y ) of each cell in the feature map using evenly spaced values between − 1 and 1 .
Revised: The output was tagged with two extra channels indicating the spatial position ( x and y ) of each cell in the feature map using evenly spaced values between − 1 and 1 .

Original: This was passed to the relational module, consisting of relational blocks, with shared parameters .
Revised: This was passed to the relational module, consisting of relational blocks, with shared parameters .

Original: Queries, keys and values were produced by 2 toattention heads and had an embedding size ( d ) of 64.
Revised: Queries, keys and values were produced by 2 toattention heads and had an embedding size ( d ) of 64.

Original: The output of this module was aggregated using a feature-wise max pooling function and passed to 4 fully connected layers, each followed by a ReLU.
Revised: The output of this module was aggregated using a feature-wise max pooling function and passed to 4 fully connected layers, each followed by a ReLU.

Original: Policy logits ( π , size 4) and baseline function ( B , size 1) were produced by a linear projection.
Revised: Policy logits ( π , size 4) and baseline function ( B , size 1) were produced by a linear projection.

Original: The policy logits were normalized and used as multinomial distribution from which the action ( a ) was sampled.
Revised: The policy logits were normalized and used as multinomial distribution from which the action ( a ) was sampled.

Original: B ASELINE AGENT ARCHITECTURE
Revised: B ASELINE AGENT ARCHITECTURE

Original: As a baseline agent we used the same architecture as the relational agent but replaced the relational module with a variable number (3 to 6) of residual-convolutional blocks.
Revised: As a baseline agent we used the same architecture as the relational agent but replaced the relational module with a variable number (3 to 6) of residual-convolutional blocks.

Original: Each residual block comprised two convolutional layers, with 3 × 3 kernels, stride of 1 and 26 output channels.
Revised: Each residual block comprised two convolutional layers, with 3 × 3 kernels, stride of 1 and 26 output channels.

Original: C S TAR C RAFT II MINI - GAMES
Revised: C S TAR C RAFT II MINI - GAMES

Original: StarCraft II agents were trained with Adam optimiser for a total of 10 billion steps using batches of 32 trajectories, each unrolled for 80 steps.
Revised: StarCraft II agents were trained with Adam optimiser for a total of 10 billion steps using batches of 32 trajectories, each unrolled for 80 steps.

Original: A linear decay was applied to the optimiser learning rate and entropy loss scaling throughout training (see Table 2 for details).
Revised: A linear decay was applied to the optimiser learning rate and entropy loss scaling throughout training (see Table 2 for details).

Original: We ran approximatelyexperiments for each mini-game, using the hyperparameter settings indicated in Table 4 combined with 3 seeds.
Revised: We ran approximatelyexperiments for each mini-game, using the hyperparameter settings indicated in Table 4 combined with 3 seeds.

Original: [Figure]
Revised: [Figure]

Original: R ELATIONAL A GENT ARCHITECTURE
Revised: R ELATIONAL A GENT ARCHITECTURE

Original: The StarCraft II (SC2) agent architecture follows closely the one we adopted in Box-World.
Revised: The StarCraft II (SC2) agent architecture follows closely the one we adopted in Box-World.

Original: Here we highlight the changes needed to satisfy SC2 constraints.
Revised: Here we highlight the changes needed to satisfy SC2 constraints.

Original: Input-preprocessing .
Revised: Input-preprocessing .

Original: At each time step agents are presented with 4 sources of information: minimap , screen , player , and previous-action .
Revised: At each time step agents are presented with 4 sources of information: minimap , screen , player , and previous-action .

Original: These tensors share the same pre-processing: numerical features are re-scaled with a logarithmic transformation and categorical features are embedded into a continuous 10-dimensional space.
Revised: These tensors share the same pre-processing: numerical features are re-scaled with a logarithmic transformation and categorical features are embedded into a continuous 10-dimensional space.

Original: State encoding .
Revised: State encoding .

Original: Spatially encoded inputs ( minimap and screen ) are tiled with binary masks denoting whether the previous action constituted a screen- or minimap-related action.
Revised: Spatially encoded inputs ( minimap and screen ) are tiled with binary masks denoting whether the previous action constituted a screen- or minimap-related action.

Original: These tensors are then fed to independent residual convolutional blocks, each consisting of one convolutional layer ( 4 ×kernels and stride 2) followed by a residual block with 2 convolutional layers ( 3 × 3 kernels and stride 1), which process and downsample the inputs to [ 8 × 8 ×
Revised: These tensors are then fed to independent residual convolutional blocks, each consisting of one convolutional layer ( 4 ×kernels and stride 2) followed by a residual block with 2 convolutional layers ( 3 × 3 kernels and stride 1), which process and downsample the inputs to [ 8 × 8 ×

Original: # channels 1 ] outputs.
Revised: # channels 1 ] outputs.

Original: These tensors are concatenated along the depth dimension to form a singular spatial input ( inputs 3 D , with shape [ 8 × 8 × # channels 1 + # channels 1 ]).
Revised: These tensors are concatenated along the depth dimension to form a singular spatial input ( inputs 3 D , with shape [ 8 × 8 × # channels 1 + # channels 1 ]).

Original: The remaining inputs ( player and previous-action ) are concatenated and passed to a 2-layer MLP (128 units, ReLU, 64 units) to form a singular non-spatial input ( inputs 2 D ).
Revised: The remaining inputs ( player and previous-action ) are concatenated and passed to a 2-layer MLP (128 units, ReLU, 64 units) to form a singular non-spatial input ( inputs 2 D ).

Original: Memory processing .
Revised: Memory processing .

Original: Next, inputs 3 D is passed to the Conv2DLSTM along with its previous state to produce a new state and outputs 3 D (shape [ 8 × 8 × # channels 2 ]), which represents an aggregated history of input observations.
Revised: Next, inputs 3 D is passed to the Conv2DLSTM along with its previous state to produce a new state and outputs 3 D (shape [ 8 × 8 × # channels 2 ]), which represents an aggregated history of input observations.

Original: Relational processing .
Revised: Relational processing .

Original: outputs 3 D is ﬂattened along the ﬁrst two dimensions (forming a 2D tensor of shape [ 64 × # channels 2 ] and passed to the stacked MHDPA blocks (see Table 3 for details).
Revised: outputs 3 D is ﬂattened along the ﬁrst two dimensions (forming a 2D tensor of shape [ 64 × # channels 2 ] and passed to the stacked MHDPA blocks (see Table 3 for details).

Original: Its output tensors (of shape [ 64 × attention embedding size ∗ number of attention heads ]) follow two separate pathways – relational-spatial : reshapes the tensors to their original spatial shape [ 8 × 8 ]; relational-nonspatial : aggregates through a feature-wise max-pooling operation (combining theembeddings into a ﬂat tensor) and further processes using a 2-layer MLP (512 units per layer, ReLU activations).
Revised: Its output tensors (of shape [ 64 × attention embedding size ∗ number of attention heads ]) follow two separate pathways – relational-spatial : reshapes the tensors to their original spatial shape [ 8 × 8 ]; relational-nonspatial : aggregates through a feature-wise max-pooling operation (combining theembeddings into a ﬂat tensor) and further processes using a 2-layer MLP (512 units per layer, ReLU activations).

Original: Output processing .
Revised: Output processing .

Original: inputs 2 D and relational-nonspatial are concatenated to form a set of shared features .
Revised: inputs 2 D and relational-nonspatial are concatenated to form a set of shared features .

Original: Policy logits are produced by feeding shared features to a 2-layer MLP (256 units, ReLU, | actions | units) and masking unavailable actions (following Vinyals et al.
Revised: Policy logits are produced by feeding shared features to a 2-layer MLP (256 units, ReLU, | actions | units) and masking unavailable actions (following Vinyals et al.

Original: Similarly, baselines values V are generated by feeding shared features to a separate 2-layer MLP (256 units, ReLU,unit).
Revised: Similarly, baselines values V are generated by feeding shared features to a separate 2-layer MLP (256 units, ReLU,unit).

Original: Actions are sampled using computed policy logits and embedded into a 16 dimensional vector.
Revised: Actions are sampled using computed policy logits and embedded into a 16 dimensional vector.

Original: This embedding is used to condition shared features and generate logits for non-spatial arguments ( Args ) through independent linear combinations (one for each argument).
Revised: This embedding is used to condition shared features and generate logits for non-spatial arguments ( Args ) through independent linear combinations (one for each argument).

Original: Finally, spatial arguments ( Args x,y ) are obtained by ﬁrst deconvolving relational-spatial to [ 32 × 32 × # channels 3 ] tensors using Conv2DTranspose layers, conditioned by tiling the action embedding along the depth dimension
Revised: Finally, spatial arguments ( Args x,y ) are obtained by ﬁrst deconvolving relational-spatial to [ 32 × 32 × # channels 3 ] tensors using Conv2DTranspose layers, conditioned by tiling the action embedding along the depth dimension

Original: [Table] [Table]
Revised: [Table] [Table]

Original: and passed by 1 × 1 × 1 convolution layers (one for each spatial argument).
Revised: and passed by 1 × 1 × 1 convolution layers (one for each spatial argument).

Original: Spatial arguments ( x , y ) are produced by sampling resulting tensors and selecting the corresponding row and column
Revised: Spatial arguments ( x , y ) are produced by sampling resulting tensors and selecting the corresponding row and column

Original: C ONTROL AGENT ARCHITECTURE
Revised: C ONTROL AGENT ARCHITECTURE

Original: The control agent architecture only differs on the relational processing part of the pipeline.
Revised: The control agent architecture only differs on the relational processing part of the pipeline.

Original: Analogous to the relational agent, outputs 2 D are obtained from Conv2DLSTM layers.
Revised: Analogous to the relational agent, outputs 2 D are obtained from Conv2DLSTM layers.

Original: These tensors are ﬁrst passed to a 12-layer deep residual model – comprising 4 blocks of 3 convolutions layers (32 output channels, 4 × 4 kernel for the ﬁrst convolution and 3 × 3 for the second and third, and stride 1) interleaved with ReLU activations and skip-connections – as proposed by He et al.
Revised: These tensors are ﬁrst passed to a 12-layer deep residual model – comprising 4 blocks of 3 convolutions layers (32 output channels, 4 × 4 kernel for the ﬁrst convolution and 3 × 3 for the second and third, and stride 1) interleaved with ReLU activations and skip-connections – as proposed by He et al.

Original: (2016), to form the relational-spatial outputs.
Revised: (2016), to form the relational-spatial outputs.

Original: These tensors also follow a separate pathway where they are ﬂattened and passed to a 2-layer MLP (512 units per layer, ReLU activations) to produce what we refer to above as relational-nonspatial .
Revised: These tensors also follow a separate pathway where they are ﬂattened and passed to a 2-layer MLP (512 units per layer, ReLU activations) to produce what we refer to above as relational-nonspatial .

Original: The remaining architecture is identical to the relational agent.
Revised: The remaining architecture is identical to the relational agent.

Original: [Table]
Revised: [Table]