Source-linked AI summary

The Unsurprising Effectiveness of Pre-Trained Vision Models for Control

Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, Abhinav Gupta

arXiv:2203.03580v2cs.CVcs.AIcs.LGcs.RO

TL;DR

Control still largely trains visuo-motor policies from scratch, leaving the role of pretrained visual representations in control an open question. This paper studies out-of-domain pretrained representations and proposes a multi-layer PVR, finding it competitive with or better than ground-truth features across four evaluation domains.

  • Problem

    Control still largely trains visuo-motor policies from scratch, leaving the role of pretrained visual representations in control insufficiently understood.

    Source: §1 · p. 1
  • Method

    The paper studies out-of-domain pretrained visual representations across control domains, examining training methods, augmentations, and feature hierarchies, and proposes a multi-layer PVR.

    Source: §1 · p. 2, §5 · p. 8
  • Results

    A multi-layer crop-augmented MoCo PVR was competitive with or outperformed ground-truth features across all four evaluation domains.

    Source: §5 · p. 8
  • Takeaways & Limitations

    The authors argue that control research should focus more on visual-input policies using pretrained perception modules than on hand-designed state features.

    Source: §5 · p. 8

    Experiments freeze vision models and exclude on-the-fly representation fine-tuning; full fine-tuning is left for future work.

    Source: §5 · p. 8

Abstract

from arXiv · show

Recent years have seen the emergence of pre-trained representations as a powerful abstraction for AI applications in computer vision, natural language, and speech. However, policy learning for control is still dominated by a tabula-rasa learning paradigm, with visuo-motor policies often trained from scratch using data from deployment environments. In this context, we revisit and study the role of pre-trained visual representations for control, and in particular representations trained on large-scale computer vision datasets. Through extensive empirical evaluation in diverse control domains (Habitat, DeepMind Control, Adroit, Franka Kitchen), we isolate and study the importance of different representation training methods, data augmentations, and feature hierarchies. Overall, we find that pre-trained visual representations can be competitive or even better than ground-truth state representations to train control policies. This is in spite of using only out-of-domain data from standard vision datasets, without any in-domain data from the deployment environments. Source code and more at https://sites.google.com/view/pvr-control.

1. Introduction

The paper asks whether frozen visual representations trained only on out-of-domain data can serve as perception for control, challenging a field dominated by learning visuo-motor policies from scratch.

1. Introduction

  • The study evaluates frozen pre-trained representations as perception modules across Habitat, DeepMind Control, Adroit, and Franka Kitchen.The perception module is trained on out-of-domain data, frozen, and reused while policies are trained on deployment tasks.Source: §1 · p. 2, Figure 2 · p. 2
  • Frozen PVRs trained entirely out of domain can match or outperform ground-truth features for imitation learning, and a single multi-layer PVR does so across all studied domains.The representations saw no evaluation-environment frames during pre-training.Source: §1 · p. 2
  • Self-supervised learning provides better control-policy features than supervised learning.Source: §1 · p. 2
  • For SSL in control, crop augmentations appear more important than color augmentations.This finding is consistent with prior work studying representation learning alongside policy learning.Source: §1 · p. 2
  • Early convolutional features work better for fine-grained MuJoCo control, while later-layer features work better for semantic Habitat ImageNav tasks.Source: §1 · p. 2

2. Related Work

Prior work established transferable visual representations and explored visual control largely with representations learned from deployment data; this study instead evaluates frozen, out-of-domain encoders and reports one representation can work across its four domains.

Representation Learning

  • Pre-trained deep representations transfer across vision, speech, and NLP, improving downstream data efficiency and/or performance (Goyal et al., 2019).In vision, representations can be learned with supervised labels or self-supervision and frozen “off-the-shelf”; prior uses include detection, segmentation, captioning, and action recognition, motivating this control-policy test.Source: §2 · p. 2

Policy Learning

  • Neural RL and IL policies have succeeded across locomotion, manipulation, and indoor navigation, while continuous-control work often uses privileged state or learns visual representations from deployment frames.Visual approaches include data augmentation, contrastive learning, and predictive world models; this work focuses on visuo-motor policy learning with imitation learning.Source: §2 · pp. 2–3

Pre-trained Visual Encoders in Control

  • Compared with earlier studies, this work spans four visually diverse domains, more pre-trained representations, and varied visual invariances, finding a single representation can succeed across all four.Earlier findings were mixed: pre-training in DeepMind Control used data from that domain, CLIP improved navigation results, and ResNet embeddings showed benefits or mixed outcomes depending on task and transfer choices.Source: §2 · p. 3

3. Experiments Setup

The experiments evaluate pre-trained visual representations across four control domains, comparing supervised and self-supervised models with frozen, from-scratch, and ground-truth baselines. Policies are trained from optimal trajectories and evaluated in environment-specific rollouts.

3.1. Environments

  • Experiments span Habitat ImageNav across five Replica scenes, five DeepMind Control tasks, two difficult Adroit tasks, and five Franka Kitchen tasks.The Adroit tasks are Relocate and Reorient Pen; Kitchen randomizes the arm’s starting pose but not the scene.Source: §3.1 · pp. 3–4

3.2. Models

  • The evaluated pre-trained representations include ImageNet-supervised ResNet-34 and ResNet-50, MoCo ResNet-50, and CLIP-pretrained ResNet-50 and ViT models.MoCo learns representations through instance discrimination with augmented image views; CLIP jointly trains visual and textual representations on web image-text pairs.Source: §3.2 · p. 4
  • The baselines are frozen random convolutional features, end-to-end training from scratch, and simulator-provided ground-truth state features treated as an oracle.The paper characterizes from-scratch visual learning as data-hungry and seeks to compete with the ground-truth state baseline.Source: §3.2 · p. 4

3.3. Policy Learning and Evaluation with PVRs

  • 10,000 shortest-path trajectories per Habitat scene yield approximately 2.1 million data points, while MuJoCo tasks use 25–100 RL-generated optimal trajectories each.Policies use behavioral cloning with frozen pre-trained perception modules; Habitat uses an LSTM and MuJoCo an MLP with a fixed history window. Success is measured by destination reach in Habitat and task success in Adroit and Kitchen; DMC returns are rescaled to [0, 100].Source: §3.3 · pp. 4–5

4. Experiments Results and Discussion

Experiments find that pre-trained visual representations can support control across diverse domains, and that performance depends on the training data, augmentations, and feature layers. A full hierarchy of MoCo features trained only on out-of-domain data solves all four domains, sometimes outperforming ground-truth features.

4.1. How do Off-the-Shelf Models Perform for Control?

  • Any off-the-shelf PVR beats frozen random features and scratch-trained perception in the small-data regime, but no PVR dominates across all four domains.On average, self-supervised MoCo models outperform supervised RN50 and CLIP; MoCo is competitive with ground-truth features in Habitat but not MuJoCo.Source: §4.1 · p. 5

4.2. Datasets and Domain Gap

  • In-domain pre-training helps over training from scratch but is not much better than ImageNet or Places; MuJoCo expert-trajectory PVRs are not competitive with those out-of-domain alternatives.The in-domain datasets use expert-policy trajectories; larger or more diverse deployment-environment datasets might help but would conflict with the goal of data-efficient learning.Source: §4.2 · p. 5, Figure 7 · p. 6

4.3. Recognition vs. Control: Two Tales of Invariances

  • Color-only augmentation underperforms in every domain except DMC, while crop-only augmentation performs as well as or better than other representations.The authors hypothesize that crops emphasize relative displacement between the agent and objects, a useful control cue; control may therefore require different invariances than semantic recognition.Source: §4.3 · p. 6, Figure 8 · p. 6

4.4. Feature Hierarchies for Control

  • Early-layer features are more effective for fine-grained MuJoCo control and can match or outperform ground-truth features, whereas later layers work better for Habitat navigation.The authors attribute Habitat’s later-layer advantage to navigation’s need for semantic information, such as detecting walls and obstacles.Source: §4.4 · pp. 6–7, Figure 9 · p. 7

4.5. Full-Hierarchy Models

  • The full MoCo hierarchy using layers 3, 4, and 5 performs best on average and solves all four domains, sometimes outperforming ground-truth features despite using only out-of-domain training data.The policy can exploit useful features from different layers: layer 3 is strong in MuJoCo, while layer 5 performs well in Habitat.Source: Figure 10 · p. 7, §4.5 · p. 7

5. Discussion and Conclusion

Off-the-shelf visual representations trained on out-of-domain data can support control across diverse tasks, and the paper identifies training methods, augmentations, and feature hierarchies that help them work. The authors advocate pursuing universal pre-trained representations while noting limits of their evaluation.

PVR: Freezing vs. Fine-Tuning

  • The experiments froze vision models without on-the-fly representation fine-tuning, leaving evaluation in the full fine-tuning regime to future work.The authors describe this as similar in spirit to the linear-classification probe protocol for evaluating computer-vision representations.Source: §5 · pp. 7–8

Imitation Learning vs. Reinforcement Learning

  • The authors chose behavior-cloning imitation learning because pre-trained representations matter in sparse-data regimes and evaluating many models with current RL was prohibitively expensive.They note that leading RL algorithms require millions or billions of samples and hope their findings can help improve RL for control (Yarats et al., 2021b; Wijmans et al., 2020; Hafner et al., 2020).Source: §5 · p. 8

Summary of Our Contibutions

  • The study found SSL features better than supervised features, crop-induced translation and occlusion invariance more relevant than illumination and color invariance, and feature depth should match task type.Early convolutional features worked better for fine-grained MuJoCo control, while later convolutional features worked better for semantic Habitat tasks.Source: §5 · p. 8

Towards Universal Representations for Control

  • The multi-layer, crop-augmented MoCo PVR trained on out-of-domain data was competitive with or outperformed ground-truth features across all four evaluation domains.The PVR combined features from multiple layers of the crop-augmented model.Source: §5 · p. 8
  • The authors advocate pre-trained visual perception over hand-designed state features for real-world control, where state estimation is hard, citing reduced data needs, improved performance, and simpler frozen-policy training.They hope the results encourage development of a single out-of-domain PVR usable as the perception module for any control task.Source: §5 · p. 8

A. Training Details · A.1. Habitat Details

Habitat policies use separately encoded agent-view and target images, concatenated before policy processing, with ground-truth and random-feature baselines. Training uses recurrent policies optimized with RMSProp, and performance is evaluated over repeated online trajectories.

A.1. Habitat Details

  • Habitat PVRs encode two 64×64 RGB images independently—one agent view and one target-location view—then concatenate their encodings for the policy.The ground-truth baseline is a 12-dimensional vector of agent position and quaternion, target position, scene ID, and version; the random-feature baseline follows Parisi et al. (2021).Source: §A.1 · p. 12
  • PVR policies apply batch normalization, a two-layer MLP, a two-layer LSTM, and a final MLP, with 1,024 units in all hidden layers.Ground-truth features omit batch normalization because it significantly harmed performance.Source: §A.1 · p. 12
  • Policies train for 125,000 updates using RMSProp at learning rate 0.0001, with 16 mini-batches of 100 consecutive steps and gradient norm clipped at 40.The optimization setup follows Parisi et al. (2021).Source: §A.1 · p. 12
  • Success rate is estimated over 50 online trajectories and averaged across the final six policy updates, totaling 300 trajectories per seed.Source: §A.1 · p. 12
  • Imitation-learning data comprise 50,000 optimal trajectories collected with Habitat’s native solver, or approximately 2,100,000 samples.The trajectories include 10,000 per scene.Source: §A.1 · p. 12

A.2. MuJoCo Details

The appendix specifies visual-only PVR inputs alongside ground-truth and random feature setups, plus policy training, evaluation, and imitation-data protocols.

A.2. MuJoCo Details

  • Policies use the last three 256×256 RGB observations, encode each independently with PVR, and fuse them through latent differences without proprioceptive inputs.Ground-truth features are simulator-provided vectors encoding agent and environment information; the random-feature baseline uses a four-layer convolutional network.Source: §A.2 · p. 12
  • The fused PVR passes through batch normalization and a three-layer, 256-unit MLP, trained with Adam on batches of 256 for 100 epochs.Policies are evaluated every two epochs over 100 online trajectories, with reported performance averaged over the three best epochs.Source: §A.2 · p. 12
  • Imitation datasets comprised 100 trajectories per Adroit task, 100 per DeepMind Control task, and 25 per Franka Kitchen task.Adroit horizons were 100 steps for Reorient Pen and 200 for Relocate, totaling 10,000 and 20,000 samples; DeepMind Control used action repeat 2 and 500-step horizons, totaling 50,000 samples per task; Franka Kitchen used 50-step horizons, totaling 6,250 samples overall (1,250 per task).Source: §A.2 · p. 12

A.3. PVRs Details · Datasets

The PVR datasets span ImageNet, Places, Habitat, and MuJoCo images. Habitat data came from sampled optimal trajectories across Replica scenes, while MuJoCo data came from expert trajectories, with extra Kitchen trajectories used only to train PVRs.

Datasets

  • 250,000 DeepMind Control images, 30,000 Adroit images, and 25,000 Kitchen images were collected for MuJoCo PVR training.Adroit and DeepMind Control images came from the same expert trajectories used for imitation learning.Source: Datasets · p. 12
  • 20,000 optimal trajectories were collected from all 18 Replica scenes for Habitat, retaining one frame every three for diversity.Source: Datasets · p. 12
  • 1.2 million ImageNet images were included in the PVR datasets.Source: Datasets · p. 12
  • 1.8 million Places images were included in the PVR datasets.Source: Datasets · p. 12
  • Kitchen PVR training used additional expert-policy trajectories because its imitation-learning dataset contained only 6,250 examples.Source: Datasets · p. 12
  • The additional Kitchen trajectories were used only to train PVRs, not the policy.Source: Datasets · p. 12

Vision Models · A.4. Intermediate Layers Compression · A.5. Compute Details

The appendix identifies the vision-model implementations and describes how intermediate-layer features were compressed for fair comparisons. It also reports the compute resources and training times used for pretraining and policy imitation.

Vision Models

  • The implementations used PyTorch Vision ResNet, MoCo v2, and OpenAI CLIP in ViT-B/32 and RN50 versions.The appendix gives the corresponding GitHub repositories for each model family.Source: Vision Models · p. 12

A.4. Intermediate Layers Compression

  • Intermediate-layer features were compressed to the fifth-layer representation’s size using two residual blocks: one reduces channels and the next restores them.This was intended to avoid prohibitively expensive compute and enable fair comparisons across layers.Source: §A.4 · p. 13
  • After random initialization, the added blocks were fine-tuned on the original pretraining task, and the first block’s output served as the compressed feature representation.The compressed features were then used in the experiments.Source: §A.4 · p. 13

A.5. Compute Details

  • Vision-model pretraining and layer compression were distributed across two SLURM-cluster nodes, each with four NVIDIA GeForce GTX 1080 Ti GPUs.Pretraining one PVR model took 1–3 days, depending on training method, model size, and dataset.Source: §A.5 · p. 13
  • Policy imitation learning used a NVIDIA Quadro GP100 GPU on a SLURM-based cluster.Training one policy, including evaluation, took 8–24 hours depending on the PVR and environment.Source: §A.5 · p. 13
  • Pretraining a PVR model took 1–3 days, with duration varying by training method, model size, and dataset.This timing applies to PVR pretraining, not policy imitation learning.Source: §A.5 · p. 13
Loading 2203.03580v2…