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 investigates the effectiveness of the language model for training the policy in embodied environments. The authors use a pre-trained GPT-2 to initialize the policy, then show the generalization effect in policy learning. In the experiments, the authors demonstrated the language model shows a better generalization effect with simple baseline and ablation studies.