Transformer Adapters for Robot Learning
Abstract: Large transformer-based architectures are capable of complex robot task planning and low-level control.
In the natural language processing (NLP) community, ﬁne-tuning large pretrained models (PTMs) such as GPT-3 and PaLM is the de-facto standard.
With the scalability of transformer models and growing availability of large-scale multimodal robot data, we investigate pretraining large backbone models to capture useful behavioral priors that enable efﬁcient few-shot transfer to downstream robot tasks.
We explore the setting of modular reinforcement learning (RL) in which each downstream task is encapsulated by an independently learned module.
With many downstream tasks, ﬁne-tuning or training separate copies of these large PTMs become computationally and memory intensive.
We propose to pretrain a large transformer backbone on task-agnostic data and learn small task-speciﬁc adapters using few-shot imitation learning to quickly adapt to downstream tasks.
We evaluate on complex robot manipulation tasks in the Metaworld environment and demonstrate that adapter training is a parameterefﬁcient approach for modular RL.
Introduction
As large pretrained models (PTMs) are being adopted in robotics and eventually deployed onto physical robot systems, how can non-industry practitioners adapt these large PTMs to custom tasks beyond just relying on their zero-shot capabilities?
We present a scalable approach, inspired by work in the NLP community, for a large pretrained generalist robot model to acquire diverse behavior and quickly adapt to new downstream tasks.
Large pretrained models such as PaLM [1], DALL-E [2] and GPT-3 [3] have demonstrated impressive capabilities across many domains ranging from natural language processing (NLP) to vision [2] and even robotics [4] [5] [6] [7].
Fine-tuning these large PTMs can incur signiﬁcant training and storage costs.
GPT-3, for example, has 175 billion parameters, cost $10 million and a few months of training, and takes up 350 GB of storage [3].
The prevalent paradigm for overcoming large model training costs in learning new NLP tasks is to train small neural modules called adapters [8] on a relatively small amounts of data speciﬁc to the task.
Adapters are embedded inside Transformer blocks of large PTMs.
This design facilitates weight sharing and accelerates learning of new NLP tasks by leveraging pretrained linguistic priors.
We advocate for a similar framework whereby an autoregressive Transformer-based policy is pretrained on a large, diverse task-agnostic robot dataset.
Small task-speciﬁc adapter modules are then learned for each new downstream task while sharing a common transformer backbone.
We demonstrate on a challenging robot manipulation environment that a pretrained Transformer model can zero-shot perform unseen tasks and task-speciﬁc adapters can quickly adapt to new tasks with a few demonstrations using < 2% of the full model parameters.
[Figure]
Transformer
Policy
[Figure]
Approach
We train an imitation learning policy for pretraining transformers using task-agnostic experience dataset.
Thereafter, we learn adapters with frozen transformer backbone for downstream task adaptation using few-shot task demonstrations.
2.1 Preliminaries
Transformer Policy.
Decision Transformers (DT) [9] demonstrate that self-attention based models can be applied to solve RL problems due to its efﬁciency and scalability in modeling long sequential data.
Rather than learning explicit value functions, DT treats policy learning as a causal sequence modeling problem.
DT uses a Transformer backbone, speciﬁcally GPT-2 [10], trained on a causal modeling objective of predicting the next action token given the context history.
Trajectories are sequences of states s t , actions a t , and reward-to-go ˆ r t tokens.
Unlike DT, which learns a reward-conditioned policy, we pretrain our model on valid trajectories only τ = { s t , a t } Tt =1 (Figure 2), followed by goal-conditioned few-shot imitation learning.
Adapters introduce a small set of new parameters between the layers of a large pretrained model.
During downstream training, the pretrained model weights are ﬁxed while the adapter weights learn to encode task-speciﬁc representations.
Adapters efﬁciently share a majority of the weights with the pretrained model, and this facilitates information sharing that improves downstream learning.
In practice, adapters are very lightweight, typically using only 0.5-8% of the full backbone model parameters.
Houlsby et al.
[8] empirically show that a two-layer feed-forward bottleneck architecture works the best for the adapter design.
Adapters are deﬁned as:
[Equation]
where x is the input hidden state, W lup , W ldown are the weights for the feed-forward up and downprojection for layer l respectively.
Adapters learn to encapsulate task-speciﬁc information that is non-destructive to the original model [8].
Adapters are modular by design.
2.2 Problem Formulation
We assume access to a task-agnostic dataset D PT =
{ τ i } Ni =1 for pretraining the backbone model.
Each trajectory τ i is collected by an agent interacting with a Markov Decision process (MDP)
de-
[Figure]
ﬁned by a tuple M = {S , A , R , T } of states, actions, rewards, and transition probability.
Data can be collected from an expert planner, play interactions, or even a random policy.
We focus on an imitation learning setting in which the MDP does not have an explicitly deﬁned reward function, i.e
M \ R .
The objective is to learn an imitation policy π ( a | s ) such that the agent behaves like the expert provided a set of expert demonstrations for a task.
During downstream task adaptation, we assume that there is a dataset D test containing a few demonstrations for each target task T i ∈ T test .
2.3 Transformer Adapters
Our approach has two stages: pretraining and downstream task learning.
During pretraining, the model learns to capture useful behavioral priors from the task-agnostic data, such as how to manipulate different objects.
By leveraging the capacity of large models, we capture diverse prior knowledge that can accelerate learning of downstream tasks.
Following [9], our pretraining objective is to minimize the mean-squared error between the predicted action and ground truth actions:
[Equation]
where H is the context length, θ are the policy parameters.
During the ﬁne-tuning stage, we use the pretrained model to bootstrap the learning of new downstream tasks via few-shot imitation learning.
The agent is provided with a small set of expert demonstrations { τ i } Ki =1 for each target task.
We insert adapter modules into the layers of the pretrained transformer (see Figure 2), and only train the weights of the adapter with the pretraining objective using the expert demonstrations.
Experiments and Results
We aim to answer two main questions: 1) Does a transformer policy pretrained on large, diverse ofﬂine data provide a strong backbone for transfer to downstream unseen tasks?
Are adapters a parameter-efﬁcient and scalable approach for few-shot imitation learning of new robot tasks?
Environment and Data We evaluate our method in the MetaWorld robot manipulation benchmark [11].
MetaWorld contains 50 different robot control and manipulation tasks with a Sawyer arm in the Mujoco environment.
We evaluate on the Meta-Learning 45 ( ML45 ) setting, where we pretrain our policy on 45 tasks and evaluate downstream adaptation to 5 held-out tasks.
The task splits are hand-selected in [11] such that the training tasks are structurally similar to the test tasks.
Object and goal positions are randomized at the start of every episode.
The action space A ∈ R 4 , is the ∆( xyz ) of the end-effector, and a continuous scalar value for gripper torque.
The state S ∈ R 39 , contain the 3D positions of the end-effector, ﬁrst object, second object and the goal.
The goal position is masked during pretraining and provided to the model during downstream adaptation.
We use a scripted policy to collect 10 trajectories for each of the 50 tasks.
We pretrain on 450 trajectories from the training split and few-shot imitation learned on 10 demonstrations per test task.
We report averaged task success across all 5 test tasks over 10 evaluation rollouts and 3 seeds.
[Figure]
• MLP w/ prediction head ( MLP ).
Fine-tune a copy of the action prediction head for each downstream task.
This comparison measures the advantage of using self-attention based Transformer models for capturing strong behavioral priors.
• Transformer w/ prediction head ( ActHead ).
Fine-tune a copy of the action prediction head for each new task.
This measures the beneﬁt of training intermediate representations that are dynamically stitched into a pretrained model versus training the last few layers.
• Fine-tuning entire Transformer ( FT Full ).
Fine-tune the weights of the full model.
This is an upper bound for downstream performance if we do not limit the model capacity.
Results and Discussion Learning task-speciﬁc adapter modules on
a pretrained transformer backbone ( Adapter ) can almost match the performance of ﬁne-tuning the entire model ( FT Full ) while using less than 2% of the pretrained model parameters (Figure 3).
Compared to the full transformer backbone with 200 Mb of storage, each individual task adapter requires less than 2 Mb of storage.
This reduced storage overhead becomes more substantial when working with large foundation models and deployment on physical robot systems.
Adapter converges slower than FT Full as adapter weights are initialized to zero and trained from scratch.
The MLP model, which has slightly more trainable weights than the Adapter model, fails to converge to good average task performance.
Methods using a frozen pretrained transformer backbone ( Adapter , ActHead , FT Full ) exhibit positive zero-shot performance on several of the evaluation tasks, indicating that the pretrained model may have captured some useful behavioral priors that can generalize without additional ﬁnetuning.
In comparison, the MLP backbone pretrained on the same data is unable to generalize zeroshot to unseen tasks.
By pretraining on even more diverse task-agnostic data, we hypothesize that a transformer-based backbone models can provide even stronger zero-shot performance gains.
Additionally, zero-shot performance can accelerate online learning because the learned priors enable more guided exploration for the agent to quickly discover reward states.
Conclusions and Future Work
We present a two-stage scalable framework for robot learning: task-agnostic transformer policy pretraining followed by parameter-efﬁcient few-shot downstream task adaption using adapters.
Taskspeciﬁc adapters outperform task-speciﬁc action heads with the same transformer backbone, while being comparable to ﬁne-tuning the full model.
In the future, we hope to extend our setup to other multi-task environments and conduct experiments with visual perception-based states.
There is a great breadth of work in the NLP community on parameter-efﬁcient ﬁne-tuning that can be adapted for robot learning.
We can explore orthogonal approaches such as LoRA [12] and Compacter [13] that are even more optimized for training efﬁciency and reducing storage cost.
Moreover, we can extend our method to handle long-horizon, compositional tasks using AdapterFusion [14]. Acknowledgments
The authors would like to thank Abrar Anwar, Tejas Srinivasan and others from the GLAMOR lab for their advice and fruitful discussions.
A. Related Works
Transformers for Robot Learning.
Recently, there has been a surge of research in the robotics community applying large Transformer models to a variety of complex robot learning problems.
Decision Transformers (DT) [9] showed that the causal self-attention mechanism of Transformer models can be used to perform credit assignment in sequential decision making.
Perceiver-Actor [6] uses a large multi-modal Transformer to learn a multi-task policy for solving robotic manipulation task.
SayCan [4] show that large pretrained language models (PTLMs) encapsulate a wealth of commonsense knowledge that can be used for efﬁcient high-level task planning in the real-world.
VIMA [7] shows that a Transformer-based model trained on large amounts of expert demonstration can zero-shot generalize to new tasks using multimodal prompts.
Scaling model and data.
In the natural language processing (NLP) community, there is a recent trend towards increasingly larger Transformer-style models trained on large, diverse text corpuses.
ELMo, one of the earliest works on contextualized word representations, published in 2018 had roughly 94 million parameters.
Within the span of three years, OpenAI released GPT-3 [3] ( billion parameters) and Microsoft and NVIDIA introduced Megatron-Turing [15] (530 billion parameter model).
GPT-4 is expected to have about 100 trillion parameters.
GPT-3 is reported to be trained on 500 billion tokens of internet text.
In contrast, the quantity of robot data available is nowhere near that magnitude hence the slow adoption of the pretraining-ﬁnetuning paradigm.
There have been several recent attempts to curate and publish larger, more diverse robot datasets to address this gap (e.g. RoboNet [16], Bridge Data [17], VIMA [7], etc).
Parameter-efﬁcient Fine-tuning in NLP.
As Transformer models become increasingly large and more resource intensive, it is infeasible to ﬁne-tune the entire model for each new downstream task.
Several alternatives have been proposed that update only a small number of extra parameters while keeping the backbone model parameters frozen.
For example, Adapter [8] tuning inserts a small set of trainable weights between each layer of the pretrained Transformer.
[18] empirically show that Adapters are 60% faster than full-model tuning in terms of computational efﬁcient because of the decrease in overhead in gradient computation and are signiﬁcantly more storage efﬁcient.
More recently, LoRA [12] and Compacter [13] learns low-rank matrices to approximate parameter updates.
B. MetaWorld Downstream Tasks
Object and goal positions are randomized at the start of every episode during data collection and time.
C. Hyperparameters
[Table]
Value context length learning rate 1elearning rate decay 1enumber of layers number of attention
headsdropout 0.
embedding dimension batch size max episode length number of epochs