Real-World Robot Learning with Masked Visual Pre-training

Ilija Radosavovic, Tete Xiao, Stephen James, Pieter Abbeel, Jitendra Malik, Trevor Darrell

Seminal October 2022

PickSink: the xArm reaches into a cluttered toy sink for the detergent bottle.

TL;DR Pre-train a vision transformer as a masked autoencoder on 4.5M Internet and egocentric images, freeze it, and train small behavior-cloning policies on top for real robot tasks. Across 981 real-world trials it beats CLIP by up to 75% and supervised ImageNet and from-scratch encoders by up to 81%.

What they do

They test whether a single vision encoder, pre-trained with self-supervision on in-the-wild images, can serve as the perception module for many real robot tasks. The encoder is trained once as a masked autoencoder (MAE), frozen, and reused unchanged across six xArm tasks and two Allegro hand tasks, each learned by behavior cloning from at most 80 demonstrations.

How it works

Data. 4.5M images: 2.6M frames from Ego4D, 1.2M from ImageNet, and 700k hand-object interaction images from Epic Kitchens, Something-Something and 100 Days of Hands.
Pre-training. MAE masks random patches (e.g. 75%) and reconstructs them. Encoders are ViT-S, ViT-B and ViT-L (22M, 86M and 307M parameters).
Policy. A four-layer MLP takes the frozen image features from a wrist-mounted RGB camera plus joint positions and outputs delta joint angles at 5 Hz. 80 demonstrations per task, mostly from HTC Vive VR teleoperation.
Evaluation. Each task uses the same 16 object placements for every model, run back to back to hold lighting and robot condition fixed.

Results

Over 981 real trials, the MAE ViT-B beats CLIP (the strongest baseline), supervised ImageNet and from-scratch ViT-B encoders on reach, push and pick, and matches CLIP’s 80-demo performance with half the demonstrations. It also beats CLIP on the fridge, fruit and sink tasks.
On PickFruit: CLIP ViT-B 18.8%, R3M ResNet-50 31.3%, MAE ViT-S 68.8%, ViT-B 93.8%, ViT-L 100%.
Scaling the model from ViT-S to ViT-B helps only when the data grows too; ViT-L improves further, most on PickSink. A wrist camera beats a third-person one, and dropping either images or proprioception gives zero success.
Allegro hand: about 50% success reaching 8 seen and 45 unseen objects, and 50% over 30 cube-flipping trials.

Limitations

Mostly toy objects in clean lab scenes. Each task involves a single object and needs neither feedback control nor multi-step planning. The authors also note that robot up-time can affect results.

Videos

PickCube: the xArm grasps and lifts a yellow cube.

PushCube: the xArm pushes a wooden cube across the bin.

Figures

From Figure 1: Visual representations are learned with a masked autoencoder on over 4.5 million in-the-wild images fr... — From Figure 1: Visual representations are learned with a masked autoencoder on over 4.5 million in-the-wild images from five data sources; the encoder is then frozen and control policies for real robot tasks are learned on top.

From Figure 3: The six tasks on a 7-DoF arm with a parallel jaw gripper: reaching a red block, pushing a wooden cube,... — From Figure 3: The six tasks on a 7-DoF arm with a parallel jaw gripper: reaching a red block, pushing a wooden cube, picking a yellow cube, closing a toy fridge, picking fruits, and picking a detergent bottle from a cluttered sink.

From Figure 4: Success rates on ReachBlock, PushCube and PickCube for the MAE encoder versus CLIP, supervised ImageNe... — From Figure 4: Success rates on ReachBlock, PushCube and PickCube for the MAE encoder versus CLIP, supervised ImageNet and from-scratch encoders, all ViT-Base; the MAE encoder is best on every task.

From Figure 7: On PickFruit, going from ViT-S to ViT-B on the HoI images alone does not help, while scaling model and... — From Figure 7: On PickFruit, going from ViT-S to ViT-B on the HoI images alone does not help, while scaling model and data together does (left). Success keeps rising from the 86M ViT-B to the 307M ViT-L, with a larger gain on PickSink (middle, right).

Abstract

In this work, we explore self-supervised visual pre-training on images from diverse, in-the-wild videos for real-world robotic tasks. Like prior work, our visual representations are pre-trained via a masked autoencoder (MAE), frozen, and then passed into a learnable control module. Unlike prior work, we show that the pre-trained representations are effective across a range of real-world robotic tasks and embodiments. We find that our encoder consistently outperforms CLIP (up to 75%), supervised ImageNet pre-training (up to 81%), and training from scratch (up to 81%). Finally, we train a 307M parameter vision transformer on a massive collection of 4.5M images from the Internet and egocentric videos, and demonstrate clearly the benefits of scaling visual pre-training for robot learning.

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