A Generalist Agent

Scott Reed, Konrad Żołna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Giménez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, Nando de Freitas

Seminal May 2022

TL;DR One 1.2B-parameter transformer, with a single set of weights, is trained to predict the next token on data from 604 control tasks plus text and images, and it plays Atari, captions images, chats and stacks blocks with a real robot arm. It scores above 50% of the expert on over 450 of the 604 tasks.

What they do

Gato applies the large-language-model recipe to control. Every modality (text, image patches, proprioception, joint torques, button presses) is serialized into one flat token sequence, and a single decoder-only transformer is trained on it by supervised next-token prediction. The same weights play Atari, navigate DM Lab, caption images, chat and stack blocks with a real Sawyer arm. Model size was held to about 1.2B parameters so it can control a real robot in real time.

How it works

Tokenization. SentencePiece text (32,000 subwords); 16x16 image patches embedded by a ResNet block; continuous values mu-law encoded and binned into 1024 levels. Each timestep is observation tokens, a separator, then action tokens.
Model and loss. 1.18B-parameter transformer, 24 layers, 1024-token context. The loss covers only text and logged actions, not image or observation tokens.
Data. Trajectories from specialist RL agents, filtered to episodes with at least 80% of expert return, mixed with MassiveText and vision-language datasets. Trained for 1M steps on a 16x16 TPU v3 slice (about 4 days).
Task specification. No task IDs: 25% of training sequences get a prompt from another episode of the same task, and at evaluation Gato is prompted with a successful demonstration.

Results

Above 50% of expert score on over 450 of 604 tasks; human level or better on 23 Atari games; over 50% on 44 of 45 Meta-World tasks.
Real-robot RGB Stacking with unseen object shapes: 50.2% average success, against 49% for the single-task BC-IMP baseline. With fine-tuning it matches the expert after 10 episodes and exceeds it at 100-1000.
A new real task, stacking blue on green: after fine-tuning on 500 demonstrations (2 h 45 min), 60% success, against 0.5% for BC trained from scratch.
For the same number of training tokens, larger models score higher (79M < 364M < 1.18B).

Limitations

Gato learns only by imitating RL-agent data, and no web-scale control dataset exists. With images, the 1024-token context holds only a few timesteps, so demonstration prompts did not significantly help on new environments; adapting needed fine-tuning. On Atari the single-task experts still beat it, pretraining gave no benefit on the held-out boxing game, and its dialogue and captioning are rudimentary.

Figures

From Figure 1: One network with one set of weights senses and acts across many embodiments and environments; Gato was... — From Figure 1: One network with one set of weights senses and acts across many embodiments and environments; Gato was trained on 604 distinct tasks with different modalities, observations and action specifications.

From Figure 2: Data from every task and modality is serialized into a flat token sequence, batched and processed by a... — From Figure 2: Data from every task and modality is serialized into a flat token sequence, batched and processed by a transformer like a large language model; the loss is applied only to text and action targets.

From Figure 5: Number of simulated control tasks, by domain, on which the pretrained Gato scores at or above each per... — From Figure 5: Number of simulated control tasks, by domain, on which the pretrained Gato scores at or above each percentage of expert score (0 is a random agent).

From Figure 10: Real-robot Skill Generalization success rate, averaged over test triplets, against the number of fine... — From Figure 10: Real-robot Skill Generalization success rate, averaged over test triplets, against the number of fine-tuning episodes for Gato, the expert, and CRR trained on 35k expert episodes (upper bound).

Abstract

Inspired by progress in large-scale language modeling, we apply a similar approach towards building a single generalist agent beyond the realm of text outputs. The agent, which we refer to as Gato, works as a multi-modal, multi-task, multi-embodiment generalist policy. The same network with the same weights can play Atari, caption images, chat, stack blocks with a real robot arm and much more, deciding based on its context whether to output text, joint torques, button presses, or other tokens. In this report we describe the model and the data, and document the current capabilities of Gato.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong