Because all policies of a model-based agent are derived from a shared world model, any modifications made to this model would affect all of them. This feature can be exploited to reduce human effort in controlling this type of agent. Specifically, suppose we concern $m$ tasks in the environment, necessitating $m$ policies. If there is a change in the environment dynamics, a model-based agent only needs task-agnostic data to replicate this change in its world model. It can then re-optimize its policies with respect to the updated model. Meanwhile, a model-free agent needs to collect task-specific data to re-train all of its $m$ policies. The data collection cost of the model-free approach scales with $m$, whereas that of the model-based approach is independent of $m$, since the policy re-optimization step uses only data generated by the world model.

The dominant approach to world modeling learns a function $M_{\theta}(s_{t+1},r_{t+1},d_{t+1}\mid h_{t})$ parameterized by a neural network $\theta$ and conditioned on a history $h_{t}=(s_{1},r_{1},d_{1},a_{1},\ldots,s_{t},r_{t},d_{t},a_{t})$. We refer to this class of models as observational world models because they can be adapted with only observational data, through either in-weight learning (updating the model parameters to fit a dataset of observations), or in-context learning (plugging in a history of observations).

Relying on observation-based adaptation leads to two drawbacks. First, controlling these models is difficult because observations are inadequate for conveying complex, abstract human intentions. Second, collecting observations requires taking real actions in the environment, which can be expensive, time-consuming, and risky.

In Messenger, a state $s$ is represented by an $H\times W$ grid with $C$ channels (an $H\times W\times C$ tensor), where each channel corresponds to an entity. In each channel $c$, there is a single non-zero cell $s(h,w,c)$ that represents the identity of the entity. The position of this cell is the location of the entity in the grid. We note that this is an idealized representation that disentangles the entities. Even so, the problem remains challenging, as the model needs to recognize attributes mentioned in the manual and associate them with the right entity token. This requires a special attention mechanism, which we will introduce shortly. Meanwhile, learning entity-disentangled representations for pixel-based environments remain an open problem, which we defer to future work.

Our model (illustrated in 1(b)) is an encoder-decoder Transformer (Vaswani et al., 2017) which encodes a manual $\bm{\ell}$ and decodes a trajectory $\tau$. We transform the trajectory into a long sequence of tokens and train the model as a sequence generator.

Concretely, our model processes a data point $(\tau,\bm{\ell})$ as follows. For the manual $\bm{\ell}=\{l_{i}\}_{i=1}^{N}$, we first use a pre-trained BERT model to convert each description $l_{i}$ into a sequence of hidden vectors. We feed each sequence to a Transformer encoder, which outputs a tensor $\bm{m}^{\textrm{enc}}$ of size $N\times L\times D$, where $N=C$ is the number of descriptions, $L$ is the maximum number of words in a description, and $D$ is the hidden size.

For the trajectory, we convert each tuple $(a_{t-1},s_{t},r_{t},d_{t})$ into a token block $B_{t}$. The first action $a_{0}$ is set to be a special <s> token. Each state $s_{t}$ is mapped to $3C$ tokens $(i_{t}^{1},h^{1}_{t},w^{1}_{t},\cdots,i^{C}_{t},h^{C}_{t},w^{C}_{t})$, which represents each of the $C$ entities by its identity $i$ followed by its location $(h,w)$. The real-valued reward $r_{t}$ is discretized into an integer label, and the termination signal $d_{t}$ is translated into a binary label. In the end, $B_{t}$ consists of $3C+3$ tokens $(a_{t-1},i_{t}^{1},h^{1}_{t},w^{1}_{t},\cdots,i^{C}_{t},h^{C}_{t},w^{C}_{t},r_{t},d_{t})$. Finally, we concatenate all $T$ blocks in the trajectory into a sequence of $T\times(3C+3)$ tokens, embed them into a $T\times(3C+3)\times D$ tensor, and add positional embeddings. We will use bold notation (e.g., $\bm{a}$, $\bm{i}$) to refer to the resultant embeddings of the tokens.

We implement a variant of EMMA (Hanjie et al. (2021)) that first identifies the description that mentions each entity and extracts from it words corresponding to the attributes of the entity. From the tensor $\bm{m}_{n}^{\textrm{enc}}$ computed by the encoder, we generate a key tensor $\bm{m}^{\textrm{key}}$ and a value tensor $\bm{m}^{\textrm{val}}$, both of which are of size $N\times L\times D$, where

for $1\leq n\leq N$. Here, $\texttt{Linear}^{D\rightarrow 1}_{\textrm{key}}$ and $\texttt{Linear}^{D\rightarrow 1}_{\textrm{val}}$ are linear layers that transform the input’s last dimension from $D$ to 1, and $\texttt{Softmax}(\cdot)$ applies the softmax function to the last dimension. Intuitively, we want each $\bm{m}^{\textrm{key}}_{n}$ to retain words that signal the identity of the entity mentioned in the $n$-th description (e.g., ferry, plane, researcher), and $\bm{m}^{\textrm{val}}_{n}$ to retrieve words depicting the other attributes (e.g., approaching, deadly, fleeing).

Let $\bm{i}^{c}_{t}$ be the embedding of the identity of entity $c$. We perform a dot-product attention with $\bm{i}^{c}_{t}$ as the query, $\bm{m}^{\textrm{key}}$ as the set of keys, and $\bm{m}^{\textrm{val}}$ as the set of values to compute the attribute features of $c$

The features are added to the identity tokens $\bm{i}^{c}_{t}$. The final input of the model is as follows:

Unlike the standard encoder-decoder Transformer, our architecture does not perform cross-attention between the encoder and the decoder because information from the encoder has already been incorporated into the decoder through EMMA.

We train the model to minimize cross-entropy loss with respect to the ground-truth (tokenized) trajectories in the training set. The label at each output position is the next token in the ground-truth sequence. In particular, we do not compute the losses at the positions of the action tokens and the first block’s tokens, because those tokens will be set during inference.

World models have a rich history dating back to the 1980s (Werbos, 1987). The base architecture has evolved from feed-forward neural networks (Werbos, 1987), to recurrent neural networks (Schmidhuber, 1990a, b, 1991), and most recently, Transformers (Robine et al., 2023; Micheli et al., 2023). In RL settings, world models are the key component of model-based approaches, which train policies in simulation to reduce the amount of interactions with real environments. Model-based RL has been successful in a variety of robotic tasks (Finn and Levine, 2017) and video games (Hafner et al., 2019, 2020, 2023). However, the incorporation of language information into world models has been underexplored. Cowen-Rivers and Naradowsky (2020) propose language-conditioned world models but focus on emergent language rather than human language. Poudel et al. (2023) incorporate features language into the representations of the model. These approaches, however, do not use language to control a world model.

Language information has been incorporated into various aspects of learning. In instruction following (Bisk et al., 2016; Misra et al., 2018; Anderson et al., 2018; Nguyen and Daumé III, 2019), agents are given descriptions of the desired behaviors and learn to interpret them to perform tasks. Language-based learning (Nguyen et al., 2021; Scheurer et al., 2023) employs language-based feedback to train models. Another line of work uses language descriptions of environment dynamics to improve policy learning (Narasimhan et al., 2018; Branavan, 2012; Hanjie et al., 2021; Wu et al., 2023a; Nottingham et al., 2022; Zhong et al., 2020). Rather than using texts to directly improve a policy, our work leverages them to enhance a model of an environment. Recently, several papers propose agents that can read text manuals to play games (Wu et al., 2023a, b). Our work differs from these papers in that we aim to build models that capture exactly the transition function of an environment.

Lin et al. (2024) model a variety of text-augmented environments but do not demonstrate the generalizability of their approach in Messenger. Recent work (Zhao et al., 2022; Du et al., 2024; Zhou et al., 2024; Zhang et al., 2024) has developed LWMs with compositional generalizability. While these papers operate on more visually realistic domains than ours, the language they study is simpler, focusing on concepts that correspond to straightforward mappings from input to output such as colors and objects. In contrast, the concepts in Messenger are more intricate, regarding interactions among multiple entities.