Deep reinforcement learning with relational inductive biases

1 INTRODUCTION . 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 flexibility in how they learn and exploit the statistical structure underlying observations and reward signals . But the downsides to such flexibility often include low sample efficiency and poor transfer beyond the specifics of the training environment ( Zhang et al. , 2018 ; Lake et al. , 2017 ; Kansky et al. , 2017 ) . Various structured approaches to RL ( e.g . Dzeroski et al . ( 2001 ) ; Driessens & Dzeroski ( 2004 ) ; Diuk et al . ( 2008 ) ; Garnelo et al . ( 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 . Although these approaches are often highly efficient , they constrain the representations and admissible learning algorithms , they struggle to learn rich representations , and they are therefore confined to relatively simple tasks and data conditions . To strike favorable tradeoffs between flexibility and efficiency , a number of recent approaches have explored using relational inductive biases in deep learning , to reap the benefits of flexible statistical learning and more structured approaches . Methods such as “ graph networks ” ( Scarselli et al. , 2009 ; 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 . Because they are implemented using deep neural networks , they can learn transformations from input ∗Equal contribution . observations into task-relevant entities , as well as functions for computing rich interaction among these entities . 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 ; Raposo et al. , 2017 ; Dai et al. , 2017 ; Wang et al. , 2018 ; Selsam et al. , 2018 ; Hamrick et al. , 2018 ; Kipf et al. , 2018 ; Sanchez-Gonzalez et al. , 2018 ) . 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 . In contrast with prior work exploring relational inductive biases in deep RL ( e.g. , 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 . 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 ) . 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 ; Gilmer et al. , 2017 ) . 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 ) . 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 . 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 . Our agent reaches higher ceiling performance , more efficiently , 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 . 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 . 2 RELATIONAL DEEP RL AGENT ARCHITECTURE . RL ALGORITHM . 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 . 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 ” ) . 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 . 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 . For further algorithmic details , including the distributed actor/learner framework and off-policy correction methods , please see Espeholt et al . ( 2018 ) . INPUT MODULE . Our agent takes as input an image of a scene , transforms it into an embedded state representation , S , which is a spatial feature map returned by a convolutional neural network ( CNN ) . 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 , ei , corresponded to a feature vector , sx , y , at a particular x , y location in the embedded scene . Crucially , this allowed for non-local computation between entities ( Wang et al. , 2017 ) , unconstrained by their coordinates in the spatial feature map . Because removing the spatial map ’ s coordinate structure could prevent spatial information from being available for downstream computation , we concatenated the spatial x , y coordinates of each sx , y onto the corresponding ei to ensure it was preserved . 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 ; Watters et al. , 2017 ; Wang et al. , 2017 ) process input images . RELATIONAL MODULE . To perform one step of relational reasoning , our architecture computed pairwise interactions between each entity and all others ( including itself ) , denoted pi , j , and updated each entity by accumulating information about all of its interactions , ẽi ← { pi , j } j=1 : N . 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 ) . 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 . ( 2018 ) ) should also be viable . Our specific self-attention implementation was based on Vaswani et al . ( 2017 ) ’ s multi-head dotproduct attention ( MHDPA ) , which we explain in detail to make clear its role as a mechanism of relational reasoning . 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 ) ) . The similarities between query qi and all keys , kj=1 : N are computed by a dot-product , which are normalized into attention weights , wi , via a softmax function , and then used to compute the pairwise interaction terms , pi , j = wi , jvj . The accumulation of the interactions for each entity is the sum of these pairwise interactions , ai = ∑ j=1 : N pi , j . This can be efficiently computed using matrix multiplications , A = softmax ( d− 1 2QKT ) V , where d is the dimensionality of the query and key vectors ( i.e . Q and K have shape N × d ) . Like Vaswani et al . ( 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 ) . The updated entities are computed as a function of the accumulated interactions , ẽi = gθ ( ah=1 : Hi ) ( where h indexes the head ) by concatenating all ah=1 : Hi together , passing them to a multilayer perceptron ( MLP ) , summing the output with ei ( i.e . a residual connection ) , and applying a final stage of layer normalization ( Ba et al. , 2016 ) . We included this non-linear transformation ( gθ ) to facilitate the computation of more complex relationships between the entities . This is analogous to what is done in Santoro et al . ( 2017 ) . 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 ) . Here we refer to the stack of multiple relational blocks as the relational module . OUTPUT MODULE . The relational module ’ s final output , Ẽ , is then used to compute π and B as follows . The Ẽ matrix , with shape N × f , is reduced to an f -dimensional vector by max-pooling over the entity dimension . This pooled vector is then passed to a small MLP , which returns an ( c+ 1 ) -dimensional vector . 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 . The π logits are normalized using a softmax function , and used as probability parameters for a multinomial distribution , from which an action is randomly sampled . 3 EXPERIMENTS AND RESULTS . 3.1 BOX-WORLD . TASK DESCRIPTION . We first introduce a controlled navigation environment , called Box-World1 , which we designed to be perceptually simple , but combinatorially complex , and require abstract relational reasoning and planning . It consists of a 12× 12 pixel room with keys and boxes scattered randomly . 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 ) . Keys are represented by a single colored pixel . 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 . Boxes are represented by two adjacent colored pixels – the pixel on the right 1The Box-World environment will be made publicly available online . 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 . To collect the content of a box the agent must first 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 . At this point the content of the box becomes accessible and can be picked up by the agent . Most boxes contain keys that , if made accessible , can be used to open other boxes . One of the boxes contains a gem , represented by a single white pixel . The agent ’ s goal is to acquire the gem by unlocking the box that contains it and picking it up by walking over it . Keys that an agent has in its possession are depicted in the input observation as a pixel in the top-left corner . In each level there is a unique sequence of boxes that need to be opened in order to reach the gem . Opening one wrong box ( a distractor box ) consumes the held key and leads to a dead-end , where the gem can not be reached and the level becomes unsolvable . There are three user-controlled parameters that contribute to the difficulty 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 . In general , the task is computationally difficult for a few reasons . 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 . 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 .

Summary: The authors present a deep reinforcement learning approach that uses a “self-attention”/“transformer”-style model to incorporate a strong relational inductive bias. Experiments are performed on a synthetic “BoxWorld” environment, which is specifically designed (in a compelling way) to emphasize the need for relational reasoning. The experiments on the BoxWorld environment clearly demonstrate the improvement gained by incorporating a relational inductive bias, including compelling results on generalization. Further experimental results are provided on the StarCraft minigames domain. While the results on StarCraft are more equivocal regarding the importance of the relational module—the authors do set a new state of the art and the results are suggestive of the potential utility of relational inductive biases in more general RL settings.