Pretraining for Language Conditioned Imitation with Transformers

1 INTRODUCTION . Whereas most contemporary reinforcement learning ( RL ) methods require programming a reward function , we are interested in training RL agents which understand language , enabling a flexible and interpretable interface and facilitating human-AI interaction . Natural language often carries tremendous amounts of information and can thus both provide strong learning signals and describe complex scenarios . Furthermore , agents which are able to learn from a diverse set of tasks – such as the space of tasks specified by language – can take advantage of broader data for supervision , access more diversity in an environment , and perform more effective few-shot learning . One path towards more general agents is data-driven reinforcement learning , which promises to leverage diverse data for improved generalization in RL ( Levine et al. , 2020 ) , similar to approaches in natural language . Recent approaches to data-driven RL have proposed leveraging powerful sequence models ( Chen et al. , 2021 ; Janner et al. , 2021 ) that are now mainstream in language ( Devlin et al. , 2018 ; Radford et al. , 2019 ) . While initial approaches to training language models on large-scale unsupervised data focused on learning individual word embeddings ( Pennington et al. , 2014 ) , newer models , especially those based on the transformer architecture ( Vaswani et al. , 2017 ) , can also transfer sequential aspects learned during pretraining . Similar work in RL has found great benefit to transferring learned embeddings via a vision encoder ( Yarats et al. , 2021 ) , but transferring sequential behavior knowledge is relatively unstudied . However , using transformer sequence modeling approaches , we can develop methods which can finetune more effectively from offline data . To better study this , we propose a new benchmark and release a new dataset1 , Text-Conditioned Frostbite , in which the agent must perform tasks specified by language instructions in the Atari Frostbite environment ( Bellemare et al. , 2013 ) . Crucially , our dataset consists of a diverse set of tasks , ranging from easier tasks such as “ try to stay on the right side ” to harder tasks such as “ spend as much time as possible on the fourth ice floe ” . Inspired by the proposal from Lake et al . ( 2017 ) , we believe Frostbite is a diverse and dynamic environment requiring temporally extended strategies , whereas traditional evaluations underutilize the environment by finding agents with single , narrow strategies . We focus on Atari Frostbite as it allows for combining this naturalistic language with complex goals and behaviors using a familiar , well-studied environment . We introduce a transformer operating on text-state-action sequences – which we call Text Decision Transformer ( TDT ) – adapting the single-stream Decision Transformer ( Chen et al. , 2021 ) to our setting of conditioning on language goals by replacing the reward tokens with text tokens . We perform empirical evaluations on our environment using the TDT , finding it serves as a reasonable baseline architecture for study , and evaluate performance with unsupervised pretraining . Our contributions are as follows : 1 . We propose a new multimodal benchmark , Text-Conditioned Frostbite , and provide a dataset of 5M labelled transitions with 500 different tasks . 2 . We evaluate a single-stream architecture , Text Decision Transformer , and find it outperforms MLP baselines , highlighting sequence modeling-based approaches . 3 . We investigate unsupervised pretraining , both on language and on trajectories , and find trajectorybased pretraining yields improved performance in low-data settings . Furthermore , we investigate scaling laws , finding the task is complex enough that TDT improves with more data . 2 PRELIMINARIES . 2.1 LANGUAGE-CONDITIONED IMITATION LEARNING . In the problem of Language-Conditioned Imitation Learning , we are provided with a tuple of ( S , A , M , P , R ) i , which correspond to the states , actions , tasks , transition dynamics , and a reward or score function . A trajectory τ = ( s1 , a1 , ... st , at ) is a sequence of states and actions generated by the transition dynamics and the agent ’ s action at every timestep , where t is less than or equal to the maximum episode length . The score function R is prompted with a trajectory τ and a text 1Dataset and code are available at : < github will be made public at publication > , temporary link to dataset and code for reproduction are in supplementary material task m and outputs the degree of success that the trajectory had with respect to the task – note we only use R for evaluation purposes . A policy/model πθ ( at|s≤t , a < t , m ) is tasked with outputting actions which maximize the score function , i.e . Ea∼πθ [ R ( τ , m ) ] . The model has access to a dataset of expert trajectory-language pairs ( τ , m ) for training . 2.2 TRANSFORMERS . Transformers were originally proposed by Vaswani et al . ( 2017 ) and have since become the frontrunner in solving sequence modelling problems . The models are built with sequential self-attention modules . Each self-attention module is given a sequence of n embeddings and outputs another n embeddings with the same dimensions . The self-attention modules work by performing the inner product between the current embedding xi with the query matrixQ to give qi and with the key matrix K to give the key ki . The model takes the inner product between the key and query for j ∈ [ 1 , n ] , yielding n logits , over which the model takes the softmax and yields a probability distribution . The final result is then a convex combination of the value vectors vi , with the weights dictated by the probability distribution . More concisely , the i-th component of the output zi is given by zi = n∑ j=1 softmax ( { 〈qi , kj′〉 } nj′=1 ) j · vj . ( 1 ) 2.3 DECISION TRANSFORMER . Our work builds on Decision Transformer ( DT ) ( Chen et al. , 2021 ) , which uses a GPT-2 ( Radford et al. , 2019 ) transformer architecture , leveraging a causal mask where the i-th prediction depends only on input tokens up to i. DT models a trajectory sequence ( R̂1 , s1 , a1 , . . . , R̂t , st , at ) autoregressively , where R̂t is the returns-to-go from timestep t to the end of the episode . Test-time evaluation proceeds by initializing a behavior specification – DT uses a desired return R̂1 – and alternating generation between states from the environment and actions sampled from the model . To reduce computation , DT only uses the last K timesteps as input to the model , called the context size . A context size of 1 represents a Markovian policy . 3 ENVIRONMENT AND DATASET . The first contribution of our work is in developing a diverse , yet accessible environment for evaluating the performance of text-conditioned imitation models . We build on the proposal from Lake et al . ( 2017 ) , in which the authors propose the “ Frostbite Challenge ” . The authors propose Frostbite is challenging due to the requirement of temporally extended strategies and the diverse landscape of possible tasks to complete . To construct a dataset for this task , we manually collected several hundred tasks and trajectories . These tasks range from “ don ’ t move ” to “ stay as long as you can on the fourth ice floe ” to “ die to a crab on your first life ” and so on . 3.1 ENVIRONMENT DETAILS . Evaluation via RAM state . The first challenge with such an environment is building deterministic evaluation strategies . For this reason , evaluation of the agent is performed by accessing the RAM state of the game , calculating the level , position , score , and other related quantities of the agent , and comparing them with the desired outcomes . This evaluation loop runs after a trajectory is completed and can be prompted with a number of different desired tasks . The interface is malleable enough to even generate rewards for specific tasks , though we have not yet conducted such a study . We use the underlying state-space of the game to calculate specific metrics in the environment . These metrics are then normalized to be in the range ( 0 , 1 ) , so no task gets unequal weighting . For example , in a task like “ spend as much time as possible on the first ice floe ” , the assigned success is simply ( number of timesteps on 1st ice floe ) / ( number of timesteps ) . Performance measurement . Performance of the model is measured by prompting the model with several different tasks and subsequently running the evaluation procedure on the generated trajectories . The score is normalized to be in the range [ 0 , 1 ] and is designed to be a success metric . If the model receives a score of 1 on a generated trajectory , then the model has succeeded in obeying the task , while a score of 0 indicates a failure . 3.2 DATASET DETAILS . Overview . The dataset consists of more than 700 labelled missions , for a total of 5m timesteps , each consisting of the current task , the current state , and the action taken . The Atari environment supports 18 discrete actions corresponding to directional movement and interaction . Game frames are resized and converted to black-white . The environment does not explicitly perform frame-stacking , as our agent has access to context , but frame-stacking can be implemented with the interface . We list some example tasks in Table 1 and include more in the Appendix . Collection . We collected several hundred tasks and trajectories by hand . These trajectories were generated by taking advantage of the Atari-Gym user interface and playing Frostbite with keyboard controls . At the beginning of every task , users are prompted with a task that they will try to achieve , and their states , actions , and tasks are recorded and saved . Task categorizations . For evaluation , a subset of the collected tasks are supported and separated into distinct task suites depending on the degree of control necessary to perform the task . For instance , in the easy task suite , there are tasks such as “ dont move , ” “ die on the first level , ” and “ stay on the left side. ” Such tasks are quite simple and do not require advanced visual perception or control of the agent . In the medium task suite , tasks require higher levels of perception , but are less demanding than the hard suite . The medium task suite contains “ get as close to 1000 score as you can , ” “ reach level 2 , ” and “ jump between the second and third ice floes. ” The hard suite contains tasks that require the most advanced control , such as “ reach level 5 ” and “ spend as much time as possible on the fourth ice floe. ” The task suites include only tasks that support RAM score evaluation . A list of further training tasks can be found in the Appendix ( see Table 3 ) . 4 METHOD . Here , we detail the architecture and modifications of the Text-Conditioned Decision Transformer . This model uses a causal transformer architecture in the style of GPT-2 ( Radford et al. , 2019 ) to model sequences jointly consisting of text and trajectories .

Summary: This work proposes a new instruction-following reinforcement learning environment based on the Atari Frostbite environment. They experiment with a modified decision transformer on the environment. The results largely show that performance improves with longer context, more pretraining, and more data.