Language-Driven Representation Learning for Robotics

In plain words

Robots need to see the world in a useful way before they can learn a task from a few examples. One way to get that is to pretrain a vision model on lots of short videos of people doing everyday things. This paper's model, Voltron, learns from those videos together with the caption that describes each clip. It hides parts of the video frames and tries to fill them back in, and for the same clip it either reads the caption as a hint or tries to write the caption itself.

Those two choices pull the model in different directions. Filling in hidden pixels pushes it to notice fine detail such as edges and colours. Writing the caption pushes it to notice what is happening, which is a higher-level idea. The mix between the two is a single setting, so you can lean one way or the other. The authors built a set of five robot tasks to test the result. They report that the learned features did better than two earlier approaches on the four tasks with numbers, with the biggest gains on tasks where language matters.

TL;DR

• Voltron is a masked autoencoder over video frames that also uses the paired caption, either as a condition for pixel reconstruction or as a target for caption generation §1 §3.
• A single probability alpha decides, per example, between conditioning and generating. The three reported variants use alpha 0 (VCond, VDual) and 0.5 (VGen) §3.2.
• Pretraining uses Something-Something-v2 (168K clips, 400 epochs, batch 1024) with ViT-Small/16 as the main backbone §4.
• The authors evaluate frozen features on five robot problems and report that Voltron beats their controlled reimplementations of MVP and R3M on the four quantitative ones. The fifth, intent scoring, is qualitative only. All numbers are author-reported §5 tab. 2 tab. 3 fig. 7.
• Official code is an MIT-licensed PyTorch repository by the first author. The code fixes a 0.75 mask ratio and a 0.5 caption-loss weight, neither of which the paper text states models.py.

Problem

Two families of visual pretraining for robotics had been reported to suit different tasks. Masked-autoencoding approaches keep fine spatial detail but lose high-level meaning, whereas contrastive approaches show the reverse pattern §1 [abstract]. Video data of humans is used because robot data is scarce and expensive §1.

Method

Overview. Voltron is trained on video clips with a caption. Each training example is a pair of frames from one clip plus the caption. The model masks most image patches, encodes what remains together with a language input, and decodes two outputs: the masked pixels, and optionally the caption §3 §3.1. The caption branch is switched on or off per example, and the switch decides how strongly the representation is pushed toward high-level features §3.2.

Steps.

1. Data. Each example is a video v of T frames and a caption c. The caption is tokenised, padded to a maximum length L (20 tokens in the reported setup), and paired with a <NULL> placeholder for empty context §3 app. B.2.
2. Frame sampling. For each clip, the first frame comes from the first 20% of the clip and the second from the remaining 80%. The reported dual-frame contexts use one such pair §3.2 app. D.1. The pseudocode below shows this rule.
3. Masking. Each frame is split into non-overlapping p x p patches, with p = 16 §4. A random mask hides a fixed fraction gamma of patches. The same mask is applied to every frame in the clip, so the encoder cannot read the hidden patches from the other frame §3.
4. Encoder. The encoder is a Transformer over the visible patches of all frames, the frozen DistilBERT embedding of the caption projected into the same width, and learned modality and frame-index embeddings §3.1 app. B.2. The visible length is S = (1 - gamma) k |R|, where k is the number of frames and |R| the number of patches per frame §3.1.
5. Pixel reconstruction. A decoder fills in the masked patches from the encoder output, conditioned on the caption §3.1. The targets are the pixels of each patch normalised to zero mean and unit variance, and the loss is the mean squared error on masked patches only §3.1.
6. Caption generation. For generation examples, the same decoder predicts the caption token by token. A causal prefix mask lets each caption token see the visual context and earlier caption tokens only §4 app. B.2. The encoder receives only <NULL> in this case, which prevents the encoder from copying the caption §3.1.
7. Objective. The paper writes the loss as the reconstruction term plus the caption term. Per example, with probability alpha the example is a generation example (z = 1) and otherwise a conditioning example (z = 0) §3.2. The code adds a fixed weight of 0.5 on the caption term models.py vgen.py.

# Sample a dual-frame context from one clip  [§3.2]
# Input:  frame_count (int)  # number of frames in the clip
# Output: (first, second) frame indices

def sample_dual_frames(frame_count, alpha=0.2):
    # First frame from the first 20% of the clip  [§3.2]
    # [code: voltron/preprocessing/core.py#L181]
    first = random_int(0, round(alpha * frame_count))
    # Second frame from the remaining 80% of the clip  [§3.2]
    # [code: voltron/preprocessing/core.py#L184]
    second = random_int(round((1 - alpha) * frame_count), frame_count)
    return first, second

# One Voltron pretraining step, VGen variant (alpha = 0.5)  [§3] [§3.2]
# Input:  frames (B, 2, 3, 224, 224)  # frames 0 and K of each clip
#         cond_tokens (B, 20)  # caption ids, or <NULL> if generating
#         gen_tokens (B, 20)  # caption ids to generate
#         gen_flags (B,)  # 1 = generate this example, 0 = condition
# Output: scalar loss

def training_step(frames, cond_tokens, gen_tokens, gen_flags):
    # Mask 75% of the patches, same mask for both frames  [§3]
    # [code: voltron/models/core/vgen.py#L260]
    visible, mask, restore = mask_patches(frames, mask_ratio=0.75)
    # Encode the visible patches jointly with the caption condition  [§3.1]
    hidden = encoder(visible, cond_tokens)
    # Shared decoder predicts masked pixels for both frames  [§3.1]
    pixel_pred = decoder_pixels(hidden, restore)
    # Shared decoder predicts caption tokens with a causal mask  [§4]
    token_logits = decoder_language(hidden, gen_tokens)
    # Masked MSE on per-patch normalised pixel targets  [§3.1]
    # [code: voltron/models/core/vgen.py#L622-L639]
    pixel_loss = masked_mse(pixel_pred, normalize_patches(frames), mask)
    # Mean token negative log-likelihood per caption
    # [code: voltron/models/core/vgen.py#L650]
    caption_nll = mean_token_nll(token_logits, gen_tokens)
    # Only generation examples contribute to the caption term
    # [code: voltron/datasets/datasets.py#L273]
    caption_loss = masked_mean(caption_nll, gen_flags)
    # Weighted sum; the paper gives no weight, the code uses 0.5
    # [code: voltron/conf/models.py#L405]
    return pixel_loss + 0.5 * caption_loss

Variants. The three reported models differ only in alpha and the number of frames in the context §3.2:

• VCond: alpha = 0, one frame, always conditioned on the caption.
• VDual: alpha = 0, two frames, always conditioned on the caption.
• VGen: alpha = 0.5, two frames, conditioned or generated with equal probability.

The authors did not evaluate alpha = 1, because some conditioning was always helpful in their preliminary runs §3.2. They limit contexts to k = 2 because four-frame contexts exceeded the memory of their hardware §3.2.

Evaluation protocol. Pretrained encoders are frozen. Each task then trains a small head on the extracted features. Features are extracted with multi-headed attention pooling (MAP) over the patch tokens rather than the MVP <CLS> token, which the authors found to be better for every ViT model they tested §4 app. D.2. The adaptation heads and optimisers are listed in the Hyperparameters table app. E.

Hyperparameters

Values marked "code" are not stated in the paper text. They come from the official repository at code_commit, which the card's ## Paper vs code section compares with the paper.

Results

Author-reported. Pretraining on Something-Something-v2 with the controlled reproductions. Baselines are R-MVP and R-R3M, reimplemented on the same frames §4. Gray rows in the paper use models pretrained on Ego4D, which are not directly comparable §4.

Grasp affordance prediction (ARC grasping, average precision, percent) tab. 2 §5.1:

Referring expression grounding (OCID-Ref, average precision at 0.25 IoU, percent) tab. 3 §5.2:

Ablations on three tasks tab. 4 §6. The grasp column reports the top-1% score, the refer column total accuracy, and the imitate column the success rate with 25 demonstrations (mean and standard deviation, as printed):

Single-task control (Franka Kitchen, 5 tasks, 2 viewpoints, 3 seeds) fig. 5 §5.3. The figure plots success rate against 5, 10 and 25 demonstrations. The authors report that most approaches perform similarly. VDual and VGen beat the single-frame models. The absolute success rates stay low.

Language-conditioned imitation on a real robot (five tasks, 20 demos each) fig. 6 §5.4. Success is scored with partial credit: 0.25 for each milestone reached (reach, interact, transport, complete). Voltron models score highest across the board. Under visual distractors, Voltron and R-MVP stay robust to the most extreme distractors §6.

Zero-shot intent scoring (qualitative) fig. 7 app. C §5.5. VGen's caption likelihood tracks the progress of "opening the faucet" in WHiRL videos of humans and robots. CLIP and R3M produce noisy scores. No numbers are reported.

Paper vs code

Code: https://github.com/siddk/voltron-robotics at commit 1b299bf5cfa06673a3738aa6e15423b92a9922cd, MIT license. The owner is the first author (the LICENSE names Siddharth Karamcheti), and the paper links the repository §4 app. B. Paths below are relative to the repository root at that commit.

• Mask ratio. The paper gives a fixed ratio gamma without a value §3.1. The code sets 0.75 as the default for every Voltron model models.py models.py models.py models.py.
• Caption loss weight. The paper writes the objective as an unweighted sum of the pixel and caption terms §3.2. The code weights the caption term by 0.5 and documents this as a rough guide borrowed from M3AE models.py vgen.py.
• Choice of generation examples. The paper samples z ~ Bernoulli(alpha) per example §3.2. The code draws exactly int(gen_ratio x N) generation examples per epoch, without replacement datasets.py. The expected fraction is the same, but the per-epoch count is fixed.
• Null language. The paper uses a separate <NULL> token distinct from padding §3. The code keeps only the first ([CLS]) token of the caption and zeroes the rest for conditioning datasets.py datasets.py. In scoring, the [CLS] id 101 is hard-coded vgen.py.
• Intent score. The paper describes the score as a likelihood of the caption under the generator §5.5. The code returns the negative mean per-token negative log-likelihood, which normalises by caption length vgen.py.
• Pixel loss over frames. The code averages the masked MSE of frame 0 and frame K equally vgen.py. The paper does not say how the two frames are combined §3.1.
• Learning rate and schedule. The paper gives neither a base rate nor a warmup. The code uses base 1.5e-4, scaled by batch over 256, with 20 warmup epochs and cosine decay over 400 epochs vgen.py datasets.py optimization.py.
• Evaluation code. The MAP feature extractor and the five evaluation harnesses are in a separate repository linked from the paper app. E. They were not inspected for this card.

Checked and consistent with the paper: the mask shared across frames vgen.py, the reconstruction target normalisation, the RMSNorm, SwishGLU and LayerScale blocks vgen.py, and the dual-frame sampling window core.py. Not checked: the separate evaluation repository.

Pitfalls

• Mask ratio and caption weight are not in the paper. An implementation that follows only the text will have no value for gamma and no caption weight. Use 0.75 and 0.5 from the code, or state the choice explicitly.
• Fixed generation count. With the code's sampler, each epoch has exactly half generation examples at alpha 0.5, and the same examples are not repeated. An implementation that uses an independent coin flip per example will be noisier, though the mean is the same.
• Caption normalisation. The caption loss is averaged per caption, then averaged over generation examples. Long captions do not dominate the loss. A token-level average would weight them more.
• Empty generation batches. If a batch has no generation examples, the code divides by eps and the caption term goes to zero vgen.py.
• Frame-averaged pixel loss. The two frames contribute equally, regardless of how many patches each has, and only masked patches are scored vgen.py.
• Feature extraction. The gains from MAP extraction are large. The authors report that switching the MVP model to MAP almost doubled its success rate on control tasks app. D.2. Compare baselines with the same extractor.
• Pretraining cost. Pretraining used 400 epochs on TPU v3-8 with batch 1024. The authors note they could not train models at the scale of the released MVP and R3M artifacts §1 §4.
• Baselines are reimplementations. R-MVP and R-R3M are the authors' reimplementations on Something-Something-v2. Released MVP and R3M models use Ego4D and are reported only as dashed reference lines §4.
• Real-robot scores are partial credit. A 0.25 milestone score is not comparable to a binary success rate §5.4.
• Control results are low in absolute terms. The authors say learning closed-loop policies from few demonstrations remains open §5.3.
• Language encoder is frozen. DistilBERT is not fine-tuned. Captions are truncated to 20 tokens app. B.2.

Lineage

• Builds on: masked autoencoding §4, with MVP as the masked-visual baseline §1, and R3M, which uses time-contrastive and video-language alignment objectives §1 §4. No arXiv ids are given in the paper text, so none are linked here.
• Compared against: R-MVP and R-R3M reimplementations, the released MVP (EgoSoup) and R3M (Ego4D) models, and CLIP ViT-B §4.
• Related ablation: multimodal masked autoencoders (M3AE), reproduced as R-M3AE, where language modelling loss fell to zero early and training overfit §6 app. C.
• Superseded by: not stated in the paper or its code.