SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shalev Lifshitz $^{1,2*}$
shalev.lifshitz@mail.utoronto.ca
Keiran Paster $^{1,2*}$
keirp@cs.toronto.edu
Harris Chan $^{1,2\dagger}$
hchan@cs.toronto.edu
Jimmy  $\mathbf{Ba}^{1,2}$
jba@cs.toronto.edu
Sheila McIlraith<sup>1,2</sup>
sheila@cs.toronto.edu
$^{1}$ Department of Computer Science, University of Toronto, Toronto, Canada.
$^{2}$ Vector Institute for Artificial Intelligence, Toronto, Canada.
Constructing AI models that respond to text instructions is challenging, especially for sequential decision-making tasks. This work introduces a methodology, inspired by unCLIP, for instruction-tuning generative models of behavior without relying on a large dataset of instruction-labeled trajectories. Using this methodology, we create an instruction-tuned Video Pretraining (VPT) model called STEVE-1, which can follow short-horizon open-ended text and visual instructions in Minecraft™. STEVE-1 is trained in two steps: adapting the pretrained VPT model to follow commands in MineCLIP's latent space, then training a prior to predict latent codes from text. This allows us to finetune VPT through self-supervised behavioral cloning and hindsight relabeling, reducing the need for costly human text annotations, and all for only $60 of compute. By leveraging pretrained models like VPT and MineCLIP and employing best practices from text-conditioned image generation, STEVE-1 sets a new bar for open-ended instruction following in Minecraft with low-level controls (mouse and keyboard) and raw pixel inputs, far outperforming previous baselines and robustly completing 12 of 13 tasks in our early-game evaluation suite. We provide experimental evidence highlighting key factors for downstream performance, including pretraining, classifier-free guidance, and data scaling. All resources, including our model weights, training scripts, and evaluation tools are made available for further research.
The ability to use text instructions to control and interact with powerful AI models has made these models accessible and customizable for the masses. Such models include ChatGPT [41], which can respond to messages written in natural language and perform a wide array of tasks, and Stable Diffusion [50], which turns natural language into an image. While those models cost anywhere from hundreds of thousands to hundreds of millions of dollars to train, there has been an equally exciting trend whereby powerful open-source foundation models like LLaMA [59] can be finetuned with surprisingly little compute and data to become instruction-following (e.g., [58, 13]).
In this paper, we study whether such an approach could be applicable to sequential decision-making domains. Unlike in text and image domains, diverse data for sequential decision-making is very expensive and often does not come with a convenient "instruction" label like captions for images. We propose to instruction-tune pretrained generative models of behavior, mirroring the advancements seen in recent instruction-tuned LLMs like Alpaca [58], without relying on a large dataset of instruction-labeled trajectories.
Figure 1: Like unCLIP [48], our approach involves two models. First, we train the policy by finetuning VPT to achieve goals given by pretrained MineCLIP [17] visual embeddings using our gameplay dataset. Second, for the prior model, we train a CVAE [54] to sample MineCLIP visual embeddings given a text prompt. The combination of these two models enables our agent to follow text and visual instructions.
In the past year, two foundation models for the popular open-ended video game Minecraft™ were released: a foundation model for behavior called VPT [5] and a model aligning text and video clips called MineCLIP [17]. This has opened up an intriguing avenue to explore fine-tuning for instruction-following in the sequential decision-making domain of Minecraft. VPT was trained on 70k hours of Minecraft gameplay, so the agent already has vast knowledge about the Minecraft environment. However, just as the massive potential of LLMs was unlocked by aligning them to follow instructions, it is likely that the VPT model has the potential for general, controllable behavior if it is finetuned to follow instructions. In particular, our paper demonstrates a method for fine-tuning VPT to follow short-horizon text instructions with only $60 of compute and around 2,000 instruction-labeled trajectory segments.
Our method draws inspiration from unCLIP [48], the approach used to create the popular text-to-image model DALL·E 2. We decompose the problem of creating an instruction-following Minecraft agent into two models: a VPT model finetuned to achieve visual goals embedded in the MineCLIP latent space, and a prior model that translates text instructions into MineCLIP visual embeddings. We finetune VPT using behavioral cloning with self-supervised data generated with hindsight relabeling [3], avoiding the use of expensive text-instruction labels in favor of visual MineCLIP embeddings. We apply unCLIP with classifier-free guidance [23] to create our agent called STEVE-1, which sets a new bar for open-ended instruction following in Minecraft with low-level controls (mouse and keyboard) and raw pixel inputs, far outperforming the baseline set by Baker et al. [5].
- We introduce a methodology, inspired by unCLIP [48], for instruction-tuning generative models of behavior without relying on a large dataset of expensive instruction labels.
- We apply this methodology to create STEVE-1, a Minecraft agent that can follow short-horizon open-ended text and visual instructions with a high degree of accuracy, all for only $60 of compute. We perform extensive evaluations of our agent, showing that it can robustly complete 12 of 13 goal-conditioned control tasks in our early-game evaluation suite in Minecraft. For long-horizon tasks like crafting and building, we show that a basic version of prompt chaining can dramatically improve performance.
- We provide experimental evidence highlighting key factors for downstream performance, including pretraining, classifier-free guidance, data scaling, prompt-engineering, and other design choices. In particular, we show that unCLIP [48] and classifier-free guidance [23] translate well to sequential decision making and are essential for strong performance.
- We release model weights for STEVE-1 as well as training scripts and evaluation code to help foster more research into instructable, open-ended sequential decision-making agents.2
<sup>1</sup>Short-horizon tasks require few steps: e.g., go to a tree and chop it down, dig a hole. Long-horizon tasks take many steps: e.g., craft complex recipes from scratch, build a house.
$^{2}$ Model weights, training code, videos, and an interactive demo script are hosted on our project webpage at https://sites.google.com/view/steve-1.
Minecraft as a Test-bed for AI Minecraft has gained popularity as a benchmark for AI research due to its complex and dynamic environment, making it a rich test-bed for reinforcement learning and other AI methods (e.g., [26, 19, 17, 21, 40, 62, 38, 9]). We leverage the MineRL environment [19] to research the creation of agents that can follow open-ended instructions in complex visual environments using only low-level actions (mouse and keyboard). We build STEVE-1 on top of two recent foundation models. In order to align text and videos, we use MineCLIP [17], a CLIP [47] model trained on paired web videos of Minecraft gameplay and associated captions. To train STEVE-1's policy, we fine-tune VPT [5], a foundation model of Minecraft behavior that is pretrained on 70k hours of web videos of Minecraft along with estimated mouse and keyboard actions. Several prior works [61, 62] have explored the use of LLMs in creating instructable Minecraft agents. These works typically use LLMs to make high-level plans that are then executed by lower-level RL [40, 62] or scripted [46] policies. Since STEVE-1 is a far more flexible low-level policy, the combination of STEVE-1 with LLMs is a promising direction for future work. Fan et al. [17] introduced an agent trained using RL with MineCLIP as a shaping reward on 12 different tasks and conditioned on MineCLIP-embedded text-prompts. However, this agent failed to generalize beyond the original set of tasks without further RL finetuning using the MineCLIP reward function. Cai et al. [9] proposed a Goal-Sensitive Backbone architecture for goal-conditioned control in Minecraft which is trained on a fixed set of goals, while STEVE-1 learns goal-reaching behavior from a large dataset in a self-supervised way without training on an explicit set of tasks.
Foundation Models for Sequential Decision-Making Foundation models which are pretrained on vast amounts of data and then finetuned for specific tasks have recently shown great promise in a variety of domains including language [8, 14, 59], vision [48, 10, 47], and robotics [7, 53, 25, 39, 65]. GATO [49] and RT-1 [7] have demonstrated the potential of training transformers to perform both simulated and real-world robotic tasks. With the exception of Kumar et al. [30], which uses Q-learning, the vast majority of cases [32, 7, 49] where deep learning has been scaled to large, multitask offline-RL datasets have used supervised RL. Supervised RL (e.g., [42, 18, 12]) works by framing the sequential decision-making problem as a prediction problem, where the model is trained to predict the next action conditioned on some future outcome. While these approaches are simple and scale well with large amounts of compute and data, more work is needed to understand the trade-offs between supervised RL and Q-learning or policy gradient-based methods [43, 44, 6, 55]. Recent works explore the use of hindsight relabeling [3] using vision-language models [47, 2] to produce natural language relabeling instructions. DIAL [65] finetunes CLIP [47] on human-labeled trajectories, which is then used to select a hindsight instruction from a candidate set. Sumers et al. [56] uses Flamingo [2] zero-shot for hindsight relabeling by framing it as a visual-question answering (VQA) task. In contrast, STEVE-1 relabels goals using future trajectory segment embeddings given by the MineCLIP [17] visual embedding.
Text-Conditioned Generative Models There has been a recent explosion of interest in text-to-X models, including text-to-image (e.g., [48, 51, 50]), text-to-3D (e.g., [27, 35]), and even text-to-music (e.g., [1]). These models are typically either autoregressive transformers modeling sequences of discrete tokens [60, 8] or diffusion models [24]. Most related to our work is unCLIP, the method used for DALL•E 2 [48]. unCLIP works by training a generative diffusion model to sample images from CLIP [47] embeddings of those images. By combining this model with a prior that translates text to visual CLIP embeddings, unCLIP can produce photorealistic images for arbitrary text prompts. unCLIP and many other diffusion-based approaches utilize a technique called classifier-free guidance [23], which lets the model trade-off between mode-coverage and sample fidelity post-training. We utilize the basic procedure of unCLIP and classifier-free guidance for training STEVE-1.
Inspired by the rapid recent progress in instruction-tuning Large Language Models (LLMs), we choose to leverage the recently released Video Pretraining (VPT) [5] model as a starting point for our agent. Since VPT was trained on 70k hours of Minecraft gameplay, the agent already has vast knowledge about the Minecraft environment. However, just as the massive potential of LLMs was unlocked by aligning them to follow instructions, it is likely that the VPT model has the potential
Figure 2: To create goal-conditioned data for finetuning, we randomly select timesteps from episodes and use hindsight relabeling to set the intermediate goals for the trajectory segments to those visual MineCLIP embeddings. This self-supervised data teaches the agent which actions lead to which states.
for general, controllable behavior if it is finetuned to follow instructions. In this work, we present a method for finetuning VPT to follow natural, open-ended textual and visual instructions, which opens the door for a wide range of uses for VPT in Minecraft.
Our approach is inspired by unCLIP, the method behind the recent text-to-image generation model, DALL $\cdot$ E 2 [48]. Our goal is to create a generative model of behavior in Minecraft conditioned on text instructions  $y$ . To do so, we utilize a dataset of Minecraft trajectory segments, some of which contain instruction labels  $y$ :  $[( \tau_1, y_1), (\tau_2, y_2), \ldots, (\tau_n, \emptyset)]$  where  $\tau$  is a trajectory of observations and actions. We also employ a pretrained CLIP model called MineCLIP [17], which generates aligned latent variables  $z_{\tau_{t:t+16}}$ ,  $z_y$ , where  $z_{\tau_{t:t+16}}$  is an embedding of any 16 consecutive timesteps from the trajectory. MineCLIP is trained using a contrastive objective on pairs of Minecraft videos and transcripts from the web. For simplicity of notation, we refer to the MineCLIP embedding of the last 16 timesteps of a trajectory segment as  $z_{\tau_{\text{goal}}}$ . Like unCLIP [48], we utilize a hierarchical model consisting of a prior and a policy:
- A prior  $p(z_{\tau_{\mathrm{goal}}} | y)$  that produces a latent variable  $z_{\tau_{\mathrm{goal}}}$  conditioned on a text instruction  $y$ .
- A policy  $p\left( {\tau  \mid  {z}_{{\tau }_{\text{goal }}}}\right)$  that produces a trajectory conditioned on a latent variable  ${z}_{{\tau }_{\text{goal }}}$  .
These two models can then be combined to produce a generative model of behaviors conditioned on text instructions:
$$
p (\tau | y) = p (\tau , z _ {\tau_ {\text {g o a l}}} | y) = p \left(z _ {\tau_ {\text {g o a l}}} | y\right) p (\tau | z _ {\tau_ {\text {g o a l}}}) \tag {3.1}
$$
3.1 Policy
To learn our policy, we finetune VPT, a foundation model of Minecraft behaviors  $p_{\theta}(\tau)$  trained on 70k hours of Minecraft gameplay videos. Specifically, VPT consists of a ResNet [22] that processes frames of dimension  $128 \times 128 \times 3$ , and a Transformer-XL [15] which processes the frame representations and autoregressively predicts the next action using the joint hierarchical action space described in Baker et al. [5]. In order to modify the architecture to condition on goal information, we add an affine transformation of  $z_{\tau_{\mathrm{goal}}}$  to the output of the ResNet before passing it to the transformer:
Process Frames:  $\operatorname{ResNet}_{\theta}(o_t) \to x_t$
[+ Conditioning on MineCLIP Embedding Goal]:  $x_{t} \rightarrow x_{t} + W_{\theta} z_{\tau_{\mathrm{goal}}} + b_{\theta}$
Predict Actions: TransformerXL $_{\theta}(x_t, \ldots, x_{t+T}) \to a_{t+T}$
In order to finetune VPT to condition on goals, we finetune the model using a method inspired by supervised RL approaches like Decision Transformer [12], GLAMOR [42], and GCSL [18]. We use a modification of hindsight relabeling which we call packed hindsight relabeling (see Figure 2) to generate a new dataset of trajectories with goals pulled from future states that periodically switch. In contrast with hindsight relabeling, packed hindsight relabeling packs multiple relabeled sequences into a single sequence. Specifically, our method to generate this dataset consists of two steps:
1. Given a trajectory  $\tau$  with  $T$  timesteps, randomly generate indices to select goals from:  $i_1, i_2, \ldots, i_n$ . These indices are chosen by starting at the first timestep and repeatedly sampling a new timestep by adding a random value to the previous timestep. This ensures that the data reflects that some goals may take longer to achieve than others.
2. For each chosen goal at timestep  $i_j$ , set the goals for timesteps  $i_{j-1} + 1, \ldots, i_j$  to be the goal at timestep  $i_j$ , denoted  $z_{\tau_{i_j}}$ .
Our final dataset  $\mathcal{D}_{\mathrm{relabeled}}$  consists of observation sequences  $(o_1, \ldots, o_T)$ , action sequences  $(a_1, \ldots, a_T)$ , and packed hindsight relabeled goals  $(z_1, \ldots, z_T)$ . We then finetune VPT on this dataset using a supervised loss to predict each action autoregressively using a causal attention mask:
$$
\mathcal {L} _ {\text {p o l i c y}} (\theta) = \mathbb {E} _ {\mathcal {D} _ {\text {r e l a b e l e d}}} \left[ - \log p _ {\theta} \left(a _ {t} \mid o _ {1 \dots t}, z _ {1 \dots t}\right) \right] \tag {3.2}
$$
3.2 Prior
In order to condition not only on embeddings of visual goals but on latent goals, we need the prior, a model that produces a latent variable  $z_{\tau_{\mathrm{goal}}}$  conditioned on a text instruction  $y$ . Our model is a simple conditional variational autoencoder (CVAE) [54, 29] with a Gaussian prior and a Gaussian posterior. Rather than learn to condition directly on text, we choose to condition on frozen text representations from MineCLIP  $z_y$ . Thus, the prior learns a function to translate from a text embedding  $z_y$  to a visual embedding  $z_{\tau_{\mathrm{goal}}}$  (see Appendix C.5). Both the encoder and decoder of our CVAE are parameterized as two-layer MLPs with 512 hidden units and layer normalization [4]. We train the model on our dataset, for which we have text labels  $\mathcal{D}_{\mathrm{labels}}$  using the following loss:
$$
\mathcal {L} _ {\text {p r i o r}} (\phi) = \mathbb {E} _ {\left(z _ {\tau_ {\text {g o a l}}}, z _ {y}\right) \sim \mathcal {D} _ {\text {l a b e l s}}} \left[ \mathrm {K L} \left(q _ {\phi} \left(z _ {\tau_ {\text {g o a l}}} \mid z _ {y}\right) \| p \left(z _ {\tau_ {\text {g o a l}}}\right)\right) - \mathbb {E} _ {c \sim q _ {\phi} \left(z _ {\tau_ {\text {g o a l}}} \mid z _ {y}\right)} \left[ \log p _ {\phi} \left(z _ {\tau_ {\text {g o a l}}} \mid c, z _ {y}\right) \right] \right] \tag {3.3}
$$
3.3 Datasets
To train our policy, we gather a gameplay dataset with 54M frames ( $\approx$  1 month at 20FPS) of Minecraft gameplay along with associated actions from two sources: contractor gameplay and VPT-generated gameplay. To train our prior, we use a small dataset of text-video pairs gathered by humans and augmented using the OpenAI API gpt-3.5-turbo model [41] and MineCLIP. See Appendix D for more detailed dataset information.
OpenAI Contractor Dataset We use 39M frames sourced from the contractor dataset which VPT [5] used to train its inverse dynamics model and finetune its policy. The dataset was gathered by hiring human contractors to play Minecraft and complete tasks such as house building or obtaining a diamond pickaxe. During gameplay, keypresses and mouse movements are recorded. We use the same preprocessing as VPT, including filtering out null actions.
VPT-Generated Dataset We generate an additional dataset of 15M frames by generating random trajectories using the various pretrained VPT agents. The diversity of this dataset is improved by randomly switching between models during trajectories [44], randomly resetting the agent's memory, and randomly turning the agent to face a new direction.
Text-Video Pair Dataset To train our prior model, which learns a mapping between text embeddings and visual embeddings, we also manually gather a small dataset of 2,000 text instructions paired with 16-frame video segments (less than a second) from our gameplay dataset. This dataset corresponds to less than 30 minutes of gameplay and takes just a few hours to collect. We augment this dataset by using the alignment between text and video embeddings from MineCLIP. For each text instruction, we find the top  $k$  most similar gameplay segments in our dataset and use the corresponding 16-frame segment as additional training data. For augmentation, we also add 8,000 text-instructions generated by the OpenAI API gpt-3.5-turbo model [41], in addition to our 2,000 hand-labeled instructions.
At inference time, we use the prior to sample a latent goal  $z_{\tau_{\mathrm{goal}}}$  from the text instruction  $y$ . We then use the policy to autoregressively sample actions  $a_t$  conditioned on the observation history  $o_{1\dots t}$  and the latent goal  $z_{\tau_{\mathrm{goal}}}$ . Similar to the observation in Appendix I of Baker et al. [5], even with conditioning the policy often fails to follow its instruction and simply acts according to its prior behavior. To mitigate this, we borrow another trick used in image generation models: classifier-free guidance. Specifically, during inference we simultaneously compute logits for the policy conditioned on the goal  $f(o_t,\dots,o_{t + 1},z_{\tau_{\mathrm{goal}}})$  and for the unconditional policy  $f(o_t,\dots,o_{t + 1})$ . We then compute a combination of the two logits using a  $\lambda$  parameter to trade-off between the two:
$$
\text {l o g i t s} = (1 + \lambda) \underbrace {f _ {\theta} \left(o _ {t} , \dots , o _ {t + 1} , z _ {\tau_ {\text {g o a l}}}\right)} _ {\text {c o n d i t i o n a l l o g i t s}} - \lambda \underbrace {f _ {\theta} \left(o _ {t} , \dots , o _ {t + 1}\right)} _ {\text {u n c o n d i t i o n a l l o g i t s}} \tag {3.4}
$$
By setting a higher value of  $\lambda$ , we can encourage the policy to follow actions that are more likely when conditioned on the goal and, as demonstrated in Section 4.5, this significantly improves performance. Also, in order to train the policy to generate these unconditional logits, we occasionally dropout the goal embedding  $z_{\tau_{\mathrm{goal}}}$  from the policy's input (with probability 0.1). This lets us generate both the conditional and unconditional logits using the same model with batch processing at inference time.
Evaluating the performance of our agent is a challenging task due to the wide variety of instructions that are possible and the difficulty of evaluating whether the agent has successfully achieved its task. We use a combination of programmatic evaluation metrics and automatic MineCLIP evaluation metrics to get a sense of the agent's capability level. We collectively refer to all of our evaluation tasks including the 11 evaluation tasks from Figure 3 and the two prompt chaining tasks from Section 4.3 as our early-game evaluation suite.
Programmatic Evaluation We compute programmatic evaluation metrics by monitoring the MineRL [19] environment state throughout each evaluation episode. As done in VPT [5], we compute multiple programmatic metrics including travel distance and early-game item collection. The travel distance is the maximum displacement of the agent along on the horizontal (X-Z) plane, measured from the initial spawn point. For early-game inventory counts, we store the maximum number of log, seed, and dirt items seen in the agent's inventory during the episode.
MineCLIP Evaluation We explore the use of text-visual alignment in MineCLIP latent space between trajectories and text or visual goals to evaluate our agent over a wider variety of tasks where programmatic evaluation isn't practical. To determine the degree to which a task has been completed at all during an evaluation episode, we record the minimum cosine distance between the (text or visual) goal embedding and the visual MineCLIP embedding at any timestep during an episode.
In our experiments, we aim to answer the following questions:
1. How well does STEVE-1 perform at achieving both text and visual goals in Minecraft?
2. How does our method scale with more data?
3. What choices are important for the performance of our method?
We base our implementation off of the official VPT codebase<sup>3</sup>. The main STEVE-1 is trained using Pytorch [45] distributed data parallel on four A40 GPUs for 160M frames, or just under three epochs of our gameplay dataset. Hyperparameters are selected to match those in Baker et al. [5] with the exception of learning rate, which we set to 4e-5. Our models are optimized using AdamW [37]. See Table 1 for a full list of hyperparameters.
4.2 Performance on Textual and Visual Goals