Source-linked AI summary
The Unsurprising Effectiveness of Pre-Trained Vision Models for Control
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, Abhinav Gupta
TL;DR
Control policy learning commonly trains visuo-motor policies from scratch on deployment-environment data, while the transferability of vision representations to control remains underexplored. The paper evaluates frozen pre-trained visual representations across diverse control domains and finds that carefully chosen representations can match or exceed ground-truth features despite entirely out-of-domain pre-training.
Problem
Control policy learning is dominated by tabula-rasa training from deployment-environment data, leaving open whether out-of-domain pre-trained visual representations can support different control tasks.
Method
The paper evaluates frozen pre-trained visual representations as perception modules across Habitat, DeepMind Control, Adroit, and Franka Kitchen, studying training methods, augmentations, and feature hierarchies.
Results
A single PVR combining multiple layers of a crop-augmented MoCo model was competitive with or outperformed ground-truth features across all four evaluation domains.
Takeaways & Limitations
Self-supervised features, crop-based invariance, and task-appropriate feature layers are central components of effective pre-trained visual representations for control.
Takeaways & Limitations
The experiments freeze vision models and leave full fine-tuning of representations for future work.
Abstract
from arXiv · showhide
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 entirely on out-of-domain data can support control across tasks, challenging tabula-rasa policy learning. Experiments identify training objectives, augmentations, and feature hierarchies associated with effective representations.
- Control policy learning remains dominated by training task-specific visuo-motor policies from scratch with extensive environment interaction.
- The paper investigates whether a single vision model trained on out-of-domain datasets can work across different control tasks.
- The proposed paradigm detaches perception from control: a vision model is trained once on out-of-domain data, frozen, and reused by policies in deployment environments.
- Frozen PVRs trained on completely out-of-domain datasets can be competitive with or outperform ground-truth state features for imitation-learning policies.
- Self-supervised learning provides better control features than supervised learning, while crop augmentations appear more important than color augmentations.
- Early convolution features suit fine-grained MuJoCo control, whereas later features suit semantic Habitat ImageNav tasks; combining layers yields a PVR competitive with or outperforming ground-truth features across studied domains.
2. Related Work
Prior work established transferable visual representations and explored visual-input control, but pre-trained encoders received limited and mixed attention. This paper broadens the evaluation across domains, representations, augmentations, and feature layers.
- Representation Learning: Representation learning transfers features from supervised or self-supervised pre-training to downstream applications, often improving data efficiency or performance.
- Policy Learning: Reinforcement and imitation learning have achieved control across locomotion, manipulation, dexterous hands, and indoor navigation, often using neural policies.
- Policy Learning: Continuous-control research has focused largely on ground-truth state features, which are available in simulation or motion capture but seldom in real-world settings.
- Pre-trained Visual Encoders in Control: Prior pre-trained-encoder studies reported improved navigation or manipulation transfer in some settings, while another study found mixed results across Adroit and DeepMind Control.
- Pre-trained Visual Encoders in Control: Compared with prior work, this study spans four visually diverse domains, more pre-trained representations, and visual invariances arising from augmentations and layers.
3. Experiments Setup
The experiments evaluate frozen pre-trained visual representations across diverse simulated control domains, using behavioral cloning and comparing them with random, from-scratch, and ground-truth feature baselines.
- Environments: The evaluation covers Habitat ImageNav, five DeepMind Control tasks, two difficult Adroit tasks, and five Franka Kitchen tasks.These domains span navigation, locomotion, manipulation, and dexterous goal-conditioned behaviors.
- Models: The study compares supervised and self-supervised pre-trained visual representation models, including ImageNet ResNets, MoCo, and CLIP.Random frozen features and end-to-end learning from scratch provide additional baselines.
- Baselines: Ground-truth simulator features serve as an oracle baseline because they compactly describe the complete agent and environment state.The comparison tests whether visual representations can approach a state-information reference.
- Policy Learning and Evaluation: After pre-training, encoders are frozen and provide perception inputs to policies trained with behavioral cloning on optimal trajectories.Success is estimated from evaluation rollouts in the environments.
- Data and Policy Setup: Habitat uses 10,000 trajectories per scene, while MuJoCo tasks use 25-100 trajectories each; DMC return is rescaled to [0, 100].Habitat policies use an LSTM, whereas MuJoCo policies use an MLP with a fixed history window.
- Evaluation: Figure 6 reports mean success over five seeds with 95% confidence intervals for off-the-shelf visual representations.The figure compares supervised learning, pre-trained representations, and ground-truth features across Habitat and MuJoCo.
4. Experiments Results and Discussion
The results show that pre-trained representations work better than scratch-trained perception in the studied small-data regime, while their effectiveness depends on training method, augmentations, and feature hierarchy.
- Off-the-Shelf Models: Any pre-trained visual representation outperforms frozen random features and end-to-end perception trained from scratch in the studied small-dataset regime.The result is consistent with representation learning being data intensive.
- Off-the-Shelf Models: Self-supervised MoCo models are better on average than supervised RN50 and CLIP models, but no representation is best across all four domains.MoCo matches ground-truth features in Habitat, whereas no off-the-shelf representation matches them in MuJoCo.
- Datasets and Domain Gap: In-domain pre-training is not much better than ImageNet or Places pre-training and is worse than those alternatives for MuJoCo.Larger or more diverse environment datasets might help, but would conflict with the goal of data-efficient deployment.
- Discussion: The findings suggest that successful control representations depend on captured invariances and feature hierarchy, not only on the pre-training dataset.This motivates customizing representations for control rather than relying solely on semantic-recognition conventions.
- Invariances: Crop-only augmentations perform as well as or better than other tested choices, while color-only augmentation underperforms outside DMC.The authors hypothesize that cropping highlights relative displacement between agents and objects.
- Feature Hierarchies: Early convolutional features are more effective for fine-grained MuJoCo control, whereas later features work better for semantic Habitat navigation.Later layers may encode semantic information useful for detecting walls and obstacles.
- Full-Hierarchy Models: Combining layers 3, 4, and 5 performs best on average and solves all four domains, sometimes outperforming ground-truth features.The policy can exploit complementary features from the full hierarchy.
5. Discussion and Conclusion
The paper evaluates frozen, out-of-domain pre-trained visual representations for control and identifies design choices supporting a broadly useful representation. It argues that such representations can rival ground-truth features while simplifying the control pipeline.
- Scope and Limitations: The experiments freeze vision models and exclude on-the-fly representation fine-tuning, leaving full fine-tuning for future work.The setup is analogous in spirit to linear-probe evaluation of representations in computer vision.
- Scope and Limitations: The study focuses on imitation learning rather than reinforcement learning, although the authors hope its insights will inform future RL work.The choice is motivated by the data intensity of strong RL algorithms and the importance of pre-trained representations in sparse-data settings.
- Discussion and Conclusion: Off-the-shelf vision models trained on out-of-domain data can be competitive with ground-truth features for control policies.The paper frames this as a departure from randomly initialized visual encoders trained from environment interactions.
- Representation Design: Self-supervised models provide better control features than supervised models, while crop augmentation is more important than color augmentation for SSL.The identified augmentation relevance is associated with translation and occlusion invariance rather than illumination and color invariance.
- Feature Hierarchies: Early convolutional features suit fine-grained MuJoCo control, whereas later convolutional features suit semantic Habitat tasks.This finding motivates combining features from multiple layers in the proposed PVR.
- Universal Representations for Control: A proposed multi-layer PVR combining crop-augmented MoCo features was competitive with or outperformed ground-truth features across all four evaluation domains.The evaluated domains were Habitat, DeepMind Control, Adroit, and Franka Kitchen.
A.1. Habitat Details
The Habitat setup uses paired visual observations for PVR policies and compares them with ground-truth and random feature baselines under a recurrent policy architecture.
- Visual Input: Habitat PVR policies receive two 64×64 RGB images: the agent’s scene view and the target location.Each image is encoded independently, and the encodings are concatenated before entering the policy.
- Baselines: The Habitat ground-truth baseline is a 12-dimensional vector containing agent pose, target position, scene ID, and scene version.The agent pose includes position and quaternion.
- Policy Architecture: The PVR policy uses batch normalization, a 2-layer MLP, a 2-layer LSTM, and a final 1-layer MLP with 1,024-unit hidden layers.The final MLP uses softmax activation, and ground-truth features omit batch normalization because it harmed performance.
A.2. MuJoCo Details
The MuJoCo experiments use sequences of visual observations fused through latent differences, compare visual and simulator-provided features, and train policies with supervised optimization.
- Visual Input: MuJoCo policies use the last three 256×256 RGB observations, encode them independently, and fuse their PVRs using latent differences.Policies rely solely on visual embeddings and use no additional proprioceptive observations such as hand joint encoders.
- Baselines: The MuJoCo ground-truth baseline is a simulator-provided low-dimensional vector encoding agent and environment information, with size dependent on task and agent.Examples include joint positions and goal positions.
- Policy Optimization: MuJoCo policies train with a 3-layer, 256-unit MLP after PVR fusion, using Adam with learning rate 0.001 for 100 epochs.Mini-batches contain 256 samples, and the number of policy updates varies with dataset size.
- Dataset Sizes: Adroit uses 100 trajectories per task, with 10,000 samples for Reorient Pen and 20,000 for Relocate.The horizons are 100 and 200 steps respectively.
- Dataset Sizes: DeepMind Control uses 100 trajectories per task and 50,000 samples per task after an action repeat of 2.Each trajectory has a 500-step horizon.
- Dataset Sizes: Franka Kitchen uses 25 trajectories per task with 50-step horizons, totaling 6,250 samples across tasks.This corresponds to 1,250 samples per task.
A.3. PVRs Details Datasets
The representation datasets span standard vision sources and environment-derived image collections, with different supervision and collection procedures across datasets.
- Standard Vision Datasets: ImageNet contains 1.2 million images, while Places contains 1.8 million images for representation pre-training.Both are standard computer vision datasets used as sources of pre-training data.
- Environment Datasets: The Habitat dataset contains approximately 2.4 million images collected from 20,000 optimal trajectories across 18 Replica scenes.Only one frame every three is retained for diversity.
- Environment Datasets: MuJoCo data includes 30,000 Adroit images, 250,000 DeepMind Control images, and 25,000 Kitchen images.Adroit and DeepMind Control images come from the expert trajectories used for imitation learning.
Vision Models
The paper uses ResNet, MoCo v2, and CLIP vision models, drawing implementations from their respective public repositories.
- ResNet implementations are sourced from github.com/pytorch/vision.
- MoCo uses the v2 implementation from github.com/facebookresearch/moco.
- CLIP uses ViT-B/32 and RN50 versions from github.com/openai/CLIP.
A.4. Intermediate Layers Compression
Intermediate-layer representations are compressed to the fifth layer’s size for computational efficiency and fair comparison, then used as features after task fine-tuning. The experiments required substantial cluster-based computation for both representation preparation and policy training.
- A.4. Intermediate Layers Compression: Intermediate-layer activations are compressed to the fifth layer’s representation size to control compute and enable fair comparisons across layers.
- A.4. Intermediate Layers Compression: After adding randomly initialized layers, the model is fine-tuned on its original pre-training task, and the first residual block outputs the experimental features.
- Pretraining one PVR model took 1-3 days, while training one policy took 8-24 hours depending on the model, method, dataset, and environment.