R3M: A Universal Visual Representation for Robot Manipulation

Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, Abhinav Gupta

Seminal March 2022

TL;DR Pretrain an image encoder on Ego4D human videos and their language annotations, then freeze it as the perception module for behavior cloning. Across 12 simulated manipulation tasks it reaches about 62% success, more than 10% above CLIP, MoCo and ImageNet features, and on a real Franka it averages 56% versus CLIP's 24% from 20 demos per task.

What they do

Robot datasets cover few tasks and scenes, so R3M pretrains a visual representation on human video instead: Ego4D, over 3500 hours from more than 70 locations, each clip paired with a language description. The encoder is then frozen and reused for policy learning in unseen environments.

How it works

Time-contrastive loss. An InfoNCE loss makes frames closer in time closer in embedding space than frames farther apart or from other videos.
Video-language alignment. A head takes the embeddings of an initial and a later frame plus a DistilBERT sentence embedding and scores whether the transition completes the description. The score should rise over the clip and be higher for matching video-language pairs.
Regularization. L1 and L2 penalties keep the ResNet50 embedding sparse and compact.
Downstream use. The frozen embedding is concatenated with proprioception and fed to a two-layer MLP trained by behavior cloning.

Results

Across 12 tasks in MetaWorld, Franka Kitchen and Adroit, R3M reaches about 62% success, over 20% above training from scratch and over 10% above CLIP, MoCo (345) and supervised ImageNet. It is best on 11 of 12 tasks.
Removing language alignment drops average success from 62.4% to 53.2% (Adroit: 65.0% to 45.6%). Removing L1 or crop augmentation costs about 2%.
MoCo trained on the same Ego4D frames gets 42.0% on Franka Kitchen versus R3M’s 53.1%, so the objective adds about 10% beyond the data.
On a Franka in an apartment with 20 demos per task, R3M averages 56% over five tasks versus 24% for CLIP (lettuce in pan 60% vs 0%, towel folding 40% vs 0%).

Limitations

Evaluation covers only behavior cloning from a few demonstrations; whether the representation helps reinforcement learning is untested. R3M gives a single-frame state representation and is not used for rewards or task specification. Simulation success rates are the best checkpoint over training.

Figures

From Figure 1: R3M is pretrained on Ego4D human video with time-contrastive learning, video-language alignment and an... — From Figure 1: R3M is pretrained on Ego4D human video with time-contrastive learning, video-language alignment and an L1 sparsity penalty, then reused for robot learning in new environments and tasks.

From Figure 2: Ego4D frames with their language annotations (left) and the training objective (right): frames close i... — From Figure 2: Ego4D frames with their language annotations (left) and the training objective (right): frames close in time are pulled together, and a language head scores whether a clip completes its description.

From Figure 4: Imitation-learning success rates with standard error on Franka Kitchen, MetaWorld and Adroit, where R3... — From Figure 4: Imitation-learning success rates with standard error on Franka Kitchen, MetaWorld and Adroit, where R3M beats MoCo (345), CLIP, supervised ImageNet features and training from scratch.

From Figure 5: Real-robot rollouts learned from 20 demonstrations: putting lettuce in a pan, pushing a mug to a goal ... — From Figure 5: Real-robot rollouts learned from 20 demonstrations: putting lettuce in a pan, pushing a mug to a goal and folding a towel.

Abstract

We study how visual representations pre-trained on diverse human video data can enable data-efficient learning of downstream robotic manipulation tasks. Concretely, we pre-train a visual representation using the Ego4D human video dataset using a combination of time-contrastive learning, video-language alignment, and an L1 penalty to encourage sparse and compact representations. The resulting representation, R3M, can be used as a frozen perception module for downstream policy learning. Across a suite of 12 simulated robot manipulation tasks, we find that R3M improves task success by over 20% compared to training from scratch and by over 10% compared to state-of-the-art visual representations like CLIP and MoCo. Furthermore, R3M enables a Franka Emika Panda arm to learn a range of manipulation tasks in a real, cluttered apartment given just 20 demonstrations. Code and pre-trained models are available at https://tinyurl.com/robotr3m.

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