Reader: Model-based language-instructed reinforcement learning (2023)

Abstract: We explore how we can build accurate world models, which are partially specified by language, and how we can plan with them in the face of novelty and uncertainty. We propose the first model-based reinforcement learning approach to tackle the environment Read To Fight Monsters (Zhong et al., 2019), a grounded policy learning problem. In RTFM an agent has to reason over a set of rules and a goal, both described in a language manual, and the observations, while taking into account the uncertainty arising from the stochasticity of the environment, in order to generalize successfully its policy to test episodes. We demonstrate the superior performance and sample efficiency of our model-based approach to the existing model-free SOTA agents in eight variants of RTFM. Furthermore, we show how the agent’s plans can be inspected, which represents progress towards more interpretable agents.

## introduction intelligent agents have the ability of re-composing known concepts to draw conclusions about new problems and this translates into the acquisition of very robust and general behaviours. current reinforcement learning (rl) agents typically lack this ability and they need to be re-trained for every new problem; in contrast language models exhibit greater generalization abilities and provide a more flexible approach to solve multiple tasks with a single model. thus language-conditioned rl is a flourishing area of research. on the other hand, language models are trained exclusively on textual inputs and struggle to ground the meaning of the words to real world dynamics. multiple interactive environments have been proposed as a testbed for learning how to ground language @xcite @xcite . while prior work mostly focuses on behavioural cloning or model-free rl, we argue for a model-based reinforcement learning (mbrl) approach, as this effectively decouples the problem of understanding how the world works from the problem of acting optimally in the world in order to solve one or more tasks. concretely, mbrl inherits the advantages of model-free rl of learning from scratch or from sub-optimal behaviour, while being orders of magnitude more sample efficient than the model-free counterpart. furthermore, it has the added value of being more interpretable and explainable. in fact, a decision made by a mbrl agent can be accompanied by human-interpretable examples of likely future trajectories that are taken into account by the model in making such a decision. in this work, we focus on read to fight monsters (rtfm), a challenging benchmark for testing grounded language understanding in the context of rl proposed by @xcite . rtfm tests the acquisition of complex reading skills in rl agents in order to solve completely new tasks based on written descriptions of the task dynamics and goal. critically, the written information provided is not enough on its own to obtain an optimal policy, but the agent needs to cross-reference it multiple times with the current state of the environment, in order to figure out a plan of action. we make the following contributions: first, we formulate a language-instructed mbrl method for solving rtfm and show how to train an agent in this environment (see fig. 1 ). our method, named reader (for reinforcement learning agent for discrete environments with written instructions), explicitly models the stochastic changes in the discrete environment and performs planning with a stochastic variant of monte carlo tree search (mcts, @xcite . we then demonstrate the effectiveness of reader on eight variants of rtfm, showing better performance and sample efficiency than the sota model-free agent txt2π from @xcite : reader surpasses txt2π on all the eight rtfm variants when both methods are trained on 1m frames of experience and is better on five and on par on the remaining three even when txt2π is trained on 100 times more frames foot_0 . furthermore, reader has comparable performance with a strong model-based transformer baseline, while being more than eight times faster during planning. finally, we show how reader's plans can be visualised and how those plans change accordingly to the written instructions, making the language-instructed model-based approach more interpretable and trustworthy than the model-free counterpart. ## read to fight monsters read to fight monsters (rtfm) is a challenging benchmark proposed by @xcite for testing grounded language understanding in the context of rl. rtfm tests the acquisition of complex reading skills in rl agents in order to solve completely new tasks based on written descriptions of the task dynamics (also called manual m) and goal g. crucially, it is not enough to consult the written information in order to obtain an optimal policy, but the agent needs to perform a multi-step reasoning between such information and the current state of the environment in order to figure out a plan of action. to elucidate the reasoning steps and reading skills needed to win an episode, we go through the concrete example reported in fig. 2 . 1. from the goal extract which team to defeat (order of the forest). the agent is given a reward of +1 if it engages the correct monster in combat while carrying the correct item, -1 in any other encounter with a monster and a reward of 0 for all intermediate steps. as every episode contains a procedurally generated set of (monster, element) pairs, (item, modifier) pairs, goal and manual entries, the agent cannot solve new episodes memorizing what is the right pair of item to take and monster to fight, but it has to learn to read the goal and the manual and cross-reference them with the environment observation. the agent's performance is tested on episodes generated in such a way that no assignments of monster-team-modifier-element are ever seen during training, to test whether the agent is able to generalize via reading to new environments with unseen dynamics. for more in depth description of the environment and how it is generated the reader can refer to @xcite . we start considering two variants of the original rtfm, the dynamical version with simple language and the dynamical version with natural language, denoted respectively dyna and dyna+nl in @xcite . first, we notice that the sota agent txt2π @xcite cannot solve even the simpler dyna from scratch; instead it has to resort to curriculum learning. second we notice that dyna and dyna+nl differ in two aspects, both concerning the way in which the manual and the goal are expressed. the first difference is that dyna uses fixed language templates like "gleaming beats fire" instead of one of multiple crowdsourced natural language reformulations, like "fire monsters are weak against gleaming items"; we denote them simple language (sl) and natural language (nl) respectively. the second difference is that in dyna the sentences of the manual are always ordered in a specific way (e.g. the first sentence always refers to monsters of the cold element and the last sentence to the star alliance team), whereas in the dyna+nl task the order of the sentences is always shuffled; we call these factors no_shuffle and shuffle respectively. taking all of this into consideration, we propose to systematically analyse the rl methods' performances under the following factors of complexity: number of monsters (one or two), kind of language (sl or nl) and sentence ordering (no_shuffle or shuffle). considering all possible combinations, we end up with eight rtfm variants; in particular we introduce the single monster variants to provide more gradual degrees of complexity to the rtfm challenge. in the rest of the paper we will use the following notation: the original variant with two monsters, simple language and no shuffling, dyna, is named (two, sl); similarly the original variant with two monsters, natural language and shuffling, dyna+nl, is named (two, nl, shuffle). a consistent notation is adopted for the remaining six variants, which do not have a corresponding counterpart in the original rtfm. ## language-conditioned world model the goal of any rl agent is to find a policy π(a|s) that maximizes the expected cumulative reward e π [ t @xmath0. in this work, we take the model-based approach to rl: we learn a language-conditioned stochastic model of the environment and use it to plan with a stochastic version of vanilla mcts. we name our method reader (for reinforcement learning agent for discrete environments with written instructions). it is composed of a world model and a mcts planning module. the world model consists of three components (see fig. 3 ). figure 3 : components of language-conditioned world model: in blue and grey the encoder and decoder of the representation model, in green the transition model and in yellow the reward model. during training the three models are trained separately (in terms of computational graphs). at planning time, we use the transition model, followed by the decoder and by the reward model to predict respectively ẑt+1 , ŝt+1 , rt+1 . note that in contrast to existing model-based agents, our world model is conditioned on the episode-specific textual descriptions: the manual m which describes the roles of the monsters and the rules of the episode and the sentence g describing the agent's goal. ## mcts planning with the learned world model to use the model for planning, we need to predict the next state ŝt+1 and reward rt+1 given the information available at the current step, that is s t , a t , m and g. we do that by first using the transition model to sample the next discrete latent code ẑt+1 , then using the representation model decoder to get ŝt+1 and finally the reward model to get rt+1 . to extend mcts to stochastic transitions, we use two types of nodes in the tree: state nodes and state-action nodes. the tree branches at state nodes by considering different actions that can be taken in the state and it branches at state-action nodes by looking at different stochastic outcomes. while we have control over the actions that we choose, the outcomes and corresponding rewards are always sampled by using the learned model. we index the possible outcomes from a state-action node based on the indices of the pair of discrete codes (z a , z b ) sampled by the transition model; if the codes have been already sampled, we continue traversing the already-expanded tree, otherwise we expand the newly-sampled state node. ## experiments training to ease the computational demands of the full rl pipeline, we consider an offline rl setup. we first collect a dataset of trajectories for each task, then we train the model on it without ever directly interacting with the real environment and finally we evaluate the mcts agent equipped with the learned model by playing 1000 episodes in the test environments. to collect the datasets we use a random policy for 50% of the episodes and a vanilla mcts policy 3 with access to a simulator of the real environment for the remaining 50% of the trajectories. the most critical for performance is obtaining sufficiently many terminal transitions where the agent interacts with items and monsters, as these require complex reasoning for reward prediction. a random policy would be inefficient and the vanilla mcts policy would not sample enough failure cases, hence we selected 50/50. we later 3 we select the number of simulations per action of the mcts policy in such a way that its performance is sub optimal. ablated reader's performance with splits 25/75 and 75/25, finding that the performance is robust to the chosen split and only slightly decreases in the case where 75% of the trajectories come from the vanilla mcts policy. we report and discuss the full results of the ablation in sec. d of the appendix. each dataset is composed of 200k episodes, or roughly 1m frames; full details are reported in ## discussion and conclusions rl environments such as rtfm are great for developing models capable of natural language understanding. prior work shows that the model-free approach might work well in those benchmarks, but it requires a large number of samples and it is not insightful with regard to the agent's understanding capabilities. in this work we show that the model-based approach can improve both the performance and the sample-efficiency of rl agents. while the sample-efficiency gains were expected, we speculate that the performance gains in rtfm are mostly driven from improvements in estimating the complex language-dependent reward function; we think that disentangling the reward prediction from the environment stochasticity by conditioning the reward prediction on the next predicted frame is a key element in this regard. moreover we provide concrete examples of how the agent's plans can be inspected and we consider this a step in the direction of more interpretable rl agents. our experiments highlight that rtfm remains a challenging benchmark and that existing methods need a large number of samples to learn the rtfm reward. therefore, more research is needed to improve the sample efficiency of models for grounded language learning; leveraging the out-of-the-box representation learning capabilities of pre-trained large language models seems a promising direction to explore. another very interesting future direction is to extend reader to more realistic visual representations; using pre-trained patch embeddings from vision foundation models like @xcite appears to be a very promising avenue.