Pretraining for Language Conditioned Imitation with Transformers
Authors: Aaron L Putterman
Kevin Lu
Igor Mordatch
Pieter Abbeel
Venue: ICLR 2022 Submitted (2022)

Abstract: We study reinforcement learning (RL) agents which can utilize language inputs. To investigate this, we propose a new multimodal benchmark -- Text-Conditioned Frostbite -- in which an agent must complete tasks specified by text instructions in the Atari Frostbite environment. We curate and release a dataset of 5M text-labelled transitions for training and to encourage further research in this direction. On this benchmark, we evaluate Text Decision Transformer (TDT), a transformer directly operating on text, state, and action tokens, and find it improves upon other baseline architectures. Furthermore, we evaluate the effect of pretraining, finding unsupervised pretraining can yield improved results in low-data settings.

Under review as a conference paper at ICLR 2022
PRETRAINING FOR LANGUAGE IMITATION WITH TRANSFORMERS
Anonymous authors Paper under double-blind review
ABSTRACT
We study reinforcement learning (RL) agents which can utilize language inputs and efficiently learn on downstream tasks. To investigate this, we propose a new multimodal benchmark Text-Condiee Coxt-Conditi Frostbite in which agent must complete tasks specified by text instructions in the Atari Frostbite environment. We curate and release dataset of 5M text-labelled transitions for training, and to encourage further research in this direction. On this benchmark, we evaluate Text Decision Transformer (TDT), transformer directly operating text, state, and action tokens, and find improves upon baseline architectures. Furthermore, we evaluate the effect of pretraining, finding unsupervised pretraining can yield improved results in low-data settings.
a
t-1
linear decoder
causal transformer
emb. pos. enc.
m
m
i+1
S
t-1
a
t-1
Decision Transformer (TDT) architecture for specifying behaviors via language inputs. Text tokens m, are prepended to the sequence of episode tokens, specifying the current task. Alternating states 8; and actions a provide context for the model and are rolled out in the environment. Actions are predicted from every state in an autoregressive fashion.
Figure 1: The
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 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 data-driven reinforcement learning, which promises 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., (017)) can also transfer sequential as- pects learned during pretraining. Similar work in RI. has found great benefit to transferring learned

---

Under review as a conference paper at ICLR 2022
100
110
110
32"
3
31
3
31 3
Task: "jump between the second and third ice floes"
Figure 2: We create new dataset on the Atari Frostbite environment, consisting of 5M timesteps of labelled text-trajectory pairs. The agent must solve diverse set of tasks specified by natural language instructions. Above are example frames representative of task from the dataset.
embeddings via a vision encoder Yarats et al. 2021). but transferring sequential behavior knowl- edge is relatively unstudied. However, using transformer sequence modeling approaches, we develop methods which can finetune more effectively from offline data.
To better study this, we propose Frostbite, in which the agent must perform tasks specified by language instructions in the Atari Frostbite environment (Bellemare 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 (2017),
benchmark and release
datase Text-Conditioned
much time as possible the fourth ice floe". Inspired by the proposal from Lake we believe Frostbite is 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 it allows for combining this naturalistic language with complex goals and behaviors using a familiar, well-studied environment.
We introduce 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 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 multimodal benchmark, Text-Conditioned Frostbite, and provide dataset of 5M labelled transitions with 500 different tasks.
We evaluate a single-stream architecture, Text Decision Transformer, and find outperforms MLPP baselines, highlighting sequence modeling-basee approaches.
3. We investigate unsupervised pretraining, both on language and on trajectories, and find trajectory- based 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.
PRELIMINARIES
2.1
IMITATION LEARNING
In the problem of Language-Conditionel Imitation Learning, we are provided with a tuple of (S, A. M, P. R),, which correspond the states, actions, tasks, transition dynamics, and reward or score function. A trajectory = ($1, ...t. 0e) is a sequence of states and actions generated by the transition dynamics and the agent's action every timestep, where is less than or equal to the maximum episode length. The score function R is prompted with trajectory and a text
Dataset and code available at: < thub wi be made publ ic at publi publication cationn tem- porary link dataset and code for reproduction in supplementary material
2

---

Under review as a conference paper at ICLR 2022
task m and outputs the degree of success that the trajectory had with respect the task note only use R for evaluation purposes. A policy/model IT [ar]sss, aci, m) is tasked with outputting actions which maximize the score function, i.c. Ba-5y|R(T m)). The model has access to a dataset of expert trajectory-languaae pairs (T,m) m) for training.
2.2 TRANSFORMERS
Transformers were originally proposed by Vaswani al. (017 and have since become the front- runner 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 v; with the query matrix Q to give d and with the key matrix K to give the key k; The model takes the inner product between the key and query for j E [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 u, with the weights dictated by the probability distribution. More concisely, the i-th component of the output 2, is given by
2,
so softmax(((ee k()i, Eg.
(1)
2.3 DECISION TRANSFORMER
Our work builds on Decision Transformer (DT) (Chen et al. 2021), which uses GPT-2 (Radford al. 2019) transformer architecture, leveraging causal mask where the i-th prediction depends only on input tokens up to DT models a trajectory sequence (R1, $1 o1..... .. Rr St, at) autoregres- sively, where R₂ is the returns-to-go from timestep to the end of the episode. Test-time evaluation proceeds by initializing a behavior specification DT uses a desired return R₁ 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 represents Markovian policy.
ENVIRONMENT AND DATASET
The first contribution of our work is in developing diverse, yet accessible environment for eval- uating 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 complete. To construct dataset for this task, manually collected several hundred tasks and trajectories. These tasks range from "don't 't move" to "stay as long as you can on the fourth ice floe" to "die a crab 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 number of different desired tasks. The interface is malleable enough to even generate rewards for specific tasks, though we have not yet conducted such study. We use the underlying state-space of the game to calculate specific metrics in the environment. These metrics are then normalized to in the range (0, 1), so no task gets uncqual weighting. For example, in task like "spend much time possible 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 trajecto- ries. The score is normalized be in the range [0, 1] and is designed to be a success metric. If the
3

---

Under review as a conference paper at ICLR 2022
Split
Example Tasks
Easy
stay the left side, die the first level, get close 500 you jump between the second and third icc floes, reach level dic on the second level reach level 5, spend as much time as possible on the fourth ice floc
Medium
Hard
Table I: Example tasks in our Text-Condition Frostbite dataset. Tasks are split into easy, medium, and hard categories for evaluation purposes.
model receives a score of I on generated trajectory, then the model has succeeded in obeying the task, while score of O indicates 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, agent has access to context, but frame-stacking can be implemented with the interface. We list some example tasks in Table 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, 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 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 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
METHOD
Here, 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.
4.1 ARCHITECTURE
Similar to Decision Transformer, we model the trajectories as a sequence of tokens. However, because operating within the framework of imitation learning, for behavior specification, we excise the return-to-go tokens from our architecture and use text tokens instead. These text tokens prepended the sequence of tokens. The original words of the task encoded using an embedding table, whereas states and actions are encoded using CNNs / MLPs, respectively. Thus, typical sequence of tokens passed into the model will resemble (see Figure I]:
T = (mi1; 72,.... Ma, 81, ; 01, 82, 02 .2....... 8t, a+)
The model is prompted with a text tokens and a sequence of states and actions. The inputs are tok- enized and passed through the model. Action logits for action a, correspond to the output from state s,, and due to the causal transformer mask, only tokens preceding 8, are visible in this prediction. For generating trajectories, we sample from the predicted action logits, and step the environment
4

---

Under review as a conference paper at ICLR 2022
with the predicted action. At the end of the trajectory, evaluate performance using the environ- ment's score function to determine if the trajectory and original text are well-aligned.
We preprocess language inputs by converting to lower case and stripping any punctuation. Words are then converted to one-hot representation and passed through a learned embedding table. See Algorithm for more details.
4.2 TRAINING
We operate with a standard supervised training procedure. Random minibatches of paired tasks and trajectories are sampled from our offline dataset. The model performs a forward pass over the joint sequence, outputting the logits of the distribution next actions in parallel. We optimize the cross entropy loss between the predictions and the true actions. Pseudocode for the model and training is included in Algorithm
4.3 UNSUPERVISED PRETRAINING
In Sections 5.2-5.4, additionally evaluate the effect of unsupervised pretraining, which refers training on unlabelled trajectories (without a text pair) drawn from the dataset. For pretraining, the model is only given input states and actions, and otherwise follows the same procedure the pervised training regime. The pretraining data is drawn from the corpus of all missions simply with text missions removed. Later, the model fine-tuned and evaluated on specific subsets of missions to measure the effectiveness of the representations learned by pretrainining. This is comparable to the standard autoregressive language modeling objective used by GPT (Radford et al., 2019), or be- havior cloning in imitation learning. We can transfer the vision encoder layers and/or the sequence model (the transformer backbone).
5 EMPIRICAL EVALUATIONS
5.1 DOES THE TRANSFORMER ARCHITECTURE MODEL THE JOINT SEQUENCES WELL?
We first evaluate the transformer architecture itself in the fully supervised setting. We compare to an MIP baseline which removes the ability to attend between modalities: text tokens are processed by pretrained BERT model (Devlin et al. 2018), and then concatenated to the state and action tokens, which are processed by the same encoders TDT. The MLP with context (K) of 1 is chosen to have the same number of parameters as TDT. The MLP with context 1 is outperformed by all but one TDT architecture choice, outperformed by the TDT architectures with context size of at least 10. This is despite the fact that the MLP with context size 20 uses many more parameters (since it relies on concatenating over timesteps). The TDT with one context represents a fairly similar architecture to the MLP Baseline with the key difference that the TDT model learns its own text embeddings and utilizes attention between different modalities. We see that the standard error bars are larger for the MLP. and we belive this the reason for its relatively strong performance on the hard tasks.
those with limited context. The MLP with context 20 is
Our results are shown in Figure The conclusion from this experiment is two-fold:
1. The causal transformer architecture accessing text and state tokens in a single stream outperforms concatenation-basee approach with pre-encoded text tokens.
Increasing the context of the causal transformer leads stronger performance, indicating that longer sequence modeling can lead to more powerful model.
5.2 HOw EFFECTIVE IS UNSUPERVISED PRETRAINING FOR LEARNING REPRESENTATIONS?
Perhaps the most enticing prospect of TDT is its ability to harness large-scale unlabelled data for pretraining, enabling few-shot behavior observed in language. For our unsupervised pretraining experiments, models initially pretrained via behavior cloning tasks. These training tasks are sampled from the same corpus as the labelled tasks, but do not include text. For each evaluation task, sample small subset of the labelled data, including the evaluation
fixed number of unlabelled
5

---

Under review as a conference paper at ICLR 2022
Architecture Comparison in Supervised Setting
1.0
0.8
MLP, K
MLP, K 20
Score
0.6
TDT, K
0.4
TDT, K 10
TDT, K 20
0.2
0.0
Easy tasks
Medium tasks
Hard tasks
Figure 3: Performance of transformer architecture (TDT), based on the context size K. We com- pare to concatenation-basee MLP architecture which receives encoded text latents from pretrained BERT and concatenates them with the latents of the current state from a learned vision encoder. Text Decision Transformer outperforms the ML.P baseline in every suite tasks. All experiments use seeds and show standard deviation between seeds.
Transferring Vision Encoder
From scratch
Pretrsined
Pretrained on 50
Pretrained
Pretrained on 200
Medium tasks
Hard tasks
Figure 4: Transferring vision encoder after unsupervised pretraining. After pretraining with behavior cloning on unlabelled trajectories, significant benefits finetuning for medium and hard tasks compared to a randomly initialized vision encoder. Note the context size used is 20.
task itself. For instance, when evaluating the "jump back and forth between the second and third ice floes" task, the model is trained on this task, as well as some related jumping tasks between the first and second, first and third, etc., ice floes. The data that the model is trained on for each evaluation task is the same across all variants. We finetune all the model parameters.
We first evaluate transferring the vision encoder, which represents the benefit from better image representations. This is the currently the most common method of finetuning RL policies after unsupervised exploration (Yarats et al. 202T). taking advantage larger datasets for stronger rep- resentations, but not transferring learned behaviors. Our results are shown in Figure [
We conclude that by pre-training the entire model on an unlabeled dataset of transitions, we learn important features in our visual encoder (despite no text labels). Using this visual encoder in later tasks then allows better performance when fine-tuning on specific tasks. This improvement scales with larger quantities of pre-training data.
5.3 DOES SEQUENTIAL CONTEXT IMPROVE PRETRAINING PERFORMANCE?
Additionally, due to the sequential nature of the transformer, we can further investigate how context length affects pretraining. Since TDT jointly models modalities, there are multiple possible modes of supervision. On the one hand, the vision encoder may learn to pick up on valuable features simply due to repeated actions in specific positions. For instance, whenever the player is located on the bottom ice floe, the agent may tend to jump higher. In such a case, there would be a learning signal from the environment which would allow the encoder to pick-up features indicating when the agent in the bottom ice floe. On another hand, the vision representations will be processed sequentially by transformer, and may learn features which correlate behavior across timesteps. For

---

Under review as a conference paper at ICLR 2022
instance, if the previous frames indicate agent jumping between two ice floes, knowledge of this past would allow the model to infer the task at hand.
To evaluate this, ablate TDT from Figure 4 by using context size of instead of 20. Again, we transfer the vision encoder only. Our results are shown Table 2 and Figure [ We observe that larger contexts do indeed lead to larger percent increases from pre-training. This suggests that the improvement from pre-training is not just due to seeing repeated state-action pairs, and benefits from trajectory-level inference.
Split
Context = Context 20
Hard 43.5%
181.1 %
Medium
39.0%
39.0%
Table 2: Comparing benefit of transferring vision encoder for different model context sizes. Im- provement is measured as percent increase in performance when pre-training with 200 unlabelled tasks compared an untrained model. The improved transfer of the 20 context model indicates that sequential trajectory-level information can be transferred.
Transterring Vision Encoder (Context 1]
From scratch
Pretrained 10
Pretrained 50
Pretrained 100
Pretrained
Medium tasks
Hard tasks
Figure 5: Transferring vision encoder with shorter context (size 1). While pretraining still helps, is generally less effective than with the model utilizing longer context (Figure 4) the numbers for pretraining on 200 unlabelled tasks are summarized in Table 2
5.4 CAN WE TRANSFER THE SEQUENCE MODEL BACKBONE?
In the past sections, we investigated transferring the vision encoder only from the pretraining phase. This has been standard for finetuning RL agents after unsupervised exploration because it can be difficult to transfer the policy/value parameters between the distinct exploration/fiectuing regimes. However, as we have seen in language, we can gain greater expressiveness finetuning by trans- ferring the sequential aspect of the model (Devlin et al., 2018 Radford et al. 2019). We repeat the same unsupervised pretraining phase, but now transfer the transformer backbone, which captures learned behaviors and sequential integration of the inputs.
Our results are shown in Figure 6 The sequence model incorporates additional information that is applicable to the downstream tasks: for instance, after pretraining 200 unlabelled tasks and then evaluating on the hard tasks, additionally transferring the sequence model doubles the score only transferring the vision encoder. This hints that transferring learned sequence models can serve powerful foundation for few-shot learning in RL settings.
5.5 How DOES TEXT DECISION TRANSFORMER SCALE WITH DATA?
Finally, we also show additionally scaling results in the supervised setting. In the previous experi- ments, we showed pretraining on larger amounts unsupervised data improved performance. Here, we additionally train the model on larger amounts of labelled data ranging from 10 to 700, and then evaluate the performance of these models on heldout validation set. The validation set consists of new trajectories and some new task statements not in the training data. The validation loss of the model shown in Figure 7 A low validation loss indicates that the model more likely to predict the correct action, and hence more likely to achieve higher success rates in the environment.

---

Under review as a conference paper at ICLR 2022
Transferring Model
scratch
Pretrained on 50 Pretrained Pretrained on 200
0.4
0.2
0.0
Hard
Figure 6: Transferring both the sequence model and the vision encoder after unsupervised pretrain- ing. The sequence model captures additional information that improves performance compared to only transferring the vision encoder.
Scaling Trend for Supervised Setting
3.6x 10-1
Validation
333 10-1
31x 10
101
10?
Number of training examples
Figure 7: Validation loss scaling the supervised setting based on number labelled training trajectories. The relationship between loss and data is approximately linear log-log plot.
5.6 QUALITATIVE EXAMPLEES
See Figure X for an example rollout with the task "jump between the second and third ice floes". The agent can be seen jumping between the two aforementioned ice floes.
RELATED WORK
Language-conditione IL/RL. BabyAl (Chevalier-Boisvertt et al., 2019) also proposes a bench- mark for evaluating language learning griworld task. While BabyAl includes combinatorially diverse set of tasks using language, the tasks are procedurally generated and thus somewhat less "natural", the tasks and example trajectories approach, where the dataset consists of human human-generatee -enerated tasks with no enforced structure of the text, in hopes of training more realistic and capable models. We also liant on skill acquisition, as the model must able to string together complex action sequences to reach high success rates. CLIPort (Shridhar et al. 2021) studies language for robotic tasks, using CLIP (Radford et al., 2021) to transfer supervision from pure language to images, complementary to our investigation of transferring supervision from pure trajectories to language. Additionally, both
diverse. We propose more open-ended
environment
BabyAl Al and CLIPort, like many other works (Chaplot et al. 20018; Stepputtis et al. 20200 Abramson
Figure 8: Key frames from a generated agent rollout corresponding to task "jump between the second and third ice floes". States represent summary of the behavior; Appendix for longer trajectory.
8

---

Under review as a conference paper at ICLR 2022
et al. 2020 Nair al. 2021). multi-strean approach, embedding different inputs/modali- ties with separate models, in contrast to our single-stream approach which simply extends Decision Transformer (Chen al. (02T). Our environment differs from CLIPort in that does not require advanced attention mechanisms to be successful. Indeed, the Atari interface allows for impressive performance
with CNNs and MLPs. Another environment addressing language-conditione imitation is ALFRED (Shridhar et al,, 2020). Unlike ours, this environment provides an overarch- ing task, along with step-by-step instructions perform the task. Additionally, the time horizon in ALFRED is smaller than in our work (some missions take thousands of steps). Our dataset contains more image-action pairs, and focuses on oonger-horizon tasks, but in simpler environment without sub-tasks. In comparison other works such R2R (Anderson et al. 2018) and Touchdown (Chen ct al. 2020), our work requires less visual feature engineering due to the simpler input, but retains complex language and allows for longer time horizon missions.
Transformers. Our approach modeling trajectories follows Chen et al. (2021) and Janner et al. (202I), except condition behaviors by conditioning language tokens rather than target returns, which can yield a more diverse set of trajectories. These approaches use sequence mod- eling, rather than TD learning other RL transformers works have (Parisotto al. 2020; Ritter jet al.) (200)), which may be simpler and more stable. Transformers have also been more broadly applied to multimodal tasks, such vision and text (Lu et al. 2019 Radford al.) [021]) Like the above discussion, most multimodal approaches use separate models for different models for differ- ent modalities, e.g. CLIP (Radford al. 2021 learns separate models to embed text and images into a similar representation space. Our approach is akin to DALL-E (Ramesh et al. 202I), which uses single-stream approach model text and images, where we replaces images with trajectories, and do not predict the text. Prompt tuning (Lester et al.) 2021 Li and Liang 2021 has the potential turn a single moddlltyyy model into multimodal one with minimal finetuning (Lu al.) [02T]] Tsimpoukelli al., (021). and would be interesting future work. We showcase positive results from ability to transfer sequential behavior from unlabeled pretraining, which parallels trends seen in unsupervised language pretraining (Devlin et al. 2018).
7 CONCLUSION
We developed a new benchmark, Text-Conditioned Frostbite, and released corresponding dataset, with the aim accelerating progress the development of RL agents which capably utilize lan- guage. We also proposed Text Decision Transformer, multimodal architecture operating text, state, and action tokens. We showed results showing that unsupervised pretraining can im- prove downstream few-shot learning, both when transferring the vision encoder and the sequence backbone. We released code and the dataset used, and hope future work will continue exploring approaches for effective few-shot learning in language and RL settings.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems, arXiv preprint arXiv:2005.01643 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv:2106.01345 2021.
Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning as one big sequence mod- cling problem. arXiv preprint arXiv:2106.02039, 2021.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:/810.04800 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Jeffrey Pennington, Richard Socher, and Christopher D Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 conference empirical methods in natural language processing (EMNLP), pages 1532-15433 1543, 2014.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with pro- totypical representations. arXiv preprint orXiv:2102. 1271,, 2021.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Brenden M. Lake, Tomer D. Ullman, Joshua B. Tenenbaum, and Samuel J. Gershman. Building machines that learn and think like people. Behavioral and Brain Sciences, 40:e253, 2017. doi: 10.1017/50140555166000 16010161
Maxime Chevaller-Boisveer Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: A platform to study the sample efficiency of grounded language learning, 2019.
Mohit Shridhar, Lucas Manuelli, and Dieter Fox. CLIPort: What and where pathways for
robotic manipulation. In 5th Anmual Conference on Robot Learning, 2021. URL https:
wwwp:reviweww.n//oorrr.eeeee www.orrrmmmiiiiiii.immmm HRSIL
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. arXiv preprint arXiv:2103 0000000000 2021.
Devendra Singh Chaplot, Kanthashree Mysore Sathyendra, Rama Kumar Pasumarthi, Dheeraj Ra- jagopal, and Ruslan Salakhutdinov. Gated-attention architectures for task-oriented language grounding. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018
Simon Stepputtis, Joseph Camphell, Mariano Phielipp, Stefan Lee, Chitta Baral, and Heni Ben Amor. suaguage-conditionee imitation learning for robot manipulation tasks. arXiv preprint 0rXiv:2010.12083. 2020.
Josh Abramson, Arun Ahuja, lain Barr, Arthur Brussee, Federico Carnevale, Mary Cassin, Rachita Chhaparia, Stephen Clark, Bogdan Damoc, Andrew Dudzik, et al. Imitating interactive intelli- gence. arXiv preprint arXiv:20/2.05672. 2020.
Suraj Nair, Eric Mitchell, Kevin Chen, Brian Ichter, Silvio Savarese, and Chelsea Finn. Learn ing language-conditioned robot behavior from offline data and crowd-sourced annotation. arXiv preprint 109.02 2021.
10

---

Under review as a conference paper at ICLR 2022
Mohit Shridhar, Jesse Thomason, Daniel Gordon, Yonatan Bisk, Winson Han, Roozbeh Mottaghi, Luke Zettlemoyer, and Dieter Fox. Alfred: A benchmark for interpreting grounded instructions for everyday tasks, 2020.
Peter Anderson, Qi Wu, Damien Teney, Jake Bruce, Mark Johnson, Niko Sünderhauf, Ian Reid, Stephen Gould, and Anton van den Hengel. Vision-and-language navigation: Interpreting visually-groundee navigation instructions in real environments. Proceedings of the IEEE Con- ference on Computer Vision and Pattern Recognition (CVPR), 2018.
Howard Chen, Alane Suhr, Dipendra Misra, Noah Snavely, and Yoav Artzi Touchdown: Natural language navigation and spatial reasoning visual street environments, 2020.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg. Raphael Lopez Kaufman, Aidan Clark, Seb Noury, al. Stabilizing transformers for reinforcement learning. In International Conference on Machine Learning, pages 7487-7498. PMLR, 2020.
Sam Ritter, Ryan Faulkner, Laurent Sartran, Adam Santoro, Matt Botvinick, and David Raposo. Rapid task- in novel environments. arXiv preprint urXiv:2006.036663 2020.
Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee. Vilbert: Pretraining task-agnostic visiolin-
guistic representations for vision-and- -ann-uuue tasks. arXiv preprint arXiv:1908.0226.. 02265, 2019.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. arXiv preprint arXiv:2/02./2099, 2021.
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint
2021.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2101.00190 2021.
Kevin Lu, Aditya Grover, Pieter Abbeel, and Igor Mordatch. Pretrained transformers as universal computation engines. arXiv preprint arXiv:2103.05247, 2021.
Maria Tsimpoukelli, Jacob Menick, Serkan Cabi, SM Eslami, Oriol Vinyals, and Felix Hill. Multi- modal few-shot learning with frozen language models. arXiv preprint 1 13884, 2021.

---

Under review as a conference paper at ICLR 2022
Example Tasks
dont finish level minimize the time spent alive start playing once there are 30 seconds left jump via the second ice floe let the timer run out in every life
do not go below the first ice floe get to the third level make sure to collect fish complete the first igloo run into the left border dont
hit birds
reach each level as fast as possible
dont enter the first igloo
hit a bird on the third level
fall in the water every life
do not pass the second ice floe prioritize getting as many fish as possible
build the first igloos only play on your first life
dont enter the third igloo
run to the left before starting build but dont enter the third igloo
dont die from crabs hit bird before entering the first igloo
start playing when there are 20 seconds left die from water 4 times
Table 3: Example tasks in the Text-Conditioned Frostbite dataset. Some of these tasks are
included in evaluation so are not broken into easy, medium, and hard suites.
A EXPERIMENTAL DETAILS
In this section, detail performance on task-by-task basis, and list the tasks which are used for the experiments.
A.1 TASK SUITES
In the easy task suite, the evaluation tasks are 'dont move', 'stay the left side', 'try to stay on the right side', "dic on the first level', and "get as close to 500 points as you can'. the medium task suite, the evaluation tasks "jump back and forth between the second and third ice floes', 'get as close as you can to 1000 points', 'spend as much time as possible on the first ice floe',
A.2 BASELINES BY TASK
In this section, we include the performances of the baselines by task.
Performancc
0.3
MLP Baseline
1context
Scare
3context
10context
0.4
20context
you-OM Cin
00r4 nove
git mair porno 25 Soozal sooiio
RPrenate jumating senweee secomst and thas Roos
meach lwwel's
de on the second Iool
IRach lave
de the tnd mool
oot do0.0o000 soor 20-
the riy
as pocditie on ce
dort booee first
spennd muchh tive piociille Frrst
200
doot 15000
by AAd
Figure 9: Performance vs. Baseline. Above we see the performance of the model relative to
concatenation based MLP architecture on per-task basis.
12

---

Under review as a conference paper at ICLR 2022
Figure 10: Select frames from generated agent rollout corresponding to mission "jump between the second and third ice floes". States are sampled from every frames, and visualized left to right, top bottom. Note that many jumps still lost by subsampling.
A.3 FULL TRAJECTORY VISUALIZATION
Please see figure T0 for a more complete sequence a generated trajectory.
A.4 HYPERPARAMETERS FOR BASELINE TEXT DT
In this section, we share the hyperparameters we used for the baseline plots with the Text Decision Transformer. Please see table 4
Table 4: Hyperparameters of TDT for Frostbite experiments.
Hyperparameter
Value
Number of layers
Number of attention heads
Embedking dimension
128
Batch size
128
Context length K
20
Text context length
20
Nonlincarity
ReLU, encoder GeLUU, otherwise
Encoder channels
32, 64,64
Eneoder filter sizes
8 8, 4 4,3 x
Encoder strides
2,
Max epochs
Dropout
0.1
Learning rate
1 10
Adam betas
(0.9, 0.95)
Grad norm clip
1.0
Weight decay
0.1
13

---

Under review as a conference paper at ICLR 2022
A.5
ALGORITHM PSEUDOCODE
Here, we provide pseudocode for our implementation.
Algorithm Text-Conditioned Decision Transformer Pseudocode
m, 5, a, t: tasks, states, actions, or timesteps transformer: transformer with causal masking (GPT) embed 8, embed_as linear embedding layers embed m, embed t: learned discrete embedding tables pred_a: linear action prediction layer
main model def restteecisionrraassffrmee (m, 8, a, t):
# compute embeddings for tokens
m embedding = embed_m m (m)
embedding embed s (s) embed (t)
a_embedding embed (a) embed_t (t)
m,n, a1, embeds stack (s embedding, embedding) # interleave (s, a) input embeds - = concatenate (m embedding, input embeds) # prepend text
interleave tokens
(m
s_K)
use transformer to get hidden states corresponding to actions hidden states transformer (input embeds=input a hidden unstack (hidden_stated) actions
) embeds)
predict action logits
return pred (a hidden)
training loop
for (m, s, a, t) in dataloader:
a_preds = Textbeciisiontrrassffoomem (m, s, a, t)
loss = cross entropy loss (a preds, a) optimizer... zero ()ad loss b)ckward optimizer.
evaluation loop
m, s, a, done [task], reset ()), [I, [1], False
while not done: # autoregressive generation/sange sampling
# sample next action
action TextDecissoonrrassformee (m, env env. step (action)
t) sample ()
new_s, r, done,
append new tokens to sequence
s, a, + [new_s], a [action], + [t[-1]+1] 5, a, s[-K:]
only keep context length of
14


Reviews:
[{"review_id": "ys6CSRtjya7", "reviewer": "Reviewer_ZRZY", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think this paper lacks important detail in that it proposes a new benchmark but declines to compare such benchmark to existing work. The claim that \"transferring sequential knowledge is relatively unstudied\" is just simply not true. There has been several works, old and recent, on instruction following. These include the large quantity of visual instruction following tasks such as R2R (https://bringmeaspoon.org/), Room-Across-Room (https://ai.google.com/research/rxr/), Touchdown (https://arxiv.org/abs/1811.12354), DRIF (https://github.com/lil-lab/drif), ALFRED (http://askforalfred.com/), as well as complex grid worlds with new language-entity associations and hard generalization requirements (https://arxiv.org/abs/1910.08210, https://arxiv.org/abs/2101.07393). In order to propose this benchmark, the paper ought to compare the proposed environment with these existing work - why is there a need for this new benchmark? In what ways is it different than previous environments? This crucial question is not addressed in this work. These related work and many others (https://arxiv.org/abs/1906.03926) are not even cited. Another point of contention I have with the way this work is posed is that it claims 5 million annotated transitions. What does this mean? It seems like there are only a few hundred labeled trajectories. Do the authors means they have a couple hundred labeled trajectories, each with 10k steps (where the task is fixed)? If so, I think this is at best misleading - no other dataset for language grounding that I am aware of lists annotation size like this.\n\nHere are some concrete examples of claims with missing analyses:\n1. \"a diverse set of tasks\" - how are these tasks diverse? What do the tasks even look like? The authors mention 500 tasks but there are only  8 listed in the appendix. Upon inspection of these tasks, they seem very unusual (die on the x level, try to get as close as you can to 1000 score). How many unique tasks are there? What does it even mean to \"get as close as you can\"? Or \"spend as much time as possible\"? How is the reward assigned?\n2. \"temporally extended strategies\" - what does this mean? what is the distribution of trajectory lengths of annotated playthroughs per task? of agent playthroughs per task?\n\nMoreover, what is the evaluation set? How is this set constructed? What skills is it evaluating? Is generalization required? If an agent just cycles through memorized action patterns from the training distribution, can it still do well on the test set?", "summary_of_the_paper": "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.", "summary_of_the_review": "No comparison to prior tasks, unclear what new proposed environment tests and how it compares to prior environments.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "RP0TTjrSUCN", "reviewer": "Reviewer_EZAA", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I agree that this paper addresses an important problem of language-conditioned imitation learning. However, It seems that the contribution of the proposed model, Text Decision Transformer (Text-DT), is slightly marginal and Text-DT seems not related to Decision Transformer. The detailed comments and questions are as follows:\n\n1. The biggest difference between standard Transformer and Decision Transformer is the return-to-go, but there is no return-to-go in Text-DT. The Text-DT, proposed in this paper, is thought to be just using the standard transformer as a conditional sequence generation method, so I am not sure if it is right to call it a decision transformer. Why is it called a (Text) decision transformer?\n\n2. The authors use Text-DT and TDT interchangeably in their papers, is there any difference? If there is a difference, it should be clearly explained, and if there is no difference, I think the terminology should be unified.\n\n3. (Figure 3) In easy/medium tasks, TDT(k=1) performs better than MLP baseline, but in hard tasks, MLP baseline shows better performance than TDT(k=1). What is the reason for this?\n\n4. As an experiment in Figure 3, how does the performance of the MLP baseline change as the context length increases? It would be better if the authors could compare and show them together.\n\n5. The authors describe the proposed Text-DT model as the main contribution, but there does not seem to be much difference from the existing works using the standard transformer architecture. Rather, it would be better to show the contribution to the point of proposing a new benchmark more intensively.\n\n", "summary_of_the_paper": "This paper presents a new multimodal benchmark for language conditioned RL settings, where an agent must complete tasks specified by text instructions in the Atari Frostbite environment. Their benchmark provides a dataset of 5M text-labeled transitions for training. Finally, the authors propose a model for language conditioned RL settings based on Transformer architecture as a baseline.", "summary_of_the_review": "I think that this paper presents an interesting benchmark for developing RL agents which can utilize language. However, the proposed model, Text-DT, is considered to have a small contribution because the Transformer architecture is simply applied to the conditional sequence generation problem with language inputs. In addition, the authors explain the proposed model based on the decision transformer, but it is not considered to have much related to the decision transformer. I would like to recommend a more detailed explanation of the proposed benchmark and possible future studies.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "FqoYcpEFNl", "reviewer": "Reviewer_hbz9", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n1. This paper introduces a quite large text corpus for the text game community. \n2. This paper presents some baselines for pretrained text decision transformers. \n\nWeakness:\n\nMay need some interpretations for some weird cases in the Table and Figures. \n\nFor example, \n\n1). In Table 4, it looks more context does not help for Medium tasks?\n\n2). Similarly, in Figure 3, it looks MLP is quite a strong baseline on the Hard tasks? ", "summary_of_the_paper": "This paper introduces a corpus of text instructions for the language guided Atari Frostbite environment, and also presents a study of pretrained Text Decision Transformer on the corpus, shows that it can improves upon baselines, yield better results in the low-data regimes. ", "summary_of_the_review": "This paper introduces a quite large corpus of text instructions for the language guided Atari Frostbite environment, which is good resources for the community, and it also presents baselines for pretrained text decision transformers (TDT), and shows some interesting empirical results. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9h8GKvbQPOr", "reviewer": "Reviewer_ZbHJ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros\n* Learning behavior policies from language supervision is an important problem setting\n* Useful to study the applicability of the transformer architecture to sequential decision making problems\n* New benchmark on language conditioned agent learning\n\nCons\n* The new proposed benchmark needs to be contextualized in existing work. Why is there a need for this benchmark and how is it different from existing language instruction following benchmarks from the embodied learning literature (Eg. VLN benchmarks like Room to room and object interaction benchmarks like ALFRED)?\n* Models and training approaches considered in the paper seem straightforward. What are the new technical contributions of the work?\n* Weak experiments \n\nI would encourage the authors to think about what’s new and interesting about the proposed new dataset, what makes it challenging and how it is related to and different from other existing benchmarks.\n\nProblem setting is unclear - In section 5.2, I don’t understand why the model is trained on the evaluation task. Also, the section ends abruptly without a conclusion. \n\nThe conclusions on pretraining and comparisons against context size 1 seem obvious and expected. I didn't learn much from these experiments. The experiments needs to provide more insight into the dataset and the results. \n\nPresentation \n- Alg 1 can be moved to the appendix. \n- The MLP baseline wasn’t clearly described. ", "summary_of_the_paper": "This paper proposes a new language conditioned imitation learning task based on the Atari Frostbite environment. In addition, the paper studies imitation learning with a transformer architecture and the effect of pre-training on game trajectories.", "summary_of_the_review": "The paper lacks novel technical contributions. The experiments are weak and do not provide new, meaningful and practically useful insights. The work further needs to be contextualized better in existing work. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The arguments the paper makes require a stronger foundation and justification. The reviewers and AC didn't find the author response sufficient. For example, in response to ZbHJ, the authors argue that their benchmark doesn't use automatically generated trajectories and therefore the language is not synthetic in some sense. It's not clear how it's related to synthetic language, but generated trajectories does create artificial regularities in the task, so an issue, but one that the author must address accurately. This argument also seems to focus on ALFRED and R2R, and ignored many other benchmarks, like the data used in DRIF (mentioned later), RxR, Touchdown, etc. There is also mis-used of technical term (e.g., Decision Transformer). Generally, the reviewers consider the work of potential, but it requires significant refinement, which the author response did not provide.
Decision: Reject