Learning Knowledge Graph-based World Models of Textual Environments
Authors: Prithviraj Ammanabrolu
Mark Riedl
Venue: NeurIPS 2021 Poster (2021)
Keywords: world models
text games
knowledge graphs
natural language processing

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 never-before-seen textual worlds shows that our methodology significantly outperforms existing textual world modeling techniques as well as the importance of each of our contributions.

Learning Knowledge Graph-based World Models of Textual Environments
Prithviraj Ammanabrolu
Mark O. Riedl
School of Interactive Computing Georgia Institute of Technology raj. ammanabr weceebroouugeee ll@@gatech. edu
School of Interactive Computing Georgia Institute of Technology riedl@cc. gatech. edu
Abstract
World models improve 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, interactive narratives, are reinforcement learning environments 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 knowledge graph; and (2) generate the set of contextually relevant natural language actions required to operate in the world. We frame this task as Set of Sequences generation problem by exploiting the inherent structure of knowledge graphs and actions and introduce both a transformer -aassd multi-tase architecture and loss function to train it. A zero-shot ablation study never- before-seen textual worlds shows that our methodology significantly outperforms existing textual world modeling techniques as well as the importance of each of our contributions.
Introduction
World models, often the form probabilistic generative models, are used in conjunction with model-based reinforcement learning to improve learning agent's ability to operate in various environments 33 7 They are inspired by human cognitive processes [15] with key hypothesis being that the ability to predict how the world will change in response to one's '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, environments that provide based 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, ie. observations of the world are incomplete. Second, the agent faces a combinatorially action space of the order of O410'4) possible actions at every step. For these reasons model-free reinforcement learning in text-based games is extremely data inefficient.
challenges for world model approaches. Text-
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, tuples such that is subject, a relation, and is an object) is critical in enabling them to operate in and model these worlds [3] 21 aiding both the challenges mentioned. These works all rely on extracting information
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

North
Loceloone Surt OBLs
You
House
the Oktj Agributess
have
marth
go north
Vieet of
House
haa
P
Jisse
West of House
(span>
are standing open field west house, with a boarded front door. There is small mailbox here.
North House
abble
You are facing the north side of white house, There no door here. and all the windows are boarded up. To the puth winds through the You are carrying: a small leaflet Prev aets north
Leaflet
You are emppy-horded
Prev act: take leaflet
Valid actss go north, go south, go west, open matilbox
go nurth, go cast, gu west, drop
Figure Two subsequent states in Zorkl consisting of: textual observations, world knowledge graphs, valid actions, and actions taken.
about one's surroundings while navigating novel environments, either through rules 3 12 question- answering [5]. or transformer-basee extraction [1] [1]] This lifted representation helps agents remember aspects of the world that become unobservable 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 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 approach to this knowledge representation problem that effectively simplifies by exploiting the inherent structure knowledge graphs-framing it to be the task of inferring the difference in knowledge graphs between subsequent states given action.
consequence of the combinatorially sized action space 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, 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 inferring which actions are valid in particular context. We hypothesize that both the challenges mentioned are closely linked and present world models that multi-task learn to tackle both simultancously questions of "What actions can perform?" and "How will the world change if perform particular action?".
to answer the
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 loss function and 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 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 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 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 g This model is then used to simulate an environment and learn control policy in it. Other contemporary works attempt also learn dynamics models using raw pixels in the

---

context of games such as Atari 226] [6]. and Super Mario Bros. [ as well as 3D simulations [[6] and robotics [7 351. We note that in all of these works, in addition to the state space being raw pixels-the 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] [4
13) but have since shifted towards transformer (34) and
knowledge graph-based representations 3 Knowledge graphs have been shown be aid in the
challenges of: (1) knowledge representation B enabling neuro-symbolic reasoning approaches
over graph-bast based state representations (2 : (2) combinatorial state aattinccionn spaces
I: and (3)
22 9 Two of these works are perhaps closest in spirit to ours. Yao et [3] 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 problem where agents ask questions to identify
incorporating external knowledge sources for commonsense reasoning [3] 21
construction in text games as question-answeriin 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 vec- tors while maintaining their inherent structure [7
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 (3). 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.
38
3 Background
Dataset. We use the JerichoWorld Dataset 4 It contains 24, 19, mappings between rich natural language observations and: (1) knowledge graphs in the form of set of tuples (s, r, o) (such that is a subject, is a relation, and o is an object) that reflect the world state in the form of map: (2) set of natural language actions that are guaranteed cause change in that particular world state. An example of the mapping between rich natural language observations and structured knowledge illustrated in Figure I 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 tuple of the form (S,, A, R) where S, and S 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 O2 E Sr, ground truth knowledge graph G2 E St, and valid actions for that state V, E Sy. This data was collected by oracle agents, i.e. agents that can perfectly solve a game, exploring using mix of an oracle and random policy to ensure high coverage of game's 's state space. full sample is found in Appendix [A.]
Tasks. Given this dataset, we focus on two tasks within it as formally defined by BerichoWorlde As mentioned in Section successful world model will be able accomplish both of these tasks.
1. Knowledge Graph Generation: this task involves predicting the graph at time step / + 1: G++1 E Se+1 given the textual observations, valid actions, and graph at time step & O₁, Vi, G2 E Sy, 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 + 1: Vet1 E S given the textual observations, valid actions, and graph at time step t O,, Vz, G E St, and action A for all samples in the dataset.
3

---

You
Jindon
House
Vou
1Bosen
You
Kandon
hawa
1orhh
Vius
Houne
north
Ggg11
Figure 2: The transformation between subsequent world knowledge graphs G2 and G1+1 based on the states in Figure The green (Gr) outlined portions in the center are additions to G2 to get G++1 (i.e. G++1 G+) and the red (R) portions similarly represent deletions to G. (i.e. G. - G(+1).
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, transformer-basee architecture. and end-to-end training method saabuodnnncluiing 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 E 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 set of tuples as (s, r.,o, 0) such that is a subject, is a relation, and is object. Let the knowledge graphs representing the world state at two subsequent steps be G, and G.+1. At every step, tuples are either added or deleted from the graph G, to update the belief state about the world and turn it into graph G Using this observation, we can simplify the knowledge graph generation problem. Instead of predicting G++1 given G2 and prior context, can instead predict the differences between the two graphs.
In Figure 2 between steps t and ++ 1, we see that G G2 the set of tuples that are added to Ge and G2 G++1 the set of tuples are are deleted from G.. 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.c. the positions of locations with respect to each other does not change; (2) objects and characters can only be in one location at 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 G, G++1 given triples to be added to the graph G - Gt. Additional implementation details can be found in Appendix
A.22 A.2]
Taken together, the Knowledge Graph Generation task can be cast as follows: predict the nodes to be added to the graph G2 at time step & G++1 G. (a much smaller set than G by itself) to transform it into graph G++1 given the textual observations, valid actions, and graph at time step f O,, V1, G, E S, and action A for all samples in the dataset.
4.2 Multi-task Architecture
The Worldformer a multi-task world model that simultancously learns to perform both knowledge graph and valid action generation. It is built 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 the basis of the previously executed action.
Figure describes the architecture of the Worldformer. The inputs to the architecture are textual observations, valid actions, and graph at time step L: O,, Vr, G, E S.. O, and Vú, are encoded through a bidirectional text encoder into O, 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. O is the output of the final hidden layer. The

---

Utwp kaald (ACT) so matd
cross attention
Bidirectional
Text
Action Decoder
Encoder
O
[OBS] West of House You are standing
[ACT] an murth [ACT] ynum maillbas
open field west of white house...
Aggre- gallor
St
[ACT] drop lesflet [ACT] po west
you
House
Graph Ensoder
Gt
G
Graph Decoder
croes-atlention
[GRAPH]] you, in. West of House (TRIPLE) West of house, has, mailbox [TRIPL.
mailhox gpenable
Gr-1
[GRAPH]
you in North of House
Figure 3: The Worldformer architecture. The text encoder (B) and graph encoder (R) have similar architecture but different pre-training strategies. Both the decoders are not pre-trained and have identical architectures. Solid black lines indicate gradient flow.
graph encoder receives G; and encodes it into Gt. It is also similar BERT, but is pre-trained knowledge graphs found in the training data using MLM loss with phrase-level masking scheme where whole components of (s, o) graph triple (individual underlined portions in Figure 3 are masked at once. Again, Gt is the output of the final hidden layer.
O2 and G. passed into representation aggregator which then sends the combined encoded state representation S, to one of two autoregressive decoders that have the same general internal architecture as GPT-2 (7)). During training, the first decoder is conditioned through cross-attention and takes in the valid actions of the next state V++1 as input, learning to predict the same input sequence shifted to the right as -equencc-to-sequencc models do. Similarly, the second decoder is conditioned on S, directly and G. through cross-attention and takes in the knowledge graph of the next state G.41 input.
S, directly and O
4.3 Set of Sequences Generation and Training
We observe that both the knowledge graph difference G.+1 - G2 and the valid actions V4+1 are both Sets of Sequences where the ordering of the sequence of tokens within action graph triple matters but the ordering of all the actions and triples does not. Standard autoregressive decoding used in settuennco----euennee (Seq2Seq) models [32] does not account for such permutation invariance. We frame the graph and action prediction tasks as generation of a Set of Sequences (SOS) wobleem-eppanding
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 then formulate training methodology that lets autoregressive decoders better account for the SOS structure.
For both of the decoders in Figure 3 we are given target sequence Y = (V1, UM } and some input context via the encoders X. Standard autoregressive techniques factor the distribution over the target sequence into chain of conditional probabilities with causal left to right structure.
M+1
P(Y|X; II 1o X;0)
Where represents the overall network parameters. This can then be used to formulate a maximum likelihood training loss with cross-entropy at every step.
M+1
Lang = log
Σ
logo(s: (100011000-11,...
(2)
5

---

In our setting, we can group elements in Y into its set of sequences form
Yos f/ J E V4+1 or 1// E G4+1 - GG, M' - M
where v/. (ye-..9441), Σ len(u/s) M
Via the decoders, we seek to learn a transformation from S, IR (the input d-dimensiioal state representation vector) and sos E y (decoder inputs in the space of all possible decoder inputs y) that map to the permutation invariant target set of sequences Yoss. This function then be defined as f IR. U2 + 2 as the permutation invariance of part of the domain and range of this function makes the power set of y
Combining this definition of permutation invariant functions with Eq. we can factorize the distribution over the output Set of Sequences as the following chain of probabilities:
M+1
P(Y (x;0== = II
p(y/ x;
(4a)
1=1
I+n
p(/((:00 = II
Mu 1:
(4b)
km
where / len(y/)), n n= len(y()
With the key intuition here being that Eq. 4a factorizes the distribution such that each element of Yos independent other elements the set, but tokens of an element s/f in the set are conditioned on preceding tokens within the element (Eq. 4bp.
This in gives maximum likelihood Set of Sequences loss that can be used train model to output a Set of Sequences.
M+1
Leos logPlYssoll (X;0) =
Σ
1=1
M+1
Σ
logp(uk Ogp(uklurk-2) 11 X;;)
(5a)
i=1
k=
where 1 Σ
len(v/), -len(v()
(5b)
In our formulation, have observation sequences timestep t: Oz, Vi encoded into Ot, graph Ge encoded into Gr. and all of them combined into St, with the output Sets of Sequences at timestep t being the graph difference G++1 G2 and valid actions V Across the two decoders, this gives us a combined loss:
Lworld logP G.|Se, G): 1o O 0)
(6)
This loss is used to multi-task train the Worldformer simultaneously across the two tasks.
5
Evaluation
We evaluate the Worldformer by comparing it on both of the tasks across never-before-seen testing games against strong baselines. We further present ablation studies in each task to determine the relative importance each of the techniques presented the previous section.
Metrics. Across both the tasks, we use the same metrics as defined in berichoWorld 4 For
knowledge graph generation, we report two types of metrics (Exact Match or EM and F1) operating on two different levels at graph tuple level and another at token level. The graph level metrics are based on matching the set of (subject, relation, object) triples within the graph, all three tokens in particular triple must match triple within the ground truth graph to count as a true positive. The

---

Expt.
Met-
Game
zork1
lib.
det.
bal.
pent.
ztuu
ludi.
deep.
temp.
overall
1294
7836
Knowledge Graph Prediction
Gr.
EM
3.72
7.61
1.39
9.17
6.44
4,94
5.10
0.49
2.48
4.70
Rules
EM
6.08
10.33
7.51
32.53
16.48
14.40
14.47
3.34
7.42
13.08
Tok.
8.42
26.74
10.23
36.09
23.36
21.74
18.48
3.86
9.44
17.50
FM
24.56
29.14
34.45
41 22
28.96
22.17
.44
4.42
36.84
32.79
QBBERT
F1
24.88
31.46
36.23
41.85
30.12
26.26
46.74
4.66
39.86
35.48
QQuestion
Tok.
EM
43.93
49.78
60.28
85.8
65.02
49. 44
57.58
9.31
48.98
53.58
Answering)
F1
48.31
52.76
63.21
86.18
69.54
49.82
610.95
9.84
49.17
55.74
EM
12.44
18.42
26.80 8G
8.19
22.18
16.89
12.94
8.38
16.48
14.29
Gr.
Seq2Seq
F1
12.96
18.89
29:48
9.04
23.54
16.89
14.18
10.47
18.52
15.54
Tok.
F1
21.12
20.84
35.86
18. 18.86
27.72
17.87
15.42
13.25
24.34
19.96
Gr:
EM
22.30
24.72
21.72
23.68
22.81
27.00
24.55
23.76
24.52
24.06
25.34
22.14
26.54
27.63
27.00
24.55
23.76
25.19
GATA-W
Tok.
EM
33.09
33.88
25.64
34.64
37.71
35.88
35.94
32.48
40.899
35.31
39.59
38.88
43.99
37.10
Worldformer
Gr.
EM
21.62
34.39
41.05
50.41
30.00
41.56
40.10
41.87
42.43
39.15"
24.44
34.39
44.53
52.43
34,30
42.20
41.65
42.74
4..177
41.06"
EM
$4.39
51.29
50.04
51.32
Tok.
F1
48.12
41.98
59.13
62.22
49.00
52.24
$1.29
50.04
$4.96
$2.45
Valid Action Prediction
Seq2Seq
Act
EM
16.65
15.13
18.19
16.19
23.39
14.75
20.10
14.71
20.34
18.10
F1
17.85
16.88
21.12
18.23
25.87
15.13
20.86
14.86
22.14
19.44
CALM
Act
EM
18.67
11.18
17.37
10.04
13.77
11.29
13.49
10.31
13.13
13.79
F1
18.90
25.49
34.42
12.16
34.40
9.95
20.94
7.84
18.57
19.11
Worldformer
Act EM
23.08
22.35
20.97
29.08
27.05
20.71
21.36
24.04
22.80
23.22
31.32
23,66
22.27
25.54'
Table 1: Results across both the tasks for the models specified. Overall indicates a size weighted average. All experiments were performed over three random seeds, with standard deviations not exceeding =3.2 in any of the overall categories for KG prediction and 1.2 for valid action prediction. Bolded results indicate highest overall scores. Asterisk (*) indicates when the top result is significantly higher (p < 0.05 with an ANOVA test followed by a post-hoc pair-wise Tukey test) over all alternatives. indicates this result is not significantly higher than Worldformer.
token level metrics operate measuring unigram overlap in the graphs, any relations 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 actions. In other words, positive EM F1 happens only when all tokens in 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 a harmonic mean of predicted precision and recall.
5.1 Knowledge Graph Generation
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 train and test times. Additional details and hyperparameters for the models are found in Appendix A.2 A.2
Rules. Following Ammanabrolu and Hausknecht we extract graph information from the observa- tion using information extraction tools such as OpenlE 6 in addition to some hand-authored rules to account for the irregularities of text games.
(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. uses the ALBERT [8] variant of the BERT [10] natural language transformer answer questions and populate the knowledge graph via a few authored rules from the answers. Examples of questions asked include: "What is my current location?" ?" "What objects around me?".
Auestion-Answeriing
Seq2Seq. This single-task model is provided as baseline by Jericho World and performs sequence learning by encoding the observation and graph with single bidirectional BERT-based encoder
7

---

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
GoTA-World. We adapt the Graph-Aided Transformer Agent [ 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, nodel, node2, relation) or (del, nodel, node2, relation). This is equivalen to predicting (G++1 - G.) U (G, - G It is trained with the Seq2Seq cross-entropy loss (Eq. 2
Table I 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 Jericho World developers note that these approaches are prone 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 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
Table 3 present the results of an ablation study testing the relative importance of the three main components of the Worldformer: graph differ- ence prediction, multi-task training, and the SOS loss. We note that a model without any of these components is equivalent to the Seq2Seq ap- proach described previously. We see significant drops in performance, particularly on the graph level metrics, when any single one of these com- ponents are removed. This indicates that all three components are necessary for the Worldformer to achieve state-of-the-art performance.
Diff
Task
EM
F1
EM
F1
14.29
13.54
18.80
19,96
29.29
31.41
39.99
41.02
32.60
34.65
42.74
44.35
35.94
36.17
48.82
50.18
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 de- viations not exceeding 2322 in any category.
In particular, 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 G.+1 G, is much smaller than G.+1. There are on average 3.42 triples or 10.42 tokens per state across the VerichoWorke test dataset for G - G, but a mean of 8.71 triples or 26. 13 tokens per state for G.+1. This also explains the increased performance of the GATA-W W over the baseline Seq2Seq agent-this 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.
5.2 Valid Action Generation
Similarly to the other task, we compare the Worldformer prediction. All models use a fixed vocabulary of size 1 11,056 at train and test times.
existing baseline for valid action
Seq2Seq. This single-task model is provided as baseline by Jericho 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-actioon pairs created by humans on the ClubFloyd online Interactive Narrative forum appears in both Ammanabrolu and Hausknecht 2 and Yao et al.
339 with the latter using it to tune GPT-2 model for valid action prediction using a GPT-2 based Seq2Seq valid action model dubbed CALM This model takes in O,, A, O.+1 and targets V141-
In Table I we see that the Worldformer significantly outperforms the Seq2Seq baseline on all the games and CALM overall. Fach valid action in text game requires 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
http://www..lllhiiiigeqememmnnnt http:///cccccocccctcccttcccioono fiction.html
8

---

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 Chih is trained on a different dataset-arr are comparable on F1 scores but Seq2Seq is better overall for exact matches. CALM also has relatively higher variance in performance across the test on some games such as zork/ and detective it outperforms
games than the other two methods the Seq2Seq and is not too far off the Worldformer especially in terms Fl 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 compo- nents 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 swggessgg 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 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 navigation actions before learning more fine-graing actions. E.g. zzuu, deephome, and balances have a high number of gold standard average valid actions while pentari, ludicorp, detective, and temple which have a low numher of average valid actions. While the latter set of games have generally higher performance both the Seq2Seq and Worldformer models, the gap is significantly less pronounced with the Worldformer. We hypothesize that this is due to the representations now contain enough information
Ablation
Act
Multi Task
SOS
EM
Loss
18.10
19,44
20.78
22.42
20.12
21.28
23.22
25.54
Table 3:
Worldformer abla-
tions to test the impact of its two main components for ac- tion prediction. All results are size weighted averages over all test games over three ran- dom seeds, with standard devi- ations not exceeding 1..
multi-task training of the Worldformer- regarding the next knowledge graph to alleviates the label imbalance of the actions and enable prediction of more fine-grained actions.
Conclusions
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 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 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.
7 Broader Impacts
We view text-games as an platform on which 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 BerichoWorld-ant further verified by us ds-data is collected from games containing situations of non-normative language laage-deserrbingg 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 [3] and Jericho [4. This is based on manual vetting and (existing) crowd-sourced reviews on the popular interactive narrative forum IFDB (https://ttdb.org/
Broadly speaking, the most relevant downstream task for this work is model- 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 lakks-downstreal 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).

---

References
[1] A. Adhikari, X. Yuan, M.-A. Côté, M. Zelinka, M.-A. Rondeau, R. Laroche, P. Poupart, J. Tang.
A. Trischler, and W. Hamilton. Learning dynamic belief graphs to generalize on text-based games. Advances in Neural Information Processing Systems, 33, 2020.
[2] P. Ammanabrolu and M. Hausknecht. Graph Constrained Reinforcement Learning for Natural Language Action Spaces. In International Conference on Learning Representations, 2020. URL
https:/opeeneeiiewnefffoooo
[3] P. Ammanabrolu and M. O. Riedl. Playing text-adventure games with graph-based deep reinforcement learning- Proceedings of 2019 Annual Conference of the North American Chapter of the Association for Computational Linguistics; Human Language Technologies, NAACL-HLT 2019, 2019.
[4] Ammanabrolu and M. O. Riedl. Modeling worlds in text. OpenReview Preprint, 2021. URL
http://ppnnreviee.ee///oromm "oce/ooummdddYYYYYSMAAAY"
[5] P. Ammanabrolu, E. Tien, M. Hausknecht, and M. O. Riedl. How avoid being eaten by a grue: Structured exploration strategies for textual worlds. arXiv preprint arXiv:2006.074000 2020.
[6] G. Angeli, J. Premkumar, M. Jose, and C. D. Manning. Leveraging Linguistic Structure For Open Domain Information Extraction. In Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), 2015.
[7] K. Arulkumaran, M. P. Deisenroth, M. Brundage, and A. A. Bharath. Deep reinforcement learning: A brief survey. IEEE Signal Processing Magazine, 34(6):26-38, 2017. doi: 10. 10.1109/ MSP:2017.27432200.
[8] G. Brockman, V. Cheung, I.. Pettersson, Schneider, J. Schulman, Tang, and W. Zaremba. Openai gym, 2016.
[9] S. Dambekodi, S. Frazier, P. Ammanabrolu, and M. O. Riedl. Playing text-based games with common sense. arXiv preprint arXiv:2012.02757, 2020.
[10] J. Devlin, M. Chang, K. Lee, and K. Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs/1810. 2018.
[11] M. Guzdial, B. Harrison, B. Li, and M. Riedl. Crowdsourcing Open Interactive Narrative. In FDG, 2015.
[12] D. Ha and J. Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https Pape-rr...
[13] M. Hausknecht, P. Ammanabrolu, M.-A. Côté, and X. Yuan. Interactive fiction games: A 2020. colossal adventure. In Thirty-Fourth AAAI Conference on Artificial Intelligence (AAAI), 2020. URL httph:s///arxiwworr//abba wttppa/s////0999995999
[14] J. He, J. Chen, X. He, J. Gao, L. Li, L. Deng, and M. Ostendorf. Deep reinforcement lcarning with natural language action space. In ACL, 2016.
[15] D. Jancke. Orientation formed by spot's trajectory: two-dimensional population approach in primary visual cortex. Journal of Neuroscience, 20(14): 2000.
[16] T. Kipf, E. van der Pol, and M. Welling. Contrastive learning of structured world models. In International Conference on Learning Representations, 2020. URL https://openneeiewwneell
Vorum?id=HIl
gax6V1DB gaxxVio
[17] T. N. Kipf and M. Welling, Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations (ICLR), 2017.
[18] Z. Lan, M. Chen, S. Goodman, K. Gimpel, Sharma, and R. Soricut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations, 2020. URL,
[19] Y. Li, O. Vinyals, C. Dyer, R. Pascanu, and P. Battaglia. Learning deep generative models of 2018.
graphs. arXiv preprint
10

---

[20] G. A. Miller. WordNet: A Lexical Database for English. Commmiccaiions of the ACM, 38(11):
39-41, 1995.
[21] K. Murugesan, M. Atzeni, P. Shukla, M. Sachan, P. Kapanipathi, and K. Talamadupula. Enhanc- ing text-based reinforcement learning agents with commonsense knowledge. arXiv preprint arXiv:20050088/ 2020.
[22] K. Murugesan, M. Atzeni, P. Kapanipathi. P. Shukla, S. Kumaravel. G. Tesauro, K. Tala- madupula, M. Sachan, and M. Campbell. Text-based RL Agents with Commonsense Knowl- edge: New Challenges, Environments and Baselines. In Thirty Fifth AAAI Conference on Artificial Intelligence, 2021.
[23] M. S. A. Nahian, S. Frazier, M. Riedl, and B. Harrison. Learning norms from stories: A prior for value aligned agents. In Proceedings of the AAAI/ACM Conference on AI, Ethics, and Society, AIES '20, page 124-130, New York, NY. USA, 2020, Association for Computing Machinery. ISBN 978 1450371 100. doi: 10. 1LL14//337562733355222 URL
3375825
[24] K. Narasimhan, T. D. Kulkarni, and R. Barzilay. Language understanding for text-based games using deep reinforcement learning. In EMNLP, pages 1-11, 2015.
[25] Niu, C. Zhang, C. Ré, and Shavlik. Elementary: Large-scale knowledge-base construction via machine learning and statistical inference. International Journal on Semantic Web and Information Systems (LISWIS), 8(3):42-73, 2012.
[26] J. Oh, X. Guo, H. Lec, R. L. Lewis, and S. Singh, Action-Conditionll Video Prediction using Deep Networks in Atari Games. In C. Cortes, N. D. Lawrence, D. D. Lee, M. Sugiyama, R. Garnett, and R. Garnett, editors, Advances in Neural Information Processing Systems 28, pages 2345-2853. Curran Associates, Inc., 2015.
[27] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and Sutskever. Language Models are Unsupervised Multitask Learners. 2019.
[28] P. Rajpurkar, R. Jia, and P. Liang. Know what you don't know: Unanswerable questions for SQuAD. In Proceedings of the 56th Annual Meeting of the Association for Com- putational Linguistics (Volume 2: Short Papers), pages 784-789, Melbourne, Australia, July 2018. Association for Computational Linguistics. doi: 10. 1/118//2///111 URL
[29] C. Sautier, J.
J. Agravante, and M. Tatsubori. State Prediction in TextWorld with a Predicate- Logic Pointer Network Architecture. In In Workshop Knowledge-bused Reinforement Learning at IJCAI-20, 2020. URL https::////bbbttgiiiiee
[30] Schrittwieser, I. Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lock- hart, D. Hassabis, T. Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. arXiv preprint 111::1911.00265: 2019.
[31] E. Sheng, K.-W. Chang, P. Natarajan, and N. Peng. The woman worked as a babysitter: On biases in language generation. Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing (EMNLP-IJCNLP), pages 3407-3412. Association for Computational Linguistics, 2019. doi: 10. 11.118653/1/D111.. 1339. URL http://wwww.wwlwebb
org/anthology/I19-1333.
[32] I. SStskever, O. Vinyals, and Q Le. Sequence to sequence learning with neural networks. In Proceedings of the 27th International Conference on Neural Information Processing Systems- Volume 2, NIPS' 14, page 3104-3112, Cambridge, MA, USA, 2014. MIT Press.
[33] R. S. Sutton and A. G. Barto. Introduction to Reinforcement Learning. MIT Press, Cambridge, MA, USA, 1st edition, 1998. ISBN 0262193993
[34] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, t.. Kaiser, and
I. Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
[35] N. Wahlström, T. B. Schön, and M. P. Deisenroth. From pixels to torques: Policy learning with deep dynamical models. arXiv preprint arXiv: 1502.0250 2015.
[36] Q Wang, Z. Mao, B. Wang, and L. Guo. Knowledge graph embedding: A survey of approaches and applications. IEEE Transactions on Knowledge and Data Engineering, 29(12):2724-274) 2017.
11

---

[37] M. Watter, J. Springenberg. J. Boedecker, and M. Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In C. Cortes, N. Lawrence, Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015. URL httss//ppoeeddnnnssrrrrsssc///op (ate55cccSS4A4aS5722990016727662
[38] Z. Wu, S. Pan, F. Chen, G. Long, C. Zhang. and P. S. Yu. A comprehensive survey graph neural networks. IEEE Transactions Neural Networks and Learning Systems, 32(1):4-24, 2021. doi: 10.
[39] S. Yao, R. Rao, M. Hauskmecht, and K. Narasimhan. Keep CALM and explore: Language models for action generation in text-based games. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 8736-8754, Online, Nov. 2020. Association for Computational Linguistics. doi:
wtw.:0000..mmmmmmmi...0..0 704.
URL http://wwwwwwo//////oo//////////
[40] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. Smola. Deep sets. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vish- wanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, vol- ume 30. Curran Associates, Inc., 2017. URL https://poooeednnsss.uo../c////p 122e4747da 4 Paper..df
12


Reviews:
[{"review_id": "wpQ0HmAizGP", "reviewer": "Reviewer_FHsa", "summary": "This paper proposes a novel knowledge graph construction task based on attribute extraction from text game engines; and a set-of-sequence generation model that achieves good performance on the proposed tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper proposes a novel knowledge graph construction task based on attribute extraction from text game engines. A set-of-sequence generation model, the Worldformer is proposed, which achieves good performance on the proposed tasks. I think this is a very clever way of creating knowledge graph construction tasks. And the two proposed tasks, on their own, are valid and interesting NLP tasks (though the human study might be needed, see my review at the end).\n\nThere are several major weaknesses this paper should addresses:\n\n* Clarity\n\nAt first I did not find how the groundtruth graphs G_t and G_{t+1} are constructed. Then following the provided github repo I found this paper: https://openreview.net/pdf?id=Y1YtS9MZA75. Unfortunately, the dataset description paper still does not provide sufficient details. I guess the two papers are written by the same authors. If so, it would be good if more details about the dataset construction can be provided in the revised version. Specifically, how the tuples are extracted? Also why the average graph tuple number in a game is usually so small?\n\n* The significance of the graphs prediction task\n\nFirst, I did not see any theoretical or empirical justifications of the values of the proposed world model representation. Why this is the right world model for understanding the game texts? Is this world model better than simple sequence embedding of history? Why is it better than the model in [4]? The missing of the above answers makes the motivation of this paper not quite clear, as the addressed problem has not been well-verified yet according to the submission. \n\nSecond, the proposed model does not show a clear advantage over the Q*BERT baseline. The baseline formulates the graph construction task as question answering, therefore it is not surprising the baseline does not perform well on graph metrics but works well on token metrics. This raises another problem of the task significance -- the used intermediate metrics cannot clearly show which model is the best. It would make more sense if the graph construction models were used to enhance the game-playing agents, and the agents' performance was reported.\n\n* Empirical comparison\n\nRelated to the previous bullet, the compared Q*BERT baseline was proposed to enhance a game-playing agent and performed quite well. As a fair comparison, the proposed Worldformer should be compared in the same setting.\n\nThe submission cited the related work [39], which was also proposed for valid action prediction. I am wondering why the work [39] was not compared in the experiments? Did I miss anything?\n\n* A minor question:\n\nAre the textual inputs guaranteed to be sufficient for the two tasks? I think conducting a human study would be necessary to show the quality of the tasks and their upperbound performance.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "h9ghy8DPepB", "reviewer": "Reviewer_cdHj", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper presents an original approach for modelling state and valid action spaces for text adventure games using pretrained language models. The work is clearly presented and well-motivated. My main concern has to do with the evaluation:\n\n1. Is this estimation helpful for actual game-playing?\nGiven that the ultimate goal is to be more sample efficient due to having a good world model, I think it would be good to use this world-model to learn policies for these games, the the hypothesis being that using these models improve upon prior results. The numbers in Table 1 are all fairly low IMO and I am not confident that these results are actually helpful downstream.\n\n2. I am unclear as to where the gains are coming from. Some ablations that I find missing include:\n- QA & Seq2Seq with multitasking\n- QA & Seq2Seq with set loss\n\n3. What do the evaluation metrics mean? It is difficult to interpret F1 in the context of text adventure games. For example, what portion of the actions generated are actually parseable by the game parser?\n\nSome questions for the authors:\n\n- what are \"lifted representations\"? L38\n- does this algorithm support removing nodes from the graph? some parts of the paper imply yes (L243) but the main exposition suggests no. L150\n- is phrase-level masking actually better than normal LM masking? L165", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": ["I don’t know"], "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "_DvX0kirY-r", "reviewer": "Reviewer_swGy", "summary": "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.\n\n**Challenges**: Partial observability, combinatorial action space", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Major Comments:**\n1. It would be helpful if the sample valid actions generated for some of the observations in the test set and the corresponding ground truth valid actions are shown side-by-side for comparison along with the graph additions and deletions. In addition to the numbers shown in Tables 2 and 3, qualitative examples could be very helpful in understanding the model.\n\n2.  The results in the experiment section show how well the worldformer performs on the KG prediction and valid action prediction tasks for the held-out set. Does using this model for playing the games show any advantage over the other baselines in terms of scores/rewards? How will these learned world models be used for text-based games?\n\n3. Does the checkmark in multitask mean the KG prediction and valid action task were learned independently?\n\n\n**Minor Comments:**\n1. Type in line 260 - table 2\n2. Line 286 - 5 tokens", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the broader statement explains the necessary details on the limitations and societal impacts.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "IOnWcVOAgfe", "reviewer": "Reviewer_QgLj", "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. \n\nTo learn the world model, the authors design two tasks:\n- 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}$.\n-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. \n\nSince 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]). \n\nAnother 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. \n\n", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, this is a very interesting paper. The authors aim to learn a world model represents a model's understanding to the environment it is performing in. World model is a hot topic in model-based RL community, as cited by the authors, state prediction and action prediction seem to be a promising direction towards learning a \"model\". \n\nThe authors tackle the problem by using a static dataset, JerichoWorld, a set of trajectories along with transitions at every game step collected from a set of text-based games. I like this choice because 1) it circumvents extra difficulties in the interactive settings so that they can focus on the world model learning part, 2) but it can also be seen as an offline-RL setting since the static dataset can be treated as what an online agent collected and cached into the replay buffer. \n\nI like how the authors formulate the world model as a graph --- there are prior works (such as Ammanabrolu and Riedl, 2019, Adhikari et al., 2020) suggest graph structures help text game agents as a memory and facilitate its planning. There is also a line of prior work focuses on representation learning in RL (e.g., Schwarzer et al., 2021), especially with contrastive predictive methods. This work in some sense brings the two lines of research together, explores how to learn a world model, in the format of a graph, by training with action and state prediction. \n\nI agree with the authors that the valid action prediction and state prediction tasks are set generation tasks. The authors attempts to modify the standard teacher-forcing paradigm to make it work better on such set generation tasks. \n\nThe paper is clearly written, it is easy to follow. \n\nI will put my questions and concerns in two parts:\n## General (maybe highlevel) questions:\n\n**(Q1)** Are the inputs to the world model sufficient? To answer this question, please first check if my understanding to the model input/output are correct (as listed in the summary review). As mentioned by the authors (L29-L30), text-based game environments are often partially observable. If that is the case, how does one predict $S_{t+1}$ from $O_t$, $V_t$, $G_t$ and $A$? \n\nI can understand that within the current observation range, it is possible to predict the consequence of performing an action (e.g., take leaflet may cause the leaflet moving from mailbox to inventory). However, how does one predict something not in their observation range? For instance, in Figure 1, given all the information on left, plus the $A$ (go north), how to predict the right hand side information? How does the agent know what place is to the north, and what objects that place contains? I do expect a Seq2Seq model being able to memorize this from the training games, but I do not see how this can be generalized to unseen data points. \n\n**(Q2)** If the inputs are somehow sufficient, are they all necessary? In the knowledge graph generation task proposed by the authors, the model learns: $O_t, V_t, G_t, A \\rightarrow \\delta$ where $\\delta = G_{t+1} - G_t$. I fail to understand why $O_t$ and $V_t$ is necessary here. Doesn't $G_t$ contain all information in $O_t$? Is this redundant information by design? Similarly, in the valid action generation task, can a model learn to deduce $V_{t+1}$ from $G_{t+1}$ only? If so, I could imagine a reasoning path of $G_t, A \\rightarrow \\delta  \\rightarrow G_{t+1}  \\rightarrow V_{t+1}$.\n\n## Detailed questions:\n\n**(Q3)** What reason is behind the choice of representing graphs as sequences of triples, compared to the (probably) more straightforward practice of adjacency matrix and Graph Neural Nets? Presumably, the model should learn an permutation invariant representation of set. However, there might be ordering biases introduced by the current way of encoding the set of triples? Did the authors do special modifications to the positional embeddings in their Transformer-based encoders (e.g., reset positions at the first token of each triple/action)?\n\n**(Q4)** How is $A$ integrated in the model? I fail to find it from Figure 3. With Figure 3 alone, the inputs to outputs mapping seems impossible. \n\n**(Q5)** How does the proposed system learn to generalize to unseen objects and their affordances? Games in the Jericho collection can be quite different from each other, I wonder how do the proposed methods generalize to unseen and potentially very different games, where unseen objects have unseen affordances. It would be also nice to investigate deeper on the numbers shown in Table 1, what contributed to those numbers? What can be transferred across games and what not (i.e., what's an upperbound EM/F1 score for JerichoWorld)?\n\n**(Q6)** I am curious to see a detailed analysis on the set of sequences (SOS) generation method. Ablation results in Table 2 and Table 3 suggest the SOS loss provide a non-trivial boost on performance, however, in experiments without SOS loss, do the authors sort the sub-sequences somehow? In Adhikari et al., they seem to suggest that sorting the sub-sequences are necessary in Seq2Seq learning. In addition, I am not intuitively convinced that cutting the dependencies between set members is the best way because there is arguably some kind of dependencies existing within a set (even to the least extent, set members should not be identical). \n\nThat said, what we want is essentially permutation invariant, rather than removing dependency. Given how the SOS loss is designed, a simple failure mode I can imagine is that the model may have hard time generating the first token in each of the triples/actions, because they all condition on the same special token [TRIPLE]/[ACT], as a result, does this lead to more duplicates during inference? I suggest the authors to provide more information to make this more convincing. Also, Ye et al., 2021 seems quite related, might worth checking. \n\n**(Q7)** (minor point) It would be nice to see the learned world models are helpful in the interactive setting of text-based games rather than their static dataset version.\n\n## Typos:\n1. L47: that that.\n2. L248: other other.\n3. L260: Table 3 --> Table 2.\n\nReferences:\n\n1. Playing Text-Adventure Games with Graph-Based Deep Reinforcement Learning. Ammanabrolu and Riedl, NAACL 2019.\n2. Learning Dynamic Belief Graphs to Generalize on Text-Based Games. Adhikari et al., NeurIPS 2020.\n3. Data-Efficient Reinforcement Learning with Self-Predictive Representations. Schwarzer et al., ICLR 2021.\n4. One2Set: Generating Diverse Keyphrases as a Set. Ye et al., ACL 2021. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors have discussed limitations and potential negative societal impact. I do not see a problem here.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper jointly learns a policy (action distribution) and world model (updates to a graph representation of the environment) in a text adventure. While interesting, I'm not seeing the obvious application of this (despite having read the ethics impact section of the paper), and have some concerns about text adventure games as a domain — however the meta-review stage of the review process is not an appropriate point for me to bring these concerns to the authors' attention, so I will happily discount them and focus on the reviews and discussion.

The reviewer consensus leans towards acceptance, after some discussion. I would have liked to see a follow-up response from reviewer FHsa to the author rebuttal, and have no evidence that they have read it or considered updating their score despite my prompts, which is a bit disappointing. Less crucially, it would have been nice to get an author response to the follow-on comments by reviewer cdHj, but this is not as important since that reviewer updated their score to recommend acceptance.

Ultimately, when discounting the score of reviewer FHsa, the median and mean scores are in acceptance territory. I confess I do not fully understand reviewer FHsa's argument against the paper, and felt the author response was detailed. In the absence of this reviewer's willingness to defend their appraisal, I recommend acceptance.
Decision: Accept (Poster)