The Surprising Effectiveness of Representation Learning for Visual Imitation

In plain words

Teaching a robot a task by showing it examples is hard when the examples look a little different each time. The usual approach trains one large network to do two jobs at once: work out what the camera image means, and work out what to do about it. Those jobs get tangled, so the network needs a lot of examples.

This paper pulls the two jobs apart. First, a picture model learns to describe camera images well, using the robot's own recorded frames and no labels. Second, the robot does not train a policy at all. When it sees a new image, it finds the most similar frames from its demonstrations and averages what the human did in those frames, giving more weight to the closer matches. The authors call this VINN.

The surprise is how well this simple lookup works. It opens a real cabinet door more often than a standard end-to-end imitation network trained on the same demonstrations, and it holds up better when the cabinet is partly covered.

TL;DR

VINN is a visual imitation method with two parts: a BYOL-fine-tuned ResNet-50 encoder, and a k-nearest-neighbour, softmin-weighted average of demonstration actions. It needs no policy training. On a 71-demonstration door-opening dataset and a real Hello Robot Stretch, it outperforms the paper's end-to-end and frozen-representation behavior-cloning baselines tab. I. It is simple to implement, but its nearest-neighbour search is linear in the number of demonstration frames §4.8.

Problem

Visual imitation methods usually train one model end to end on images and actions. This couples representation learning and behavior learning, so the model needs a large number of demonstrations to generalise §1. Generalising across environments and doors requires either many demonstrations, task-specific priors, or large parametric models §1.

The authors argue that self-supervised representation learning, applied to the robot's own offline frames, can supply the representation separately §1 §3.1. The action side is then a non-parametric lookup, which needs no training beyond the encoder §1.

Assumptions: the demonstrations and the offline frames come from the same task setting. The paper's self-supervised step uses only task-specific data, and task-agnostic pre-training is listed as future work §3.1 §5.

Method

Section numbers are the arabic form of the paper's IEEEtran numbering, so the paper's section III-B is cited §3.2. Appendix subsections are cited by one word: app. Architectures is "Network Architectures and Training Details", app. Demonstration is "Demonstration Collection Details", and app. Pseudocode is "VINN Pytorch Pseudocode".

VINN has two stages.

1. Encoder. Initialise a ResNet-50 (final linear layer removed, 2048-d output) from ImageNet weights. Fine-tune it with BYOL on the frames of the demonstration dataset, without action labels. BYOL uses an online network and a slowly moving average target network §3.1 §4.3 app. Architectures. The paper also reports that other self-supervised objectives can be swapped in §3.1.
2. Action prediction. Embed all demonstration frames once to get the set E. At test time, embed the query, find its k nearest demonstration embeddings under Euclidean distance, and return the weighted average of their actions. The weights are a softmin over distances, normalised over the k neighbours §3.2. The kernel in the paper is written as exp(-‖e − e^(i)‖₂) divided by the sum of the same terms over the k neighbours. The paper gives no bandwidth or temperature term §3.2 (display equation, unnumbered).

For the robot, the closed loop queries the model at every step. The model returns a translation action and a gripper state. The robot moves by c ⊙ â, where each element of c is below 1 to damp errors from the structure-from-motion labels and to ease transfer from human demonstrations §3.3.

Demonstrations are recorded with a reacher-grabber carrying a GoPro. Frames are linearised, sampled at one per second, and their camera poses are reconstructed with OpenSfM. Translations between consecutive frames become actions app. Demonstration.

Training pipeline, in the paper's order:

# Stage 1: fine-tune the encoder with BYOL, no action labels  [§3.1] [§4.3]
# Input:  demo_frames (N, H, W, 3); ImageNet-pretrained ResNet-50
# Output: encoder mapping an image to a 2048-d embedding

def train_encoder(demo_frames, imagenet_weights):
    # Initialise online and target networks from ImageNet  [§3.1]
    online = resnet50(imagenet_weights)
    target = copy(online)
    for epoch in range(100):  # 100 epochs  [§4.3] [app. Architectures]
        for frame in demo_frames:
            # Two augmented views: crop, jitter, grayscale, blur  [§4.3]
            view_a = augment(frame)
            view_b = augment(frame)
            # BYOL loss between online and target outputs  [§3.1]
            loss = byol_loss(online(view_a), target(view_b))
            # Adam step on the online network, lr 3e-4  [app. Architectures]
            online = adam_step(online, loss, learning_rate=3e-4)
            # Target is a slow moving average of the online network  [§3.1]
            target = moving_average(target, online)
    return online

Inference, with the paper's distance and weighting:

# Stage 2: non-parametric action prediction  [§3.2] [app. Pseudocode]
# Input:  query_image; demo_embeddings (N, 2048); demo_actions (N, 3)
# Output: predicted action (3,)

def predict_action(query_image, encoder, demo_embeddings, demo_actions, k):
    # Embed the observation with the fine-tuned encoder  [§3.1]
    query = encoder(query_image)                       # (2048,)
    # Euclidean (L2) distance to every demonstration embedding  [§3.2]
    distances = norm(demo_embeddings - query, axis=1)  # (N,)
    # Keep the k closest demonstration frames; k around 10  [§3.2] [§4.7]
    nearest = argsort(distances)[:k]                   # (k,)
    # SoftMin weights exp(-distance), normalised over the k neighbours  [§3.2]
    weights = softmax(-distances[nearest])             # (k,)
    # Weighted average of their actions  [§3.2]
    return weights @ demo_actions[nearest]             # (3,)

The gripper is a separate output. For behavior cloning, a linear layer predicts four classes (open, almost open, almost closed, closed) app. Architectures §3.3. For the nearest-neighbour methods, a floating value g is mapped to a gripper state with hyperparameters tuned per experiment §3.3. The paper does not say how the gripper is predicted for VINN itself, and it does not give the mapping values.

Hyperparameters

Results

Author-reported. The paper evaluates on three offline datasets (Pushing, Stacking, Door Opening) by test MSE, and on a real Hello Robot Stretch door-opening task with three cabinets §4.1.

Real robot, door opening tab. I. Success over 30 trials (10 per cabinet on three cabinets), closed loop.

Generalisation to visual modifications tab. II. Success over 10 trials on one cabinet with occlusions.

Offline test MSE (×10⁻¹) tab. III. Standard deviations are reported only for BC-Rep (three random seeds).

The low-data comparison on Pushing, Stacking and Door Opening (MSE against training-set size) is shown only as a plot fig. 4. The text says end-to-end BC catches up with large training sets, and that nearest-neighbour methods are competitive or better with few demonstrations, with VINN clearly better on Stacking and Door Opening below 20 demonstrations §4.4. The exact values are not tabulated.

Paper vs code

The authors' project page links github.com/jyopari/VINN [abstract]. We cloned it at commit fb801ee3e69a708f7fe10c4f3d9d16381228ed02. It holds only the project website (index.html, CSS, demo videos and gifs). It has no training or inference code, no configs, and no LICENSE file. So there is no official implementation to compare against, and this card records no code field. Every hyperparameter above comes from the paper and its appendix, and none of it was checked against code.

The appendix pseudocode app. Pseudocode is a sketch, not runnable code. See Pitfalls.

Pitfalls

• Kernel scale is unspecified. The weights are exp(-distance) with no temperature. Raw distances on 2048-d embeddings can be large, so the softmin can concentrate on one neighbour. Implement it with the usual max-subtraction for stability. The paper does not say whether embeddings are L2-normalised before the distance, and that choice changes the scale §3.2.
• The appendix pseudocode does not run as written. It never initialises dist_list, uses an undefined dataset_path, calls an unimported softmax, has a mis-indented return, and hard-codes a 3-d action app. Pseudocode. Use it only to confirm shapes and the sort-then-weight order.
• Gripper handling is under-specified. The mapping from the float g to a gripper state is tuned per experiment, and the values are not given. The paper does not state how gripper state is produced for VINN §3.3.
• Action scaling and normalisation are not reproducible from the paper. The scale vector c, and the normalisation of SfM-derived translations, are not given §3.3 §4.1.
• Introduction numbers differ from the tables. The introduction reports 40% success on novel scenes and 3.3% for the strongest baseline §1. Those values match the average of the three occluded rows of Table II for VINN (70, 50, 0) and BC-rep (10, 0, 0). The paper does not say that it averaged them, so treat the mapping as our reading tab. II.
• BC-rep is not beaten in the unmodified scene. BC-rep scores 90% against VINN's 80% in the baseline row tab. II. The gain shows up under visual change.
• Offline MSE is not a reliable proxy. The paper says low MSE is necessary but not sufficient for real-robot success §4.5. Table III shows VICReg better than VINN on Pushing, and the authors chose BYOL for simplicity §4.7.
• Lookup cost grows with the demonstrations. Each query is compared against every demonstration embedding, so cost is linear in the number of frames. The paper uses brute-force search and names FAISS as a faster option, without testing it §4.8.
• Timing figures lack hardware detail. The encode and lookup times §4.8 do not say which hardware produced them.
• Robot results are small-sample. 30 rollouts across three cabinets for Table I and 10 trials per row for Table II. The paper reports no confidence intervals.
• Encoder data is task-specific. The paper uses only demonstration frames for self-supervision and lists task-agnostic pre-training as future work §3.1 §5.
• Appendix pipeline details are only loosely specified. The SfM and gripper-extraction steps are described in prose, with no version, parameter or code app. Demonstration.

Lineage

• Builds on: BYOL for the self-supervised encoder (arXiv:2006.07733, cited as [byol] in the paper §3.1); locally weighted regression for the action rule §2.3.
• Related to: Playful Interactions for Representation Learning (arXiv:2107.09046), whose BC-on-representations baseline the authors say is similar to their BC-rep §4.2 §5.
• Compared against: end-to-end behavior cloning §4.2; Implicit BC (arXiv:2109.00137) §4.7; ImageNet-feature and frozen-representation variants §4.2.
• Superseded: not known from content in this wiki.