Language Model Pre-training Improves Generalization in Policy Learning

1 INTRODUCTION . In recent years , language models ( LMs ) trained on open-domain text corpora have come to play a central role in machine learning approaches to natural language processing tasks ( Devlin et al. , 2018 ) . This includes tasks that are not purely linguistic , and additionally require nontrivial planning and reasoning capabilities : examples include as vision-language navigation ( Majumdar et al. , 2020 ; Fried et al. , 2018 ; Suglia et al. , 2021 ) , instruction following ( Zhang & Chai , 2021 ; Hill et al. , 2020 ) , and visual question answering ( Tsimpoukelli et al. , 2021 ) . Indeed , some of these tasks are so remotely connected to language modeling that it is natural to ask whether the capabilities that result from LM pre-training might extend to tasks that involve no language at all—and if so , how these capabilities might be accessed in a model trained only to process and generate natural language strings . In this paper , we study these questions through the lens of embodied decision-making , investigating the effectiveness of LM pretraining as a scaffold for learning control policies for interactive tasks featuring partial observability , large action spaces , complex states , and complex dynamics . We describe a series of experiments in the VirtualHome environment ( Puig et al. , 2018 ; 2020 ) in which LMs are used to initialize policies , and show that LM pre-training substantially improves generalization across common-place tasks in household environments . In Experiment 1 ( Section 6 ) , we encode the inputs to a policy—including observations , goals , and action histories—as templated English phrases ( e.g . representing the goal on ( fork , table ) as There is a fork on the table . ) as shown in Figure 1 . A pretrained LM is then fined-tuned to produce representations of these phrases that can be used to predict subsequent actions . For i.i.d . training and evaluation tasks , we find that this approach completes tasks at a rate comparable to the same transformer-based policy trained from scratch . For generalization to out-of-distribution tasks , however , LM pretraining confers substantial benefits : it improves task completion rates by nearly 20 % for tasks involving novel initial environment states and goals ( Figure 1 “ Transformer from scratch ” and “ Pretrained transformer ( natural encoding ) ” ) . Next , we conduct two experiments aimed at clarifying the role of this string-based encoding . We design Experiment 2A that uses random strings instead of natural language inputs and Experiment 2B that uses non string-based encodings to study different ways of building interfaces between input encodings and LMs . In Experiment 2A ( Section 7 ) , we replace the “ natural ” string encodings of Experiment 1 with an arbitrary mapping between logical goals and tokens ( e.g . serializing on ( fork , table ) as brought wise character trees fine order yet ) . This random encoding substantially ( by roughly 12 % ) degrades performance on out-of-distribution tasks , indicating that LM encoders are sensitive to the form of string encodings even when fine-tuned ( Figure 1 “ Pretrained transformer ( natural encoding ) ” and “ Pretrained transformer ( random encoding ) ” ) . In Experiment 2B , we investigate whether string-based encodings are necessary at all . We replace the model ’ s word embedding layer with a randomly initialized embedding layer that maps from discretized environment observations , goal , and history actions to a sequence of LM input vectors , and fine-tune this embedding layer jointly with the LM itself . This learned encoder performs almost the same as the encoding of Experiment 1 , indicating that effective encodings for non-linguistic tasks can be learned from scratch ( Figure 1 “ Pretrained transformer ( natural encoding ) ” and “ Pretrained transformer ( learned encoding ) ” ) . These experiments offer two main conclusions . First , they show that language modeling improves generalization in policy learning : initializing a policy with a neural LM ( pre-trained on a nextword prediction task with a large text corpus ) substantially improves out-of-distribution performance on ( non-linguistic ) tasks in an interactive environment . Second , they show that language-based environment encodings are not needed to benefit from LM pretraining : it is instead possible to learn an interface between observations , actions , and model-internal representations derived from text corpora . These results point the possible effectiveness of language modeling as a generalpurpose pre-training scheme to promote structured generalization in broader machine learning applications . 2 RELATED WORK . In recent years , word and sentence representations from pre-trained LMs ( Peters et al. , 2018 ; Devlin et al. , 2018 ; Radford et al. , 2018 ) have become ubiquitious in natural language processing , playing a key role in state-of-the-art models for tasks ranging from machine translation ( Zhu et al. , 2020 ) to task-oriented dialog ( Platanios et al. , 2021 ) . Some of the most successful applications of pretraining lie at the boundary of natural language processing and other domains , as in instruction following ( Hill et al. , 2020 ) and language-guided image retrieval ( Lu et al. , 2019 ) . Building on this past work , our experiments in this paper aim to explain whether these successes result entirely from improved processing of text , or instead from domain-general representational abilities . Below , we briefly survey existing applications of pretraining that motivate the current study . Learning representations of language From nearly the earliest days of the field , natural language processing researchers have observed that representations of words derived from distributional statis- tics in large text corpora serve as useful features for downstream tasks ( Deerwester et al. , 1990 ; Dumais , 2004 ) . The earliest versions of these representation learning schemes focused on isolated word forms ( Mikolov et al. , 2013 ; Pennington et al. , 2014 ) . However , recent years have seen a number of techniques for training ( masked or autoregressive ) language models to produce contextualized word representations ( which incorporate information neighboring words in sentences and paragraphs ) via a variety of masked-word prediction objectives ( Devlin et al. , 2018 ; Yang et al. , 2019 ) . Applications In addition to producing useful representations , these language models can be finetuned to perform language processing tasks other than language modeling by casting those tasks as word-prediction problems . Successful uses of representations from pretrained models include syntactic parsing ( Kitaev et al. , 2018 ) and language-to-code translation ( Wang et al. , 2019 ) ; successful adaptations of LM prediction heads include machine translation ( Zhu et al. , 2020 ) , sentiment classification ( Brown et al. , 2020 ) and style transfer ( Keskar et al. , 2019 ) . Text-based games ( Yao et al. , 2021 ; Yuan et al. , 2018 ; Ammanabrolu & Riedl , 2018 ; Côté et al. , 2018 ) inherently involve text as both the input and the output . Recent works ( Yao et al. , 2020 ) in text-based games use GPT-2 to solve the text-based games and get significant performance improvements . However , it is hard to describe 3D information using text and most of their experiments are in 2D environments . Included in these successes are a number of tasks that integrate language and other modalities , including visual question answering and image captioning ( Yang et al. , 2020 ) . In models that condition on both text and image data , several previous approaches have found that image representations can be injected directly into language models ’ embedding layers ( Tsimpoukelli et al. , 2021 ) using a similar mechanism to the one we describe in Experiment 2B . One of our main contributions is to show that approach works even for tasks in which only non-linguistic information is relevant to model predictions . What do LMs encode ? The possibility that LMs might encode non-linguistic information useful for other downstream tasks is suggested by a number of recent “ probing ” studies aimed at understanding their predictions and the structure of their internal representations . Pre-trained LMs can answer a non-trivial fraction of queries about both factual and common-sense knowledge ( Roberts et al. , 2020 ) . Their representations encode information about perceptual relations among concepts , including visual similarity among object classes ( Ilharco et al. , 2020 ) and the structure of color spaces ( Abdou et al. , 2021 ) . Finally , they appear to be capable of basic simulation , modeling changes in entity states and relations described by text ( Li et al. , 2021 ) . LM pretraining beyond language Two recent papers consider questions closely related to the ones investigated here : ( Brown et al. , 2020 ) show that the GPT-3 model is capable of performing a limited set of arithmetic and string manipulation tasks ; ( Lu et al. , 2021 ) show that pretrained LMs require very little fine-tuning to match the performance of task-specific models on several image classification and numerical sequence processing tasks . In this paper , we focus on non-linguistic tasks where the inputs and outputs do not involve language . To the best of our knowledge , the current study is the first to demonstrate improved generalization in a non-linguistic problem over a standard neural-network baseline using a pre-trained language model . 3 LANGUAGE MODELING AND POLICY LEARNING . We begin with a brief review of the ingredients of language modeling and policy learning tasks used in our experiments . 3.1 LANGUAGE MODELING . Our experiments in this paper focus on autoregressive , transformer-based language models ( Vaswani et al. , 2017 ) . These models are trained to fit a distribution pθ ( y ) over a text sequence y by decomposing it into a sequence of tokens y = { y1 , y2 , . . . , yn } via the chain rule : log pθ ( y ) = n∑ i=1 log pθ ( yi | y1 , y2 , . . . , yi−1 ) . ( 1 ) Each conditional distribution pθ ( yi|y1 , y2 , . . . , yi−1 ) is parameterized by a transformer neural network fθ ( y1 , y2 , . . . , yi−1 ) . This network encodes each conditioned token yi into a continuous embedding ei = g ( yi ) which is then fed into the transformer architecture and encoded into a categorical distribution over token values of yi . Our experiments utilize a standard language model , GPT-2 , that is trained on Webtext dataset ( Radford et al. , 2018 ) using Huggingface library ( Wolf et al. , 2019 ) . 3.2 POMDPS AND POLICY LEARNING . Our experiments explore the application of LMs to general sequential decision-making tasks in partially observed environments . These tasks may be formalized as partially observable Markov decision processes ( POMDPs ) . A POMDP is defined by a set of states S , a set of observations O , a set of actions A , and a transition model T ( st+1|st , at ) that predicts the next state st+1 based on the current state st and an action at . Importantly , in a POMDP setting , the observation ot only captures a portion of the underlying state st , and an optimal decision-making strategy ( a policy ) must incorporate both the current observation and the previous history of observations and actions . For experiments in this paper , policies are parametric models πψ ( at|g , h , ot ) that select actions given the goals g , history information h , and partial observations ot of the current state st. All our experiments use imitation learning ( Santara et al. , 2017 ; Ng et al. , 2000 ; Peng et al. , 2018 ) , specifically behavior cloning ( Pomerleau , 1991 ; 1989 ; Torabi et al. , 2018 ) , to train πψ . We collect a dataset of N̂ expert training trajectories D = { d1 , · · · , dN̂ } , where each individual trajectory consists of a set of goal , observations , and actions , i.e . di = { o1 , a1 , · · · , aT , g } , where T is the length of an expert trajectory . We then train a policy πψ ( at|g , ht , ot ) which maximizes the likelihood pψ ( a ) of the expert actions a = { a1 , · · · , aT } taken in a trajectory using supervised learning , log pψ ( a ) = T∑ t=1 log pψ ( at | g , ht , ot ) , ( 2 ) where ht consists of all history in the environment up to timestep t .

Summary: This paper takes a transformer-based language model, pre-trained on a large text corpus (In this case GPT-2) and uses it for the symbolic version of the VirtualHome environment. The observation, goals and action history of the agent are encoded as text strings in a few various ways and fed as input to the transformer and the output of the model is pooled to predict the agent action. The paper demonstrates that in cases where the test distribution differs in some way from the training distribution, that using pre-trained transformers greatly improves performance on the task