Learning Knowledge Graph-based World Models of Textual Environments
Year: 2022

Abstract:
World models improve a learning agent’s ability to efficiently operate in interactive and situated environments. This work focuses on the task of building world models of text-based game environments. Text-based games, or interactive narratives, are reinforcement learning environments in which agents perceive and interact with the world using textual natural language. These environments contain long, multi-step puzzles or quests woven through a world that is filled with hundreds of characters, locations, and objects. Our world model learns to simultaneously: (1) predict changes in the world caused by an agent’s actions when representing the world as a knowledge graph; and (2) generate the set of contextually relevant natural language actions required to operate in the world. We frame this task as a Set of Sequences generation problem by exploiting the inherent structure of knowledge graphs and actions and introduce both a transformer-based multi-task architecture and a loss function to train it. A zero-shot ablation study on neverbefore-seen textual worlds shows that our methodology significantly outperforms existing textual world modeling techniques as well as the importance of each of our contributions.

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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.

heading: 3 Background
text: 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

heading: 4 The Worldformer
text: 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.

heading: 4.1 Knowledge Graph Difference Generation
text: 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 cannot 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.

heading: 4.2 Multi-task Architecture
text: The Worldformer is a multi-task world model that simultaneously learns to perform both knowledge graph and valid action generation. It is built on the hypothesis that each of these tasks contains information crucial to the other—the valid actions that can be executed at any timestep are entirely dependent on the current state and vice versa the state knowledge graph updates on the basis of the previously executed action.
Figure 3 describes the architecture of the Worldformer. The inputs to the architecture are textual observations, valid actions, and graph at time step t : Ot, Vt, Gt ∈ St. Ot and Vt are encoded through a bidirectional text encoder into Ot. In our work, we used an architecture similar to BERT [10] with the original pre-trained weights that are then fine-tuned using a masked language model (MLM) loss on observations taken from the training data. Ot is the output of the final hidden layer. The
graph encoder receives Gt and encodes it into Gt. It is also similar to BERT, but is pre-trained on knowledge graphs found in the training data using a MLM loss with a phrase-level masking scheme where whole components of a 〈s, r, o〉 graph triple (individual underlined portions in Figure 3) are masked at once. Again, Gt is the output of the final hidden layer.
Ot and Gt are passed into a representation aggregator which then sends the combined encoded state representation St to one of two autoregressive decoders that have the same general internal architecture as GPT-2 [27]. During training, the first decoder is conditioned on St directly and Ot through cross-attention and takes in the valid actions of the next state Vt+1 as input, learning to predict the same input sequence shifted to the right as sequence-to-sequence models do. Similarly, the second decoder is conditioned on St directly and Gt through cross-attention and takes in the knowledge graph of the next state Gt+1 as input.

heading: 4.3 Set of Sequences Generation and Training
text: We observe that both the knowledge graph difference Gt+1 − Gt and the valid actions Vt+1 are both Sets of Sequences where the ordering of the sequence of tokens within an action or a graph triple matters but the ordering of all the actions and triples does not. Standard autoregressive decoding used in sequence-to-sequence (Seq2Seq) models [32] does not account for such permutation invariance. We frame the graph and action prediction tasks as a generation of a Set of Sequences (SOS) problem—expanding on the simple set prediction problem definition proposed by works such as Deep Sets [40] to account for the specific structure of Sets of Sequences. This problem structure is used to then formulate a training methodology that lets autoregressive decoders better account for the SOS structure.
For both of the decoders in Figure 3, we are given a target sequence Y = {y1, ..., yM} and some input context via the encoders X . Standard autoregressive techniques factor the distribution over the target sequence into a chain of conditional probabilities with a causal left to right structure.
P (Y |X; θ) = M+1∏ i=1 p(yi|y0:i−1, X; θ) (1)
Where θ represents the overall network parameters. This can then be used to formulate a maximum likelihood training loss with cross-entropy at every step.
Lseq = logP (Y |X; θ) = M+1∑ i=1 logp(yi|y0:i−1, X; θ) (2)
In our setting, we can group elements in Y into its set of sequences form
Ysos = {y′1...y′M ′}, y′i ∈ Vt+1 or y′i ∈ Gt+1 −Gt,M ′ ≤M where y′i = {yk...yk+l}, ∑ j len(y′j) =M (3)
Via the decoders, we seek to learn a transformation from St ∈ IRd (the input d-dimensional state representation vector) and Ysos ∈ Y (decoder inputs in the space of all possible decoder inputs Y) that map to the permutation invariant target set of sequences Ysos. This function can then be defined as f : IRd ∪ 2Y → 2Y as the permutation invariance of part of the domain and range of this function makes it the power set of Y . Combining this definition of permutation invariant functions with Eq. 2, 3, we can factorize the distribution over the output Set of Sequences as the following chain of probabilities:
P (Ysos|X; θ) = M+1∏ i=1 p(y′i|X; θ) (4a)
p(y′i|X; θ) = l+n∏ k=l p(yk|yl:k−1, X; θ) (4b)
where l = ∑ j<i len(y′j), n = len(y ′ i)
With the key intuition here being that Eq. 4a factorizes the distribution such that each element of Ysos is independent of other elements in the set, but tokens of an element y′i in the set are conditioned on preceding tokens within the element (Eq. 4b).
This in turn gives us a maximum likelihood Set of Sequences loss that can be used to train a model to output a Set of Sequences.
Lsos = logP (Ysos|X; θ) = M+1∑ i=1 logp(y′i|X; θ)
= M+1∑ i=1 l+n∑ k=l logp(yk|yl:k−1, X; θ) (5a)
where l = ∑ j<i len(y′j), n = len(y ′ i) (5b)
In our formulation, we have observation sequences at timestep t : Ot, Vt encoded into Ot, graph Gt encoded into Gt, and all of them combined into St, with the output Sets of Sequences at timestep t+ 1 being the graph difference Gt+1 −Gt and valid actions Vt+1. Across the two decoders, this gives us a combined loss:
Lworld = logP (Gt+1 −Gt|St,Gt; θ) + logP (Vt+1|St,Ot; θ) (6)
This loss is used to multi-task train the Worldformer simultaneously across the two tasks.

heading: 5 Evaluation
text: We evaluate the Worldformer by comparing it on both of the tasks across 9 never-before-seen testing games against strong baselines. We further present ablation studies in each task to determine the relative importance of each of the techniques presented in the previous section.
Metrics. Across both the tasks, we use the same metrics as defined in JerichoWorld [4]. For knowledge graph generation, we report two types of metrics (Exact Match or EM and F1) operating on two different levels—at a graph tuple level and another at a token level. The graph level metrics are based on matching the set of 〈subject, relation, object〉 triples within the graph, all three tokens in a particular triple must match a triple within the ground truth graph to count as a true positive. The
token level metrics operate on measuring unigram overlap in the graphs, any relations or entities in the predicted tokens that match the ground truth count towards a true positive.
For valid action generation, we adapt the graph level Exact Match (EM) and F1 metrics as described in the previous task to actions. In other words, positive EM or F1 happens only when all tokens in a predicted valid action match one in the gold standard set. In all cases, EM checks for accuracy or direct overlap between the predictions and ground truth, while F1 is a harmonic mean of predicted precision and recall.

heading: 5.1 Knowledge Graph Generation
text: We compare the Worldformer to 4 baselines taken from contemporary knowledge graph-based world modeling approaches in text games. All sequence models use a fixed graph vocabulary of size 7002 that contains all unique relations and entities at train and test times. Additional details and hyperparameters for the models are found in Appendix A.2.
Rules. Following Ammanabrolu and Hausknecht [2], we extract graph information from the observation using information extraction tools such as OpenIE [6] in addition to some hand-authored rules to account for the irregularities of text games.
Question-Answering. (QA) This baseline comes from the Q*BERT agent described in [5]. It is trained on both the SQuAD 2.0 [28] the Jericho-QA text game question answering dataset [5] on the same set of training games as found in Worldformer. It uses the ALBERT [18] variant of the BERT [10] natural language transformer to answer questions and populate the knowledge graph via a few hand-authored rules from the answers. Examples of questions asked include: “What is my current location?”, “What objects are around me?”.
Seq2Seq. This single-task model is provided as a baseline by JerichoWorld and performs sequence learning by encoding the observation and graph with a single bidirectional BERT-based encoder
and using an autoregressive GPT-2-based decoder to decode the next graph. It is trained using the standard Seq2Seq cross-entropy loss (Eq. 2).
GATA-World. We adapt the Graph-Aided Transformer Agent [1] to our task. It consists of the same encoder structure as the Worldformer but contains one decoder that performs single-task Seq2Seq learning to decode both the set of tuples that must be added as well as deleted from the graph in the form of: 〈add, node1, node2, relation〉 or 〈del, node1, node2, relation〉. This is equivalent to predicting (Gt+1 −Gt) ∪ (Gt −Gt+1). It is trained with the Seq2Seq cross-entropy loss (Eq. 2). Table 1 describes the results in this task over all the games. We see that on the graph level metrics, the Worldformer performs significantly better than all other other baselines. On the token level metrics, the Worldformer and QA method are comparable—the difference between these two methods are statistically non-significant (p = 0.18) with each other but both significantly (p < 0.05) higher than all others. The QA method, and other extractive methods, highlight portions of the input observation that form the graph and are particularly well suited for the token level metrics. The JerichoWorld developers note that these approaches are prone to over-extraction, i.e. extracting more text than is strictly relevant from the input observation aiding token level overlap but resulting in a sharp drop in terms of the graph level metrics [4]. Additional failure modes of such extraction based approaches occur when the text descriptions are incomplete or hidden—e.g. the contents of a chest are revealed through the textual observation only when it is opened by a player. The Worldformer is able to make a informed guess as to the contents of the chest due to its training, providing a form of look ahead that the Rules and QA systems cannot. Ablation Graph Token
Graph Multi SOS EM F1 EM F1Diff Task Loss 14.29 15.54 18.80 19.96
X X 29.29 31.41 39.99 41.02 X X 32.60 34.65 42.74 44.35 X X 35.94 36.17 48.82 50.18 X X X 39.15 41.06 51.32 54.45
Table 2: Worldformer ablations to test the impact of its three main components for KG prediction. All results are size weighted averages over all test games over three random seeds, with standard deviations not exceeding ±3.2 in any category.
Table 3 present the results of an ablation study testing the relative importance of the three main components of the Worldformer: graph difference prediction, multi-task training, and the SOS loss. We note that a model without any of these components is equivalent to the Seq2Seq approach described previously. We see significant drops in performance, particularly on the graph level metrics, when any single one of these components are removed. This indicates that all
three components are necessary for the Worldformer to achieve state-of-the-art performance.
In particular, we note that the largest performance drop was when Worldformer did not use the graph difference simplification. In this case, the KG prediction task is simplified to predicting only; the length of the set of sequences Gt+1 − Gt is much smaller than Gt+1. There are on average 3.42 triples or 10.42 tokens per state across the JerichoWorld test dataset for Gt+1 −Gt but a mean of 8.71 triples or 26.13 tokens per state for Gt+1. This also explains the increased performance of the GATA-W over the baseline Seq2Seq agent—this agent only needs to predict on average 5.04 rules or 20.16 tokens across the testing games. Predicting a smaller number of triples and tokens per state makes the problem relatively more tractable for world modeling agents.

heading: 5.2 Valid Action Generation
text: Similarly to the other task, we compare the Worldformer to an existing baseline for valid action prediction. All models use a fixed vocabulary of size 11,056 at train and test times.
Seq2Seq. This single-task model is provided as a baseline by JerichoWorld and is identical to the Seq2Seq model described in the previous task but is single-task trained to predict valid actions.
CALM. A complementary dataset of observation-action pairs created by humans on the ClubFloyd online Interactive Narrative forum2 appears in both Ammanabrolu and Hausknecht [2] and Yao et al. [39] with the latter using it to tune a GPT-2 model for valid action prediction using a GPT-2 based Seq2Seq valid action model dubbed CALM.3 This model takes in Ot, A,Ot+1 and targets Vt+1.
In Table 1, we see that the Worldformer significantly outperforms the Seq2Seq baseline on all the games and CALM overall. Each valid action in a text game requires at most 5 tokens. This combined with an average of 10.30 valid actions per test state means that for every state we would need to generate about 52 tokens. Yet further, the vocabulary size for actions is 11, 056, larger than the graph
2http://www.allthingsjacq.com/interactive_fiction.html 3https://github.com/princeton-nlp/calm-textgame
vocabulary of 7, 002. This increase in task difficulty explains the relative decrease in the magnitude of performance metrics between KG and valid action prediction tasks. Both the Seq2Seq model and CALM—which is trained on a different dataset—are comparable on F1 scores but Seq2Seq is better overall for exact matches. CALM also has relatively higher variance in performance across the test games than the other two methods—e.g. on some games such as zork1 and detective it outperforms the Seq2Seq and is not too far off the Worldformer especially in terms of F1 score. This would appear to indicate that the Club Floyd dataset of text game transcripts that CALM was trained on is better suited for transfer to certain games than others due to differences in training set genre similarities.
Table 3 presents an ablation study that tests the two main components of the Worldformer for this task: multi-task learning, and SOS loss. As with the KG prediction task, we observe significant drops in performance when either of these components are taken away—suggesting that they are relatively critical components. The JerichoWorld developers note that there is a correlation between performance of the baseline Seq2Seq model to the average number of valid actions for the testing game (see Appendix Table 4). They attribute this to label imbalance in the dataset, stating that the model likely learns a common set of actions found across all games such as navigation actions before learning more fine-grained actions. E.g. ztuu, deephome, and balances have a high number of gold standard average valid actions while pentari, ludicorp, detective, and temple which have a low number of average valid actions. While the latter
set of games have generally higher performance on both the Seq2Seq and Worldformer models, the gap is significantly less pronounced with the Worldformer. We hypothesize that this is due to the multi-task training of the Worldformer—encoder representations now contain enough information regarding the next knowledge graph to alleviates the label imbalance of the actions and enable prediction of more fine-grained actions.

heading: 6 Conclusions
text: We presented the Worldformer, a state-of-the-art world model for text games that: maps worlds by predicting the difference in knowledge graphs between subsequent states, multi-task learns to map a world and act in it simultaneously, and frames all of these tasks as a Set of Sequences generation problem. Its state-of-the-art performance and an ablation study have three potential implications: (1) the simplification of the knowledge representation problem into that of predicting knowledge graph differences between subsequent states is a critical step in making the problem more tractable; (2) performance improvements due to multi-task training imply that acting in and mapping these worlds is inherent a highly correlated problem and benefits from being solved jointly; and (3) the performance boosts due to the SOS loss suggest that accounting for this property of graphs and actions enables more effective training than if we were to treat them as simple sequences.

heading: 7 Broader Impacts
text: We view text-games as an platform on which to teach agents how to communicate effectively using natural language, to plan via sequential decision making in situations that may not be anticipated. We seek to enable agents to more efficiently model such worlds, helping them produce more contextually relevant language in such situation. As stated in JerichoWorld—and further verified by us—data is collected from games containing situations of non-normative language usage—describing situations that fictional characters may engage in that are potentially inappropriate, and on occasion impossible, for the real world such as running a troll through with a sword. Instances of such scenarios are mitigated by careful curation of the games that the data is collected by both the authors of Jericho [13] and JerichoWorld [4]. This is based on manual vetting and (existing) crowd-sourced reviews on the popular interactive narrative forum IFDB (https://ifdb.org/).
Broadly speaking, the most relevant downstream task for this work is model-based reinforcement learning. It is applicable to many sequential tasks, some of which cannot be anticipated. World modeling for text environments is more suited for domains in which change in the world is affected via language, which mitigates physical risks—downstream lines of work are not directly relevant to robotics—but not cognitive and emotional risks, as any system capable of generating natural language is capable of accidental or intentional non-normative and biased language use [23, 31].

Review summaries:
This paper presents a method for modelling state (via graph prediction) and action distributions in text adventure games. To do the former, this work proposes predicting incremental graph updates instead of entire graphs. For both, the authors propose predicting sets of sequences as opposed to sequence generation and show that this improves the accuracy of predicted action sets and environment state.

The paper utilizes world models for text-based reinforcement learning agents to be aware of the change in the world based on one's own actions and what actions could be performed based on the current state of the world using Worldformer, a multi-task transformer for jointly learning to generate the knowledge graph differences based on the action taken and a set of valid action at the current state. Since the valid actions generated at the current step affect the action taken to change the world, these tasks are learned using the same loss function.  **Challenges**: Partial observability, combinatorial action space

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.