PLM -based World Models for Text-based Games (2022)

Abstract: World models have improved the ability of reinforcement learning agents to operate in a sample efficient manner, by being trained to predict plausible changes in the underlying environment. As the core tasks of world models are future prediction and commonsense understanding, our claim is that pre-trained language models (PLMs) already provide a strong base upon which to build world models. Worldformer is a recently proposed world model for text-based game environments, based only partially on PLM and transformers. Our distinction is to fully leverage PLMs as actionable world models in text-based game environments, by reformulating generation as constrained decoding which decomposes actions into verb templates and objects. We show that our model improves future valid action prediction and graph change prediction. Additionally, we show that our model better reflects commonsense than standard PLM.

## introduction in model-based reinforcement learning, world models @xcite , being trained to predict plausible changes in the underlying environment, has helped agents to quickly identify sensible or high-value actions. as a result, agents equipped with world models have outperformed state-of-the art model-free algorithms in atari games, while drastically improving sample efficiency @xcite . in this work, we focus on world modeling in textbased game (tbg) environments, in which players and agents must perceive and interact with the world entirely through textual natural language. as such, they present several unique challenges @xcite ; ammanabrolu bart encoder-decoder denoted in blue learns the future valid actions prediction task, while the one denoted in orange simultaneously learns the graph difference prediction task. the light yellow blocks denote generation constrained by a template, bringing generation closer to the target controlled sublanguage of the text-based game's parser. and riedl, 2021b): aside from language understanding itself, tbgs require dealing with an exponential action space: for instance, the action space size for zork1, combining five words from a vocabulary of 697 words, is o(10 14 ). tbgs also require rich and accurate knowledge representations of locations and objects, in order to facilitate navigation and interaction. finally, solving tbgs requires commonsense reasoning, including understanding object affordance, and the causal ramifications of actions. the first dataset and model for learning world models in tbgs were proposed by jeri-choworld @xcite , and worldformer @xcite , respectively. intuitively, large pre-trained language models (plms), make promising candidates for instantiating world models @xcite , as they are already trained to generate language in a much larger action space, e.g. o(40, 000 512 ), assuming a standard vocabulary of 40, 000 tokens, and context length of 512 tokens. while promising, adapting plms to tbgs is non-trivial, as in tbgs, generated actions must be executable in the game environment. while we are not generating in a formal syntax as in semantic parsing, we are still constrained by a controlled sublanguage @xcite , which is closer to natural language but follows a grammar defined by the engine's parser. furthermore, a world modeling task such as the valid action decoding task of jerichoworld, requires predicting, from a particular world state, all future valid actions. in contrast, semantic parsing focuses on learning a one-to-one mapping between natural language and the controlled sublanguage. to encapsulate these challenges, we define a concept of actionability as the main objective of the valid action generation task of jerichoworld, in that the model's generation should be actionable in the tbg environment, by 1) conforming to the parseable controlled sublanguage, and 2) ensuring the generation of valid actions consistent with the commonsense governing the dynamics of the tbg environment. in this work, we focus on improving the actionability of the world model's generations, focusing in particular on bringing the plm's commonsense and reasoning capabilities to bear on world modeling tasks. we begin by building worldformer-bart, an implementation of worldformer based on bart. while successful, commonsense in plm-based world models are not easily transferred to more formal forms, such as that of the controlled sublanguage of tbgs. motivated by these findings, we build a retrieval-augmented model which aims to enhance actionability by formalizing commonsense through templates, hence named actionable world model (awm). we show through experimental results that awm-bart significantly improves actionability over worldformer-bart. we also compare with augmenting trained models with a comet-based commonsense filter, showing that our model outperforms such approaches. ## preliminaries we now describe our base models in detail. we begin with a description of worldformer, which forms the basis of our model architectures. ## awm-bart for actionable world model we now propose our model, which aims to improve the actionability of the bart-based world model. we decompose action generation as template selection and filling, aiming to capture two benefits: through input-constrained decoding using templates, we enhance the parseability of the world model's generations. furthermore, we posit that the inductive bias from templates will aid the world model to learn more accurate object affordance and object localization, improving commonsense. more specifically, bart dec action now performs the task of generating the sequence of possible fillings y t @xmath0, ..., y t j n }, @xmath1, ..., y tm } of a template t @xmath2, of a jericho game environment. then, the template-conditional action generation task is reformulated as follows: the mask filling task is illustrated in fig. 1 . the above formulation is advantageous in that the task of filling the masks of an action template brings the generation objective close to bart's original objective. however, it does not address the issue of choosing an appropriate template to fill. we next describe how we employ multitask learning to utilize a single bart encoder-decoder as both a template retrieval and generation model. ## experiments we evaluate our models on the jerichoworld modeling tasks. ## results we report the results of our experiments in ## conclusion in this work, we build world models based on plms. towards leveraging semantic and linguistic priors learned by plms, we identified major areas in which plms need to be systematically improved, namely generation to a controlled sublanguage, and commonsense understanding. we propose an actionable world model based on templateaugmented generation, showing that both parseability and commonsense understanding can be significantly improved. as future work, we consider combining plm-based world models with rewardbased agents to learn goal-directed policies as a promising direction. finally, our hope is that our work will contribute to the active exploration of text-based games as an alternative testbed for further research on commonsense reasoning.