Grounding Language to Entities for Generalization in Reinforcement Learning

1 INTRODUCTION . Interactive game environments are useful for developing agents that learn grounded representations of language for autonomous decision making ( Golland et al. , 2010 ; Andreas & Klein , 2015 ; Bahdanau et al. , 2018 ) . The key objective in these learning setups is for the agent to utilize feedback from the environment to acquire linguistic representations ( e.g . word vectors ) that are optimized for the task . Figure 1 provides an example of such a setting , where the meaning of the word fleeing in the context is to “ move away ” , which is captured by the movements of that particular entity ( wizard ) . Learning a useful grounding of concepts can also help agents navigate new environments with previously unseen entities or dynamics . Recent research has explored this approach by grounding language descriptions to the transition and reward dynamics of an environment ( Narasimhan et al. , 2018 ; Zhong et al. , 2020 ) . While these methods demonstrate successful transfer to new settings , they require manual specification of some minimal grounding before the agent can learn ( e.g . a ground-truth mapping between individual entities and their textual symbols ) . In this paper , we propose a model to learn an effective grounding for entities and dynamics without requiring any prior mapping between text and state observations , using only scalar reward signals from the environment . To achieve this , there are two key inferences for an agent to make — ( 1 ) figure out which facts refer to which entities , and ( 2 ) understand what the facts mean to guide its decision making . To this end , we develop a new model called EMMA ( Entity Mapper with Multimodal Attention ) , which simultaneously learns to select relevant sentences in the manual for each entity in the game as well as incorporate the corresponding text description into its control policy . This is done using a multi-modal attention mechanism which uses entity representations as queries to attend to specific tokens in the manual text . EMMA then generates a text-conditioned representation which is processed further by a deep neural network to generate a policy . We train the entire model in a multi-task fashion using reinforcement learning to maximize task returns . 1Code and data are available at https : //www.dropbox.com/s/fnprjrfekbnxxru/code_data.zip ? raw=1 . To empirically validate our approach , we develop a new multi-task framework containing 1320 games with varying dynamics , where the agent is provided with a text manual in English for each individual game . The manuals contain descriptions of the entities and world dynamics obtained through crowdsourced human writers . The games are designed such that two environments may be identical except for the reward function and terminal states . This approach makes it imperative for the agent to extract the correct information from the text in order to succeed on each game . Our experiments demonstrate EMMA is able to outperform three types of baselines ( languageagnostic , attention-ablated , and Bayesian attention ) with a win rate almost 40 % higher on training tasks . More importantly , the learned grounding helps our agent generalize well to previously unseen games without any further training ( i.e . a zero-shot test ) , achieving up to a 79 % win rate . Our model is also robust to noise and linguistic variation in the manuals . For instance , when provided an additional distractor description , EMMA still achieves a win-rate of 75 % on unseen games . 2 RELATED WORK . Grounding for instruction following Grounding natural language to policies has been explored in the context of instruction following in tasks like navigation ( Chen & Mooney , 2011 ; Hermann et al. , 2017 ; Fried et al. , 2018 ; Wang et al. , 2019 ; Daniele et al. , 2017 ; Misra et al. , 2017 ; Janner et al. , 2018 ) , games ( Golland et al. , 2010 ; Reckman et al. , 2010 ; Andreas & Klein , 2015 ; Bahdanau et al. , 2018 ; Küttler et al. , 2020 ) or robotic control ( Walter et al. , 2013 ; Hemachandra et al. , 2014 ; Blukis et al. , 2019 ) ( see Luketina et al . ( 2019 ) and Tellex et al . ( 2020 ) for more detailed surveys ) . Recent work has explored several methods for enabling generalization in instruction following , including environmental variations ( Hill et al. , 2020a ) , memory structures ( Hill et al. , 2020c ) and pre-trained language models ( Hill et al. , 2020b ) . In a slightly different setting , Co-Reyes et al . ( 2018 ) use incremental guidance , where the text input is provided online , conditioned on the agent ’ s progress in the environment . Andreas et al . ( 2017 ) developed an agent that can use sub-goal specifications to deal with sparse rewards . Oh et al . ( 2017 ) use sub-task instructions and hierarchical reinforcement learning to complete tasks with long action sequences . In all these works , the text conveys the goal to the agent ( e.g . ‘ move forward five steps ’ ) , thereby encouraging a direct connection between the instruction and the control policy . This tight coupling means that any grounding learned by the agent is likely to be tailored to the types of tasks seen in training , making generalization to a new distribution of dynamics or tasks challenging . In extreme cases , the agent may even function without acquiring an appropriate grounding between language and observations ( Hu et al. , 2019 ) . In our setup , we assume that the text only provides high-level guidance without directly describing the correct actions for every game state . Language grounding by reading manuals A different line of work has explored the use of language as an auxiliary source of knowledge through text manuals . These manuals provide useful descriptions of the entities in the world and their dynamics ( e.g . how they move or interact with other entities ) that are optional for the agent to make use of and do not directly reveal the actions it has to take . Branavan et al . ( 2012 ) developed an agent to play the game of Civilization more effectively by reading the game manual . They make use of dependency parses and predicate labeling to construct feature-based representations of the text , which are then used to construct the action-value function used by the agent . Our method does not require such feature construction . Narasimhan et al . ( 2018 ) and Zhong et al . ( 2020 ) used text descriptions of game dynamics to learn policies that generalize to new environments , without requiring feature engineering . However , these works assume some form of initial grounding provided to the agent ( e.g . a mapping between object IDs and their descriptions , or the use of entity names in text as state observations ) . In contrast , our model figures out even this fundamental mapping between entity IDs in observation space and their symbols in text entirely through interaction with the environment . 3 FRAMEWORK . Our objective is to demonstrate grounding of environment dynamics and entities in a multi-task setup in order to drive generalization to unseen environments . Here and throughout , we refer to an entity as an object represented as a symbol in the observation that the agent may interact with . Environment We model decision making in each environment as a Partially-Observable Markov Decision Process ( POMDP ) with the 8-tuple ( S , A , O , P , R , E , Z , M ) . S and O are the set of all states and observations respectively where each o ∈ O contains entities from the set of entities E. At each step t , the agent takes some action at ∈ A. P ( st+1|st , at ) is the transition distribution over all possible next states st+1 conditioned on the current state st and action at . R ( st , at , st+1 ) is a function that provides the agent with a reward rt ∈ R for action at and transition from st to st+1 . Z is a set of text descriptions , with each z ∈ Z providing information about an entity e ∈ E. M is the map ze 7→ e which identifies the entity that each description describes . M , P , and R are not available to the agent . Note that there might not be a one-to-one mapping between Z and entities in the current state observation . Reinforcement Learning ( RL ) The objective of the agent is to find a policy π : O → A to maximize its cumulative reward in an episode . If π is parameterized by θ , standard deep RL approaches optimize θ to maximize the expected reward of following πθ . In our setup , we want the agent to learn a policy πθ ( a|o , Z ) that conditions its behavior on the provided text . However , in contrast to previous work ( Narasimhan et al. , 2018 ; Zhong et al. , 2020 ) , M is not available to our agent and must be learned through interaction . 4 MODEL . To learn a latent mapping between text symbols and entities , we develop a new model , EMMA ( Entity Mapper with Multi-modal Attention ) , which employs a soft-attention mechanism over the text descriptions . At a high level , for each entity description , EMMA first generates key and value vectors from their respective token embeddings obtained using a pretrained language model . Each entity attends to the descriptors via a symbol embedding that acts as the attention query . Then , instead of representing each entity with its embedding , we use the resulting attention-scaled values as a proxy for the entity . This approach helps our model learn a control policy that focuses on entity roles ( e.g . enemy , goal ) while using the entities ’ identity ( e.g . queen , mage ) to selectively read the text . We describe each component of EMMA below and in Figure 2 . Text encoder Our input consists of a h×w grid observation o ∈ O with a set of entity descriptions Z . We encode each description z ∈ Z using a BERT-base model whose parameters are fixed throughout training ( Devlin et al. , 2018 ; Wolf et al. , 2019 ) . For a description z , let t1 , ... , tn be its token embeddings generated by our encoder . We obtain key and value vectors kz , vz : kz = n∑ i=1 αiWkti + bk , α = softmax ( ( uk · tj ) nj=1 ) ( 1 ) vz = n∑ i=1 βiWvti + bv , β = softmax ( ( uv · tj ) nj=1 ) ( 2 ) The key and value vectors are simply linear combinations of Wkti + bk and Wvti + bv with weights α , β respectively , whereWk , Wv are matrices which transform each token to d dimensions and bk , bv are biases . The weights α , β are obtained by taking the softmax over the dot products ( uk · tj ) nj=1 and ( uv · tj ) nj=1 respectively . These weights imbue our model with the ability to focus on relevant tokens . All of Wk , bk , uk , Wv , bv , uv are learned parameters . Entity representation generator To get a representation for each non-avatar ( non-player ) entity e , we embed its symbol into a query vector qe of dimension d to attend to the descriptions z ∈ Z with their respective key and value vectors kz , vz . We use scaled dot-product attention ( Vaswani et al. , 2017 ) and denote the resulting representation for the entity e as xe : xe = m∑ i=1 γivzi γ = softmax ( ( qe · kzj√ d ) m j=1 ) ( 3 ) where m = |Z| is the number of descriptions in the manual . This mechanism allows EMMA to accomplish two forms of language grounding : the key and query select relevant descriptions for each object by matching entities to names ( e.g . mage ) , and the value extracts information relevant to the entities ’ behaviors in the world ( e.g . enemy , chasing ) . For each entity e in the observation , we place its representation xe into a tensor X ∈ Rh×w×d at the same coordinates as the entity position in the observation o to maintain full spatial information . The representation for the avatar ( player entity ) is simply a learned embedding of dimension d. Action Module To provide temporal information , we concatenate the outputs of the representation generator from the three most recent states to obtain a tensor X ′ ∈ Rh×w×3d . To get a distribution over the next actions π ( a|o , Z ) , we run a 2D convolution on X ′ over the h , w dimensions . The flattened feature maps are passed through a fully-connected feed-forward network terminating in a softmax over the possible actions . y = Flatten ( Conv2D ( X ′ ) ) π ( a|o , Z ) = softmax ( FFN ( y ) ) ( 4 ) In contrast to previous approaches that use global observation features to read the manual ( Zhong et al. , 2020 ) , we build a text-conditioned representation for each entity ( xe ) . One advantage is that xe can directly replace the entity embeddings typically used to embed the state observation in most models . This means EMMA ’ s action module can easily be swapped with existing models , such as txt2pi ( Zhong et al. , 2020 ) while still being completely end-to-end differentiable . Further details about EMMA and its design can be found in Appendix D .

Summary: The paper considers the task of training an agent to act following a manual expressed in natural language. The manual describes the roles and the behaviors of the entities in the environment. Each entity can be the goal, the message or the enemy and it can also be either fleeing, chasing or not moving. The agent has to bring the message to the goal while avoiding the enemy. A model called EMMA is proposed to change entity representations based on the manual, thereby making the agent aware of the entities’ roles. It is shown that EMMA is more effective than simple baselines.