Learning Knowledge Graph-based World Models of Textual Environments

1 Introduction . World models , often in the form of probabilistic generative models , are used in conjunction with model-based reinforcement learning to improve a learning agent ’ s ability to operate in various environments [ 33 , 7 ] . They are inspired by human cognitive processes [ 15 ] , with a key hypothesis being that the ability to predict how the world will change in response to one ’ s actions will help you better plan what actions to take [ 12 ] . Evidence towards this hypothesis comes in the form of studies showing that simulating trajectories using internal learned models of the world improves sample efficiency in learning to operate in an environment [ 12 , 30 ] . Text-based games , in which players perceive and interact with the world entirely through textual natural language , are environments that provide new challenges for world model approaches . Textbased games are structured as long puzzles or quests that can only be solved by navigating and interacting with potentially hundreds of locations , characters , and objects . The puzzle-like structures to the games are exacerbated by two factors . First , these environments are partially observable , i.e . observations of the world are incomplete . Second , the agent faces a combinatorially-sized action space of the order of O ( 1014 ) possible actions at every step . For these reasons model-free reinforcement learning in text-based games is extremely data inefficient . Prior work on text-based game playing agents repeatedly demonstrated that providing agents with structured memory in the form of knowledge graphs ( sets of 〈s , r , o〉 tuples such that s is a subject , r is a relation , and o is an object ) is critical in enabling them to operate in and model these worlds [ 3 , 21 , 1 , 2 ] —aiding in both the challenges mentioned . These works all rely on extracting information 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . about one ’ s surroundings while navigating novel environments , either through rules [ 3 , 2 ] , questionanswering [ 5 ] , or transformer-based extraction [ 1 , 21 ] . This lifted representation helps agents remember aspects of the world that become unobservable as the agent navigates the environment . However , we hypothesize that agents that rely on lifted representations of the world will benefit from more than just memorization but the ability to predict how the graph state representation will change . For example , by inferring that a locked chest is likely to contain treasure before it is actually revealed provides an agent with a form of look-ahead that will potentially enable it to bias its actions towards opening such a chest . We introduce an approach to this knowledge representation problem that effectively simplifies it by exploiting the inherent structure of knowledge graphs—framing it to be the task of inferring the difference in knowledge graphs between subsequent states given an action . A consequence of the combinatorially-sized action space in text-based games is that that the set of contextually relevant actions—i.e . those that are most likely to affect change in the environment—are overwhelmed by the irrelevant actions . For example , it is not illegal to try to climb a tree when there are no trees present , and the game engine will just respond with feedback that nothing happens . An aspect of world modeling that has not been considered for other games is inferring which actions are valid in a particular context . We hypothesize that both the challenges mentioned are closely linked and present world models that multi-task learn to tackle both simultaneously—i.e . to answer the questions of “ What actions can I perform ? ” and “ How will the world change if I perform a particular action ? ” . Our work has four core contributions . ( 1 ) We first show how changes in the world can be represented in the form of differences between subsequent knowledge graph state representations . ( 2 ) We present the Worldformer , a novel multi-task transformer based architecture that learns to simultaneously generate both the set of graph differences and the set of contextually relevant actions . ( 3 ) We introduce a loss function and a training methodology that enable more effective training of the Worldformer by exploiting the fact that knowledge graphs and natural language valid actions can be represented permutation invariant Sets of Sequences—wherein the ordering of tokens within an item in the set matters but the set itself lacks ordering . ( 4 ) A zero-shot ablation study conducted on diverse set of never-before-seen text games shows the significance of each of the prior three contributions in outperforming strong existing baselines . 2 Related Work . We will focus on three main areas of related work : world modeling and model-based reinforcement learning , and world modeling and knowledge graphs in text games , and general ( knowledge ) graph construction techniques . World modeling via model-based reinforcement learning often serves to learn transition models of an environment to allow for simulation without actually interacting with the environment [ 7 ] . Ha and Schmidhuber [ 12 ] use Variational Autoencoders ( VAEs ) combined with recurrent neural networks to learn compressed state representations over time of visual reinforcement learning environments [ 8 ] . This model is then used to simulate an environment and learn a control policy in it . Other contemporary works attempt to also learn dynamics models using raw pixels in the context of games such as Atari [ 26 , 16 ] , and Super Mario Bros. [ 11 ] as well as 3D simulations [ 16 ] and robotics [ 37 , 35 ] . We note that in all of these works , in addition to the state space being raw pixels—the action space is fixed and orders of magnitude smaller than in text games . In textual environments , the traditional state representations of choice have been raw text encodings via recurrent neural networks [ 24 , 14 , 13 ] but have since shifted towards transformer [ 34 ] and knowledge graph-based representations [ 3 , 1 ] . Knowledge graphs have been shown to be aid in the challenges of : ( 1 ) knowledge representation [ 3 , 1 ] , enabling neuro-symbolic reasoning approaches over graph-based state representations [ 29 ] ; ( 2 ) combinatorial state-action spaces [ 2 , 1 ] ; and ( 3 ) incorporating external knowledge sources for commonsense reasoning [ 3 , 21 , 22 , 9 ] . Two of these works are perhaps closest in spirit to ours . Yao et al . [ 39 ] train a GPT-2 model [ 27 ] to decode valid actions based on human text game transcripts found online , showing that improved valid action generation ability results in better control policies . Ammanabrolu et al . [ 5 ] frame knowledge graph construction in text games as a question-answering problem where agents ask questions to identify common objects in the world and their attributes , showing that improved knowledge graph quality results in better control policies . We note that of these works handles only one or another of the sub-tasks required for world modeling in these environments . A core aim in the field of graph representation learning is representing graphs as continuous vectors while maintaining their inherent structure [ 17 , 38 ] . Approaches to this task when applied to knowledge graphs , attempt to exploit the inherent structure of knowledge graphs to create more accurate continuous vector graph representations in the form of embeddings [ 36 ] . Building on these representation learning works is the task of automated knowledge base construction attempts to create links in knowledge graphs given text [ 25 ] . Li et al . [ 19 ] approach the graph generation problem as a sequential process , first learning a generative model of the graph and then iteratively adding nodes and edges to it using the learned model—this work does not consider cases where the graphs are conditioned on text , however . These works focus solely on graph construction and do not include the inherent interactive action-based components featured in world modeling and text games . 3 Background . Dataset . We use the JerichoWorld Dataset [ 4 ] .1 It contains 24,198 mappings between rich natural language observations and : ( 1 ) knowledge graphs in the form of a set of tuples 〈s , r , o〉 ( such that s is a subject , r is a relation , and o is an object ) that reflect the world state in the form of a map ; ( 2 ) a set of natural language actions that are guaranteed to cause a change in that particular world state . An example of the mapping between rich natural language observations and structured knowledge is illustrated in Figure 1 . The training data is collected across 27 text games in multiple genres and contains a further 7,836 heldout instances over 9 additional games in the test set . Each instance of the dataset takes the form of a tuple of the form 〈St , A , St+1 , R〉 where St and St+1 are two subsequent states with A being the action used to transition between states and R the observed reward . As mentioned earlier , each of the states in the tuple contain information regarding the observation Ot ∈ St , ground truth knowledge graph Gt ∈ St , and valid actions for that state Vt ∈ St . This data was collected by oracle agents , i.e . agents that can perfectly solve a game , exploring using a mix of an oracle and random policy to ensure high coverage of a game ’ s state space . A full sample is found in Appendix A.1 . Tasks . Given this dataset , we focus on two tasks within it as formally defined by JerichoWorld . As mentioned in Section 1 , a successful world model will be able to accomplish both of these tasks . 1 . Knowledge Graph Generation : this task involves predicting the graph at time step t+ 1 : Gt+1 ∈ St+1 given the textual observations , valid actions , and graph at time step t : Ot , Vt , Gt ∈ St , and action A for all samples in the dataset . 2 . Valid Action Generation : this task is formally defined as predicting the set of sequences of valid actions at time step t+ 1 : Vt+1 ∈ St+1 given the textual observations , valid actions , and graph at time step t : Ot , Vt , Gt ∈ St , and action A for all samples in the dataset . 1https : //github.com/JerichoWorld/JerichoWorld 4 The Worldformer . This section describes the core methodological contributions of our work in creating world models for text games . We first show how knowledge graph generation can be simplified to predicting the graph difference between agent steps . We then describe the Worldformer , a transformer-based architecture , and end-to-end training method—including an objective function—that treats both of the world modeling tasks as a Set of Sequences generation problem . 4.1 Knowledge Graph Difference Generation . Figure 2 describes the gist of our simplification of the knowledge graph generation problem . Recall that knowledge graphs are directed graphs that are stored the form of a set of tuples as 〈s , r , o〉 such that s is a subject , r is a relation , and o is an object . Let the knowledge graphs representing the world state at two subsequent steps be Gt and Gt+1 . At every step , tuples are either added or deleted from the graph Gt to update the belief state about the world and turn it into graph Gt+1 . Using this observation , we can simplify the knowledge graph generation problem . Instead of predicting Gt+1 given Gt and prior context , we can instead predict the differences between the two graphs . In Figure 2 , between steps t and t+1 , we see that Gt+1 −Gt is the set of tuples that are added to Gt and Gt−Gt+1 the set of tuples are are deleted from Gt . Together they make up the graph differences . Here , we make a second key observation that allows for yet further simplification of the problem . This observation is based on generally applicable properties of such worlds : ( 1 ) locations are fixed and unique , i.e . the positions of locations with respect to each other does not change ; ( 2 ) objects and characters can only be in one location at a time ; and ( 3 ) contradicting object attributes can be identified using a lexical dictionary such as WordNet [ 20 ] , e.g . an object can not be both open and closed at the same time . These properties let us uniquely identify the triples to be deleted from the graph Gt − Gt+1 given triples to be added to the graph Gt+1 − Gt . Additional implementation details can be found in Appendix A.2 . Taken together , the Knowledge Graph Generation task can be cast as follows : predict the nodes to be added to the graph Gt at time step t : Gt+1 − Gt ( a much smaller set than Gt+1 by itself ) to transform it into graph Gt+1 given the textual observations , valid actions , and graph at time step t : Ot , Vt , Gt ∈ St , and action A for all samples in the dataset .

Summary: In this paper, the authors propose a novel system that tries to learn world model --- represented as knowledge graphs --- for agents playing text-based adventure games.   To learn the world model, the authors design two tasks: - Knowledge graph prediction. More specifically, knowledge graph difference prediction. At a game step $t$, given the observation $O_t$, the set of valid actions $V_t$, the graph representation $G_t$, and an action $A \in V_t$, the model is required to predict the difference between $G_{t+1}$ and $G_t$, represented as a set of $<s, r, o>$ triplets. Using heuristics, the system deduces $G_{t} - G_{t+1}$ from $G_{t+1} - G_t$. Together, the set of graph difference (triplets) can be used to infer the full set of $G_{t+1}$, which represents the new game state $S_{t+1}$. -Valid action prediction. Similarly, at a game step $t$, given $O_t$, $V_t$, $G_t$ and $A \in V_t$, the model is required to predict the set of valid actions $V_{t+1}$ at next game step.   Since both of the above tasks are essentially set (of sequences) generation tasks, the authors propose a modification on the standard autoregressive NLG paradigm. Specifically, the decoder outputs a sequence of elements where each element is a multi-word chunk (either a  $<s, r, o>$ triplet or an action), the author propose to remove the dependency across elements so that when generating a word, the decoder is only conditioned on the previous tokens within the same element. From Figure 3, the authors show that such chunks are concatenated with special tokens (e.g., [ACT], [TRIPLE]).   Another contribution is that the authors propose a multi-task architecture to jointly learn the aforementioned two tasks. Because the authors use triplets (sequences of tokens) to represent graphs, both the text encoder and graph encoder can be implemented as transformer-based encoders. The authors use an aggregator to combine the encodings from both encoders, denoted as $S_t$, which is later used to perform cross-attention with each of the individual encoding as the input to their corresponding decoder.