VIP: Towards Universal Visual Reward and Representation via Value-Implicit Pre-Training

Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, Amy Zhang

Seminal September 2022

TL;DR VIP pre-trains a ResNet50 on Ego4D human videos with a goal-conditioned value objective that needs no actions, so the embedding distance to a goal image can serve as a dense reward for new robot tasks. With the frozen model as the reward, MPPI solves about 30% of FrankaKitchen tasks (44% with a larger budget), and few-shot offline RL from about 20 real demos reaches 60-100% success.

What they do

Pre-trained visual encoders (R3M, MoCo, CLIP, ImageNet ResNet) work well as policy inputs. But the distance to a goal image in their embedding space makes a poor reward. VIP treats Ego4D human videos as data for offline goal-conditioned RL and learns a value function without action labels. The negative embedding distance to a goal image then serves as a dense reward for unseen robot tasks, with no fine-tuning.

How it works

Action-free value objective. The Fenchel dual of a KL-regularized offline goal-conditioned RL objective contains no actions. The reward is -1 at every frame except the goal, so the value estimates the discounted number of steps to the goal.
Implicit time contrast. The value is the negative L2 distance between embeddings. The loss pulls a clip’s first and last frames together, and middle frames are tied only by one-step TD errors, so frames farther from the goal end up farther away and distance-to-goal curves are smooth.
Training. ResNet50 on about 72,000 Ego4D clips (4.3M frames), with R3M’s preprocessing and hyperparameters.
Use. Frozen encoder; the reward is the change in negative embedding distance to the goal image(s).

Results

FrankaKitchen (12 tasks x 3 views): with MPPI, VIP solves about 30% of tasks and is the only representation that makes non-trivial progress in the Hard setting. It reaches 44% with 256 trajectories per step, while baselines often get worse with more budget.
Online RL with NPG: 40% aggregate success. VIP’s encoder with the ground-truth sparse reward solves no tasks.
Real Franka, about 20 demos per task: VIP-RWR gets 100/90/60/90% on CloseDrawer/PushBottle/PlaceMelon/FoldTowel, vs 50/50/10/20% for VIP-BC and 80/70/0/0% for R3M-RWR. VIP trained only on the task data scores 0%.
Fraction of “bumps” (steps where distance to goal rises) on robot videos: 0.243 for VIP vs 0.323 for R3M.

Limitations

Tasks must be specified by a goal image. The symmetric distance assumes reversible environments. The model was only tested frozen, and only on manipulation. As a plain encoder for imitation learning it is not better than R3M (53.6 vs 55.9% success).

Figures

From Figure 1: VIP pre-trains on human videos by pulling the initial and goal frames of a clip together, while middle... — From Figure 1: VIP pre-trains on human videos by pulling the initial and goal frames of a clip together, while middle frames are linked only through one-step temporal differences. The frozen network then supplies reward and representation for downstream robot control.

From Figure 2: 2D embeddings of a held-out FrankaKitchen demonstration learned in-domain by VIP and by standard time ... — From Figure 2: 2D embeddings of a held-out FrankaKitchen demonstration learned in-domain by VIP and by standard time contrastive learning (TCN), colored from start (purple) to end (red). Insets show the embedding distance to the last frame.

From Figure 4: Cumulative success rate on FrankaKitchen when each pre-trained representation provides the reward for ... — From Figure 4: Cumulative success rate on FrankaKitchen when each pre-trained representation provides the reward for MPPI trajectory optimization and NPG online RL, in the Easy and Hard settings.

From Figure 6: Normalized embedding distance to the goal frame along a real-robot video for each representation.

Abstract

Reward and representation learning are two long-standing challenges for learning an expanding set of robot manipulation skills from sensory observations. Given the inherent cost and scarcity of in-domain, task-specific robot data, learning from large, diverse, offline human videos has emerged as a promising path towards acquiring a generally useful visual representation for control; however, how these human videos can be used for general-purpose reward learning remains an open question. We introduce $\textbf{V}$alue-$\textbf{I}$mplicit $\textbf{P}$re-training (VIP), a self-supervised pre-trained visual representation capable of generating dense and smooth reward functions for unseen robotic tasks. VIP casts representation learning from human videos as an offline goal-conditioned reinforcement learning problem and derives a self-supervised dual goal-conditioned value-function objective that does not depend on actions, enabling pre-training on unlabeled human videos. Theoretically, VIP can be understood as a novel implicit time contrastive objective that generates a temporally smooth embedding, enabling the value function to be implicitly defined via the embedding distance, which can then be used to construct the reward for any goal-image specified downstream task. Trained on large-scale Ego4D human videos and without any fine-tuning on in-domain, task-specific data, VIP's frozen representation can provide dense visual reward for an extensive set of simulated and $\textbf{real-robot}$ tasks, enabling diverse reward-based visual control methods and significantly outperforming all prior pre-trained representations. Notably, VIP can enable simple, $\textbf{few-shot}$ offline RL on a suite of real-world robot tasks with as few as 20 trajectories.

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