Masked Visual Pre-training for Motor Control

In plain words

Teaching a robot arm to act from camera images with trial and error takes a huge number of tries, and the learned behaviour often fails when the scene changes. The authors ask whether the robot's "eyes" can be trained first, on ordinary photos and videos with no labels, and then frozen. The robot's decision-making part is then the only thing that learns the task.

The pre-training works like a fill-in-the-blank game: most of each image is hidden, and the network learns to guess what is missing. To get good at this, it has to learn the shape of objects and how hands handle them. The picture data comes from videos of people handling objects, many of them from a first-person camera view.

Once the eyes are frozen, a small controller reads the image features plus the robot's joint readings and decides how to move. The same frozen eyes serve many different tasks, such as reaching, picking up, moving and cabinet tasks. The authors also release a set of simulated pixel-based tasks, so others can compare against this approach.

TL;DR

A ViT-S/16 masked autoencoder is pre-trained on about 700k hand-object images and then frozen. A small PPO controller reads its features and the robot's joint state, and it learns eight pixel-based Franka and Kuka tasks in a GPU simulator. The authors report that this beats a supervised ImageNet encoder on 7 of 8 tasks and often matches an oracle that sees the true object state. These results are author-reported and shown in figures, and the pre-training code is not in the repository.

Problem

Learning motor control end-to-end from pixels with reinforcement learning has high sample complexity, especially with high-dimensional inputs and actions. The learned features also tend to overfit to the scenes and objects seen during training §1 §2.2.

Auxiliary self-supervised objectives on environment data help, but the representations are still learned only from the task's own experience §1 §6. The authors want representations learned from large collections of natural images, used without fine-tuning on each task §1.

For evaluation, the authors say existing motor-control benchmarks lack at least one of: high-resolution images, realistic robots, fast physics, efficient training, and suitable reward functions and metrics §3.1 tab. 1.

The method assumes a wrist-mounted camera by default and proprioceptive readings (joint positions and velocities), both available on real robot hardware §3.4 §2.2.

Method

The paper contains no numbered equations and no algorithm boxes §2. The pseudocode below is written from the text and the official code, and it is not a verbatim copy of either.

Stage 1: masked visual pre-training §2.1.

1. Build the Human-Object Interaction (HOI) dataset: about 700k images combined from Epic-Kitchens, Something-Something and 100 Days of Hands, with frames sampled at 1 fps and 0.3 fps §2.1 §4. The paper does not use any human labels or temporal information §2.1. An ImageNet variant is used for controlled comparisons §2.1.
2. Train a masked autoencoder (MAE) with a 75% masking ratio. The encoder sees only the unmasked patches, and a lighter decoder reconstructs the pixels §2.1.
3. Use a ViT-S/16 encoder, with an auxiliary classification token used for downstream use §4.

Stage 2: control with a frozen encoder §2.2.

4. Freeze the whole visual encoder for the entire RL training run. No task-specific fine-tuning is done, and the same encoder is used for every task §2.2 §4.
5. Turn the image into a 128-d vector with a trainable linear layer over the classification-token feature, and turn the proprioceptive state into a 128-d vector with a second linear layer. Concatenate them into a 256-d input §4. The code also applies the ViT's final LayerNorm first [code: mvp/ppo/actor_critic.py#L150-L155] vit.py.
6. The policy is a four-layer SELU MLP with hidden sizes 256, 128 and 64. It outputs joint-angle deltas under position control §3.4 §4. The critic has the same architecture with its own weights §2.2.
7. Train the policy and critic with PPO, using clipped surrogate and value losses §2.2 §4 ppo.py.

# Pixel observation to joint embedding  [§2.2] [§4]
# Input:  image (3, 224, 224), normalised; state (2 * n_joints,)
# Output: joint embedding (256,)

def embed_observation(image, state):
    # Frozen ViT-S/16; take the classification-token feature (384-d) [§4]
    cls_feature = frozen_vit_cls_token(image)            # (384,)
    # Final LayerNorm, then a trainable linear layer to 128-d [§4]
    image_emb = image_projection(layer_norm(cls_feature))  # (128,)
    # Joint positions and velocities, embedded to 128-d [§2.2] [§4]
    state_emb = state_projection(state)                  # (128,)
    return concatenate(image_emb, state_emb)             # (256,)

# One PPO training iteration  [§2.2] [§4]
# Collect 32 steps from each of 256 environments per iteration [§4]

def ppo_iteration(envs, lr_init, iteration, max_iter):
    # Roll out the current policy in parallel environments [§4]
    rollouts = collect(envs, steps_per_env=32)
    # Generalised advantage estimation; gamma and lambda are not stated in
    # the paper, the config uses 0.99 and 0.95 [code: configs/ppo/train/
    # FrankaPickPixels.yaml#L31-L32]
    advantages, returns = gae(rollouts, gamma=0.99, lam=0.95)
    # Normalise advantages over the whole batch [code: mvp/ppo/storage.py#L85]
    advantages = normalise(advantages)
    for epoch in range(10):                              # 10 epochs [§4]
        # Cosine decay of the learning rate, evaluated per epoch [§4]
        # [code: mvp/ppo/ppo.py#L397-L401]
        lr = lr_init * 0.5 * (1 + cos(pi * (iteration + epoch / 10) / max_iter))
        for minibatch in split(rollouts, num_minibatches=4):   # [§4]
            # Clipped surrogate policy loss; clip range 0.1 [code: #L23]
            ratio = exp(new_log_prob - old_log_prob)
            loss_pi = -min(adv * ratio, adv * clamp(ratio, 0.9, 1.1))
            # Value loss; weight 2.0 [code: hydra_utils.py#L197]
            loss_v = (value - returns) ** 2
            # Gradient-norm clip at 1 [§4]; Adam, beta 0.9 and 0.999 [§4]
            adam_step(mean(loss_pi + 2.0 * loss_v), grad_clip=1.0, lr=lr)

Benchmark: PixMC §3.

8. The benchmark has eight tasks in NVIDIA IsaacGym, using a Franka arm with a parallel gripper or a Kuka arm with a 4-finger Allegro hand §3.2 §3.3 §5.1. The tasks are hand-designed, with dense rewards and reward-independent success metrics §3.5. Control runs at 60 Hz §3.4.
9. tab. 1 compares PixMC with RLBench, Robosuite and MetaWorld. PixMC is the only one listed with IsaacGym, and it is the only one with both a fast simulator and multi-finger hands tab. 1 §3.1.
10. Training uses 256 environments per GPU at 224 by 224 pixels, with gradients synchronised across GPUs §3.6.

Hyperparameters

Learning rate per task in the released configs (the same value is used for the supervised baseline, per the config comments): 0.001 for Franka reach, pick, cabinet and pick-object; 0.0015 for Kuka reach and Franka move; 0.0005 for Kuka pick, move, cabinet and pick-object *Pixels.yaml.

Results

Author-reported. Setup: IsaacGym simulation of the PixMC benchmark, pixel observations, ViT-S/16 encoder frozen during RL, PPO. The baseline is a supervised ViT-S trained on ImageNet §4 §5. The paper states its headline results in text and shows per-task numbers only in figures, so the per-task values below are described, not quoted §5.

Paper vs code

Official code: github.com/ir413/mvp at code_commit. The project page linked from the paper points to it §1 fig. 1. The repository holds PPO training, behaviour cloning, the PixMC tasks and configs. It does not hold the MAE pre-training code. Its pre-trained weights are downloaded from a Box link model_zoo.py. The repository has no LICENSE file, so code_license is not stated, and no code is quoted verbatim.

• Learning rate. Paper: a grid over three rates with 5 seeds each, and the best rate is reported §4. Code: one fixed rate per task, from 0.0005 to 0.0015 *Pixels.yaml. Which rate produced each reported curve is not stated.
• Actor and critic weights. Paper: "the policy and the critic do not share weights" §2.2. Code: the MLPs are separate, but the image projection and the proprioception projection feed both heads and are trained by both losses actor_critic.py actor_critic.py actor_critic.py.
• Optimiser scope. Code: Adam over all actor-critic parameters, including the log standard deviation and the two projections ppo.py.
• Value loss. Paper: not stated. Code: weight 2.0 and no clipping, unless a config sets them hydra_utils.py hydra_utils.py. The class defaults differ (weight 1.0, clipped) ppo.py ppo.py.
• Discount and clip range. Paper: not stated. Code configs: gamma 0.99 and clip 0.1. Class defaults: gamma 0.998 and clip 0.2 ppo.py.
• Advantage normalisation. Paper: not stated. Code: normalised over the whole rollout batch storage.py.
• Image feature. Paper: a linear layer projects the classification-token feature to 128-d §4. Code: the ViT's final LayerNorm runs first [code: mvp/backbones/vit.py#L46-L47] actor_critic.py.
• Image normalisation. Paper: not stated. Code: ImageNet mean and std franka_pick.py.
• Task names. The text says "pick and move" in one sentence and "relocate" in the next, for the same 2000-iteration budget §5.1. The code has FrankaMove and KukaMove FrankaMovePixels.yaml. The mapping is inferred, not stated.
• Proprioception. Paper: joint positions and velocities §2.2 §3.4. Code: scaled joint positions and velocities, with velocities multiplied by 0.1 franka_pick.py FrankaPickPixels.yaml.
• Pre-training. Paper: MAE for 1600 epochs on 16 GPUs §4. Code: no MAE pre-training script; only the released ViT-S weights are fetched model_zoo.py. The pre-training recipe cannot be reproduced from this repository.

Pitfalls

• Per-task numbers are in figures. The text gives "up to 80 points" and a few counts (7 of 8, 5 of 8). Read per-task success from fig. 5 before quoting it. Do not reconstruct it from the text.
• Learning rate choice. Final results depend on which of three rates was selected, and the code's fixed per-task rates do not show the search. Start from the code values, then re-check with a small sweep.
• Seed and rate variance. The stability claim rests on one task family (reach) with 15 runs per model §5.3 fig. 10.
• Frozen features can be cached. The code computes the encoder once per step and reuses the stored features during the update [code: mvp/ppo/actor_critic.py#L277-L280]. This saves compute, but it assumes the encoder never changes.
• Scale is not handled. The 6 cm distractor is picked at about 50%, which the authors attribute to scale ambiguity from a single first-person camera §5.2.
• Oracle comparisons. The oracle-state model cannot run on the distractor tests, because its state dimensions change §5.2.
• Compute. Each image-based model takes about 5 hours on one 2080 Ti GPU, for about 8 million environment steps §3.2. The ViT-B and unfrozen-encoder runs needed 8 GPUs instead of 1 to keep the same number of environments §5.3.
• Pre-training data. HOI images come from Epic-Kitchens, Something-Something and 100 Days of Hands §2.1. Each dataset has its own terms, which must be checked before reusing the data.
• Released assets. The repository ships Franka, Kuka and YCB assets with their own licence files licenses. The repository itself has no LICENSE file, so nothing from it may be reused without checking.
• Source version. The text compiled here is the arXiv v1 LaTeX source, in ICML 2022 style. Later versions were not checked.

Lineage

• Builds on: Masked autoencoders for the pre-training objective §2.1; PPO for the controller §2.2; ViT as the encoder §2.1 §4.
• Compared against: CURL and RAD (environment-trained contrastive and augmentation methods) §5.4; CLIP ViT-B §5.4; MoCo-v3 §5.4; supervised ViT-S on ImageNet §4. None of these has a card in this wiki yet.
• Superseded by: not known from content in this wiki.