Theia: Distilling Diverse Vision Foundation Models for Robot Learning

Authors: jinghuan  shang
karl  schmeckpeper
brandon b may
maria vittoria minniti
tarik  kelestemur
david  watkins
laura  herlant
Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-04 00:17:43

Abstract: Vision-based robot policy learning, which maps visual inputs to actions, necessitates a holistic understanding of diverse visual tasks beyond single-task needs like classification or segmentation. Inspired by this, we introduce Theia, a vision foundation model for robot learning that distills multiple off-the-shelf vision foundation models trained on varied vision tasks. Theia's rich visual representations encode diverse visual knowledge, enhancing downstream robot learning. Extensive experiments demonstrate that Theia outperforms its teacher models and prior robot learning models using less training data and smaller model sizes. Additionally, we quantify the quality of pre-trained visual representations and hypothesize that higher entropy in feature norm distributions leads to improved robot learning performance. Code, models, and demo are available at https://theia.theaiinstitute.com.

Main paper: 
 introduction: 
 Visual understanding, i.e., the process of abstracting high-dimensional visual signals like images and videos, includes many different sub-problems, from depth prediction [1] and vision-language correspondence [2, 3] , to tasks ranging from coarse to fine granularity such as classification [4, 5] and object grounding [6, 7, 8] , as well as tasks defined along spatial and temporal axes like segmentation [9, 10] and tracking [11] . Given this diversity, a long-standing effort in the vision community has been to develop models tailored to one or a few specific types of visual understanding tasks. In recent years, several models [4, 5, 10, 2, 3, 12, 13, 6, 14, 15, 10, 7, 16] have achieved remarkable generalizability to unseen domains and new tasks, and are commonly referred to as vision foundation models (VFMs).Vision-based robot policy learning, which learns action policies from visual inputs, requires strong and diverse visual comprehension. These policies involve many implicit vision tasks such as object recognition and semantic grounding, where off-the-shelf VFMs corresponding to some well-defined tasks can be easily found, but there is no single model for all of vision tasks. Studies have shown that off-the-shelf VFMs such as CLIP [2] usually under-perform relative to visual representation models tailored for specific tasks in robot learning [17, 18, 19, 20, 21] . This fact reveals a gap between the needs of robot learning and the limited visual understanding capabilities of any individual VFM. Our motivation is different than all prior works on learning visual representation models for robotics, where those works focused primarily on improving training data [17, 21, 20] , designing objective functions [17, 18] , and directly taking advances of vision architectures [19] . We uniquely focus on improving the visual representation from the angle of solving multiple implicit visual understanding tasks, which will benefit the downstream robot learning.In this work, we propose combining multiple large VFMs into a single, smaller model for robot learning that leverages diverse visual understanding abilities from VFMs. We achieve this via knowledge distillation [22] . Unlike conventional distillation from a larger model to a smaller model on the same task, we distill VFMs tailored for varied vision tasks to improve visual representation for robot learning, which is an unseen task for VFMs.We introduce Theia, a robot vision foundation model that simultaneously distills off-the-shelf VFMs such as CLIP [2] , DINOv2 [7] , and ViT [5] . Compared to off-the-shelf VFMs [2, 5, 7] and prior works [19, 20] , Theia offers both better pre-trained visual representations for higher downstream robot learning performance and reduced computational costs. Furthermore, training Theia only requires ImageNet [23] and about 150 GPU hours, in contrast to prior works which necessitate substantially more compute [19, 20, 17, 18] . To understand what makes a good visual representation for robot learning, we observe multiple factors that relate to the performance of downstream robot learning tasks. We hypothesize that higher entropy in representation norms [24] correlates with improved robot learning performance.In summary, our contributions are:• We introduce Theia, a model that combines knowledge from multiple VFMs into a single, smaller model using knowledge distillation with low training cost. • Through extensive simulated and real-world experiments, we confirm that Theia's visual representations lead to better downstream robot learning with improved computational efficiency. • We identify key factors relevant to robot learning performance, such as model size, the use of spatial tokens, and the entropy of representation norms, offering valuable insights for guiding future research on optimizing visual representations for robot learning.2 Related Work 
visual representations for robot learning: 
 Visual representations are important for vision-based robot policies to parse high-dimensional visual signals. Visual representation learning can happen at different stages, including pretraining [17, 18, 20, 21] , joint-learning with robot tasks [25, 26, 27, 28] , or a combination of both using either trainable or frozen visual representations [29, 30] . Off-the-shelf visual encoders [4, 5, 2, 31, 32] can also provide visual representations for robot learning. Additionally, an important factor in training visual representations is the choice of data. ImageNet [23] , as suggested by Dasari et al. [21] , is a particularly effective pre-training dataset, while video datasets [33, 34] are also widely used. Training objectives and auxiliary tasks for visual representation learning vary, including data augmentation [25, 26] , prediction tasks [28] , contrastive learning [27, 35, 36] , and self-supervised learning [19, 20] . Specifically, to handle invariance and equivariance in visual observations, inductive biases [37, 38] and constraints [39, 40] can be introduced into neural networks to improve visual representation quality. Unlike prior works, we build a robot vision foundation model from the perspective of merging the visual understanding abilities of VFMs via distillation. Concurrent work OpenVLA [32] uses pre-trained SigLIP [41] and DINOv2 [7] encoders, which suggests the benefit of diverse visual understanding in robot tasks. Vision foundation models trained on large-scale data exhibit strong task-specific performance and transferability to unseen domains and new tasks. VFMs can focus on single tasks, multiple tasks, or remain task-agnostic at the pre-training stage. For example, ViT [5] , DeiT [42] , and ConvNeXt [43] are designed for image classification, SAM [10] for semantic segmentation, and Depth-Anything [1] for depth prediction. When combined with Large Language Models (LLMs) [44, 3, 45, 46] , these models can solve multiple visual tasks including referring segmentation, visual question answering (VQA), and image editing. Task-agnostic models like the DINO series [6, 7] are trained through selfdistillation, while CLIP [2] is trained by aligning image-text pairs. Given the strong generalization capability of VFMs, robot learning should also benefit from leveraging the latent representations of pre-trained VFMs. which is the primary motivation for our research.Knowledge distillation [22] compresses the knowledge from one or more larger models into a single smaller model. To leverage VFMs, which are usually computationally intensive, several studies have explored distilling them into more compact models. For example, some works distill a single VFM like SAM [10] into smaller variants [47, 48, 49] . There are also relevant works on combining two models: SAM-CLIP [50] merges SAM [10] and CLIP [2] into one model, while ViLD [51] and GroundingDINO [52] combine vision and language models to perform object detection. RA-DIO [53] is an agglomerative model that simultaneously distills CLIP, SAM, and DINOv2. These studies show that combined models can improve performance on downstream tasks or enable new applications. Similarly, in this work, we investigate whether combining VFMs will benefit robot learning. Notably, RADIO [53] is the closest approach to ours. The key differences between Theia and RADIO [53] are that (1) We aim to use our representation for robot learning tasks that none of the VFMs have covered in their pretraining tasks, (2) we distill only spatial tokens rather than both spatial and [CLS] tokens, (3) we choose a different set of teacher models, and (4) we show how each teacher model contributes to robot learning performance. 
method: 
 Overview. We introduce Theia, a framework that distills the knowledge of multiple VFMs into a smaller model, producing rich spatial representations for downstream vision-based robot learning.Figure 2 shows the overall design of Theia. Our model comprises a visual encoder (backbone) and a set of feature translators for distillation. Note that only the visual encoder is used to produce latent representations for downstream robot learning tasks.Architecture. Given an input image x, the visual encoder f (•) produces a rich representation z = f (x) (called the Theia-representation), which is utilized for downstream robot learning tasks. We focus on backbone models that are smaller than typical VFMs, specifically using ViT-Tiny, Small, and Base [5, 42] 1 motivated by the limited computing resources on robotic systems. We use Theia-Tiny (Theia-T), Theia-Small (Theia-S) and Theia-Base (Theia-B) to refer to Theia models using a vision transformer backbone with the corresponding size. To train the Theia-representation, we perform feature distillation with the help of feature translators, which will be described below.The Theia-representation is a set of encoded tokens corresponding to input image patches. We choose spatial tokens because spatially-dense representations are the foundation for diverse visual understanding, as evidenced by the powerful per-patch features in DINOv2 [7] . Therefore, we aim to distill all spatial tokens and leave the [CLS] token untouched.Feature Translators. Our goal is to supervise Theia-representations with teacher representations from various VFMs. We extract teacher representations h i (x) of VFMs at the last layer for CLIP [2] ,ViT [5] and DINOv2 [7] , or before the decoders for SAM [10] and Depth-Anything [1] . Since a single representation cannot be learned to match all the teacher representations directly, feature translators g i (•) are used to map the Theia-representation, z, to each teacher representation. Feature translators are shallow CNNs to ensure that knowledge is distilled primarily into Theia's visual encoder. Details are available in the Appendix.Distillation Objective. Our training objective is matching the outputs of the feature translators with their corresponding teacher VFM representations. To achieve this, we use a combination of cosine and smooth-L1 losses [53] to match each pair of predicted and ground truth representations for the same image, taking their weighted average. Formally, our loss isEQUATIONwhere x is the input image, M is the number of teacher VFMs, α i is the loss weight for each teacher, and β is the weight for balancing cosine loss and smooth-L1 loss respectively. In general, we set α i = 1/M such that the loss weights each teacher equally. We empirically set β = 0.9 [53] .Feature Normalization. To properly accommodate the different scales of teacher representations, we first perform a normalization step. This helps us scale the loss of different teacher features evenly and avoid biasing (collapsing) to a teacher model with extremely larger norms. We perform the normalization on the teacher representations over each latent dimension, where mean and variance are calculated from all ImageNet training samples. Details are in Appendix A.Dataset. We train our model on the ImageNet [23] training set for 50 epochs. We opt to use ImageNet because of its greater diversity compared to human videos [33, 54, 55, 56] and robot datasets [57, 58, 59, 60, 61] within the same number of images. This diversity has been experimentally shown to improve visual representation learning [21] . 
benchmark and settings: 
 To evaluate pre-trained visual representations, we use simulation tasks in CortexBench [20] , which combines MuJoCo tasks (Adroit [62, 63] , DeepMind Control Suite (DMC) [64] , and Meta-World [65] ), Habitat [66, 67] tasks (ImageNav [68] , ObjectNav [69] , and MobilePick [70] .), and Trifinger [20] tasks 2 . ImageNav and MobilePick are reinforcement learning (RL) tasks and others are imitation learning (IL) tasks. We follow the experiment settings of Majumdar et al. [20] and report aggregated scores for rewards (DMC tasks) and success rates (all other tasks). For DMC tasks, raw rewards are divided by 10 to be in a scale consistent with the success rate. We also conduct real robot experiments, introduced in Section 4.3. We use the same policy heads for the each type of representations (vector or spatial tokens). Full experimental details are available in the Appendix. We comprehensively evaluate Theia and baseline pre-trained models on the MuJoCo subset of Cor-texBench [20] to provide an overall assessment of pre-trained visual representations. We consider prior works on visual representations for robot learning, including R3M [17] , VIP [18] , MVP [19] , and VC-1 [20] , as well as agglomerative models for vision tasks RADIO and E-RADIO [53] , and off-the-shelf vision foundation models ViT [5] , DINOv2 [7] , and CLIP [2] . We also test a naive concatination of three VFMs (CLIP, ViT, and DINOv2) which are used to train Theia, referred as CDV.All pre-trained representations are frozen in this experiment. Throughout this section, we answer the following questions:• How does Theia perform compared to baselines?• Which is more effective for visual representations: [CLS] or spatial tokens?• How does robot learning performance scale with the size of the visual encoder? Theia Performance. As shown in Figure 4 , Theia outperforms all evaluated models, surpassing the performance of the best prior models, R3M and MVP, as well as agglomerative models for vision tasks RA-DIO and E-RADIO. We also tested a naive approach of using multiple VFMs (CLIP, DINOv2, and ViT) simultaneously by concatenating their spatial tokens channel-wise (CDV in Figure 4 ), but this performed much worse than using just the individual VFMs. Theia models scale effectively from tiny to base sizes, with Theia-S and Theia-B being the only models to break scores of 80 on this subset of CortexBench, even though they use only a small fraction of the inference computation required by comparable models. Theia's training is very efficient, using only the 1.2M images in ImageNet with a training time of about 150 GPU hours on NVIDIA H100 GPUs, compared to approximately 5M images used in prior works [17, 18, 19, 20] and 1B images used by RADIO [53] .[CLS] Token. We evaluate Transformer-based models (all models in Figure 4 except R3M and VIP) using either their [CLS] token or spatial tokens [5] for downstream robot learning. To accommodate spatial tokens in the robot policy, we introduce extra shallow CNN layers at the input of the policy network, known as "compression layer" [71, 20] . Figure 4 shows the results of all models we evaluated, clearly showing that for Transformer-based models, providing spatial tokens is consistently better than using the [CLS] token for robot learning. This finding applies to both off-the-shelf VFMs and prior pre-trained representations, including MVP and VC-1. Scaling with Model Size. We observe that most models, including Theia, achieve better performance with larger sizes. The scaling effect is more obvious when using the [CLS] token, probably because the [CLS] token encodes less information, making the size of the feature vector more critical. Different models that use spatial tokens also scale at varying rates. VC-1 shows only minor improvements when scaling from base to large, while ViT shows much larger improvements when scaling from tiny to huge model.CortexBench Results. We compare Theia-B against baselines over the CortexBench evaluation suite 2 . The results in Table 1 confirm that Theia-B outperforms all other models.Based on simulation performance, we test Theia-B and the best-performing baseline models: MVP-L [19] , R3M [17] , VC-1-L [20] , DINOv2-L [7] , ViT-H [5] , and E-RADIO-L [53] for evaluation on real-world tasks. We employ four tasks (Figure 3 ): Door Opening, Pick-and-Place, and Toy-Microwave Cooking with a WidowX 250s arm, and Drawer Opening with a Boston Dynamics Spot.We train behavioral cloning policies on top of visual representations using conventional policy networks composed of CNNs and MLPs in the WidowX setup and diffusion policy [72] in the Spot setup. During testing, we vary the robot position for Door Opening and Drawer Opening, randomize the object position for Pick-and-Place, and randomize both the object positions and object types in Toy-Microwave Cooking. Full experimental settings, including details about the number of collected demonstrations, the interface used to collect them, and the policy architecture, are available in the Appendix (Sec. D.3).Table 2 shows the success rates on these real-world tasks. Theia-B achieves the highest success rate across all tasks except Toy-Microwave Cooking. The results also highlight that the Theiarepresentation is useful for both conventional and diffusion-based policy heads, and for either freezing or fine-tuning the visual representation. E-RADIO is the most competitive model in this setting amongst all models compared, likely due to its similar distillation of VFMs and much larger training dataset. VC-1 has high task variance, performing poorly on Door Opening but adequately on Pickand-Place. ViT-H works much better when being fine-tuned but DINOv2 does not, which could be caused by some fundamental differences in VFMs. Selection of Teacher Models. Theia is motivated by distilling many vision foundation models, but not all the models can be effectively integrated into one model nor do they contribute equally to downstream robot learning. In these experiments, we empirically identify the most effective combination of teacher models among five candidate VFMs: CLIP (referred to as C) [2] , Depth-Anything (De) [1] , DINOv2 (Di) [7] , SAM (S) [10] , and ViT (V) [5] . We select these candidates because they have been designed to perform well across various important vision tasks.To select the best combination, we train Theia-T using different sets of teachers and evaluate the learned representation on the MuJoCo subset of tasks. Figure 5 shows the results of different teacher combinations. We start with single VFM teachers and observe that C and V are the most beneficial.We then distill All teacher models altogether and discover that it performs better than four out of five single-teacher distillations, but performs worse than only distilling V, suggesting possible negative effects from some of the teacher models. We then remove each teacher model from All and distill the remaining four (All-X). The results show that removing V and C from the teachers causes the most significant performance drops while taking out S and De leads to either similar or improved performance. Given the (nearly) negative effects observed, we distill the rest of the three models (CDiV), which results in the best performance. We also try the teacher combination that was performed in RADIO [53] (CDiS), but it does not outperform CDiV.Model Design. We conducted ablation studies on distillation loss choices and model architecture.We compared two kinds of distillation losses: MSE and Cos+L1 from Ranzinger et al. [53] . As shown in Table 3a , we find that Cos+L1 performs better. For model architecture, we examined the effect of "register tokens" introduced by Darcet et al. [24] . These extra tokens are introduced at the input without supervision and are discarded when using the representation. With the same amount of spatial tokens supervised by distillation, we varied the number of register tokens of between 0, 1, 4, and 8. Note the unsupervised [CLS] token is removed in zero register token case, and it counts for 1 register token in other cases. According to Table 3b , we find that using 1 register token worked the best, while having no register tokens performed the worst. Compared to RADIO [53] , which distills both [CLS] and spatial tokens, our use of the [CLS] token as a register token provides an advantage. We also evaluated distilling both [CLS] and spatial tokens in the Theia model, and the results shown in Table 3c confirm the advantage of exclusively distilling spatial tokens. We present qualitative visualizations to demonstrate how Theia-representations can be transformed into teacher representations through feature translators. Using Theia trained with all teachers (CDeDiSV, All), we applied PCA for visualizing predicted DINOv2 [7] features, used the SAM [10] decoder to produce segmentation results, and used the Depth-Anything [1] head to produce estimated depth. Results are shown in Figure 6 with more examples in the Appendix.The visualizations indicate that our predicted representations can be decoded by the original VFM and produce reasonable results. Encouragingly, we find that predicted depth maps from the Theia-predicted representation appear to be more accurate compared to the original Depth-Anything model, particularly in the stove-top area and at the microwave door. This shows the potential benefit of combining different visual understandings. 5 What Makes Visual Representations Good for Robot Learning?Traditionally, the quality of the pre-trained visual representations is evaluated through downstream robot learning like IL or RL. However, it is unclear why different visual representations lead to varying robot learning performance outcomes. In this section, we quantify the quality of visual representations and analyze how they correlate with downstream robot learning performance.Feature Norm Distributions and Entropy. Darcet et al. [24] analyzed the norm of spatial tokens in vision Transformers and found that high-norm outlier tokens are detrimental to vision task performance. Following this, we investigate whether a similar phenomenon arises in visual representations for robot learning. We inspect the feature norms of Theia with different teacher combinations and baseline models evaluated in Section 4.2, and their corresponding performance on the MuJoCo subset tasks. We sample 1% of the MuJoCo task training set and calculate the L2-norm of each spatial token after encoding images. We measure the entropy of the feature norm distribution across all samples and patches per model and use it as a quantitative metric. To calculate the entropy of the distribution, we first discretize the distribution by a histogram. Then, we use the following formula to obtain the entropy H =i (p i * log(p i )), where p i is the probability of each discretized bin.We confirm that similar outlier tokens also appear in VC-1 corresponding to the image patches that are not task-relevant, shown in the visualizations of feature norms on the right of Figure 7 . In contrast, Theia has very few or no outlier tokens, and the tokens with higher norms are more taskrelevant even though Theia-representations are not trained on these robot images. In our quantitative analysis (Figure 7 , left), we divide the models into distilled and regular based on the observation that distilled models generally have higher entropy (fewer outliers, Figure 4 (c) in Darcet et al. [24] ). We find that there is a strong correlation (R=0.943) between entropy and robot learning performance among regular models, and a high correlation (R=0.638) among distilled models. We hypothesize that spatial token representations with high entropy (better feature diversity) encode more information that aids policy learning, while less diverse representations (low entropy) may hinder it. In the Appendix G, we discuss the results of other quantitative measurements, including feature similarity and PCA-explained variance ratios, where no strong correlations are found. 
conclusion: 
 In this work, we introduced Theia, a novel robot vision foundation model specifically distilled from multiple VFMs to enhance robot learning tasks. Theia builds a rich visual representation from diverse VFM teachers, preserving spatial features to ensure detailed visual comprehension. Through extensive evaluations on CortexBench and in the real world, Theia consistently outperforms stateof-the-art models, including all prior models for robot learning, off-the-shelf VFMs, and similarly distilled models for vision tasks. Our results highlight the effectiveness of distilling multiple VFMs into a compact model for superior performance in a variety of robot learning scenarios. Furthermore, we answer a key question about what kinds of visual representations lead to better robot learning by finding a strong correlation between the entropy of feature norms and enhanced downstream performance, offering insights for future research on optimizing visual representations for robotics. 
 
 Appendix: Backbone. We use the DeiT-Tiny, DeiT-Small, and DeiT-Base models [42] as our backbone architectures. We keep the [CLS] token in the model and in the forward pass, but there is no supervisory signal provided for it. As a result, the [CLS] token serves as a "register token" [24] , which provides some benefits for learning high quality representations. We train Theia from scratch (no pre-trained DeiT [42] weights are applied).The feature translators are composed primarily of CNNs, with a linear layer appended at the end to match the teacher's representation dimension. Pure linear transforms might not be able to map Theia-representations to all three teacher representations well, resulting in a failure of learning (See Table 6a ). Thus, we use three CNN layers to account for the fact that each teacher model's representations are very different from one another. Details are listed in Table 4 , where we show the architectural details of the translators used for our (student, teacher)-feature pairs. Feature Normalization. Formally, the normalization is:EQUATIONwhere c is the channel index for the feature, j is the index of image sample, and N is the number of samples in ImageNet.We train Theia on 8 NVIDIA H100 GPUs. The main bottlenecks in training are the data transfer speed between devices and the GPU memory bandwidth to load large spatial feature tensors, for example, of size 1280×16×16 for ViT-H and 256×64×64 for SAM. We pre-compute the features from all teacher models instead of doing inference on the fly. This approach requires extra storage space to save all the features extracted from the VFMs, but significantly saves on training time and avoids loading models with high GPU memory usage during training, such as Depth-Anything or SAM (a batch size of 16 cannot fit into 80GB of GPU memory). All training configurations are listed in Table 5 .Teacher VFM Features. We use the output representations at the last layer of ViT [5] , CLIP [2] , and DINOv2 [7] . For SAM [10] , we use its encoder output. For Depth-Anything [1] , since it is Training from Scratch Pre-trained Backbone 78.9 ± 0.8 80.8 ± 1.5initialized from DINOv2, we use the latent representation before the final convolution layer. When decoding SAM and Depth-Anything results from Theia-predicted representations, we send the predicted representations through the remaining layers of original models and obtain the output.We conduct two additional ablation studies to verify design choices in the Theia model. The first is a comparison between the current CNN-based feature translator and a linear feature translator.In Table 6a , we find that using a Linear feature translator leads to a significant performance drop.The second ablation studies whether Theia should be trained from scratch or be initialized using the pre-trained DeiT [42] weights. In Table 6b , we find that using pre-trained weights improves the downstream performance. This could be interpreted as the positive effect of incorporating knowledge from an additional useful model into the distillation process. We would expect to see similar performance improvements as more informative models are included during training. We train a behavior-cloning policy for each of the four evaluated setups and for each evaluated baseline (see Table 2 and Figure 3 ); the training hyperparameters are shown in Table 10 . To train each of the three tasks performed with the WidowX robot, i.e., Door Opening, Pick-and-Place, and Toy-Microwave Cooking, we collected human demonstrations by teleoperating the robot with a Virtual Reality (VR) controller using the setup introduced in [59] ; the number of collected demonstrations for each of the tasks is reported in Table 9 . The policy's observations are RGB images and robot joint states. Images are encoded by a pre-trained visual encoder and a randomly initialized, unfrozen feature neck ("compression layer" [71] ). We use the same feature neck as we did for the previously discussed MuJoCo tasks in Section 4.2. The encoded vector is concatenated with the robot's joint states, which is fed into a 3-layer MLP with a hidden dimension of size 256. The policy outputs end-effector commands, consisting of the end-effector's delta positions (Cartesian coordinates), delta rotations (Euler angles), and the gripper's opening/closing command; such commands are tracked by the robot at a frequency of 5 Hz. In addition to the hyperparameters listed in Table 10 , we vary the policy action prediction horizon depending of the difficulty of the task, i.e., at each step the policy predicts the next 10, 10 and 5 actions for Door Opening, Toy-Microwave Cooking, and Pick-and-Place, respectively.In the following, we give more details about the WidowX tasks showcased in our work.Door Opening. In this task the robot has to open a fridge door in a toy-kitchen setup; we identify two stages to evaluate the task's success: Open and Fully Open (see Figure 3 ). We place the robot in front of the fridge and collect 63 demonstrations to train the behavioral cloning policy. We vary the height (z-axis) of the robot base between 40-46cm, and the position (x-axis, parallel to the toy kitchen) of the robot base; samples from the initial state distribution of the demonstrations are shown in Figure 8 . At inference time, we vary the height of the robot base among {40, 42, 43, 44, 46} cm, and select between 5 randomly-picked positions along the x-axis (for all policies). For each robot position, we evaluate the policy twice, for a total of 50 runs.Pick-and-Place. In this task, the robot has to pick up a pink cup from a toy-sink and drop it into a drying rack located on the left of the sink. We collected 48 demonstrations to train the policy, where we varied the initial pose of the objects, as shown in Figure 8 . During evaluation, we vary the cup's starting position amongst a total of 10 positions, of which 8 positions are equally distributed about the perimeter, and 2 are in the center of the sink. We also roughly vary the direction of the cup handle towards the left or the right. In total, we evaluate this task for 20 runs. There are two key stages for which we measure the success rate: picking up the cup and successfully releasing it into the drying rack.Toy-Microwave Cooking. In this task, the robot has to pick up an object from within the pot on the stovetop, putting the object into a toy-microwave, and closing the microwave. In each test, we initialize the environment with the microwave door open. In this task, we collected 100 demonstrations across 10 different toy-food objects (10 demonstrations per object) with randomized object positions; examples from the initial state distributions of the collected demonstrations can be seen in Figure 8 . During evaluation, we test 40 runs on 10 seen, in-distribution objects (4 runs per object), and 10 runs on 5 unseen, out-of-distribution objects (2 each), for a total of 50 runs. Furthermore, we vary the position of the pot that holds the object. We point out that this task is characterized by a longer horizon with respect to the other two; in fact, it has three different steps: 1. picking up the object, 2. placing the object into the microwave, 3. and closing the microwave. In addition, the policy needs to generalize to the different poses of the objects on the stove-top and to out-of-distribution objects. The longer horizon and the variety of objects in this task make it particularly challenging, so using frozen visual encoders was not effective (0% success rate). However, with fine-tuning, the policies performed reasonably well.Spot Setup. For the Spot experiments, we train a diffusion policy [72] conditioned on the encoded image. The diffusion policy outputs the desired absolute positions and rotations of the end-effector and the gripper state. Hyperparameters for the policies are shown in Table 11 . Drawer Opening. In this task, the robot has to open the top drawer of a cabinet. The policy receives color images from a single forward facing camera mounted on the body of the Spot. A fiducial marker is used to enable the robot to walk to a random position and orientation for each trial. Samples from the distribution of random initial states are shown in Figure 9 . The starting locations vary by ±5 cm in x and y and ±0.2 radians in orientation. 50 successful demonstrations were collected using a scripted policy and we evaluate each policy for 20 trials. A trial is considered successful if the drawer is opened at least 10 cm. After each trial, a scripted policy uses the fiducial marker to reset the environment and the robot moves to a new random location. In Table 12 we report per-task scores of the models evaluated in Figure 4 , over the MuJoCo subset of tasks. In Table 13 , we report per-task scores of all 14 tasks we evaluated in Cortexbench, corresponding to Table 1 . Note that we perform the evaluation following the original Cortexbench [20] protocol, where there are a total of 75 runs per MuJoCo task (we increased it from 30 to 75 for DMC tasks), 100 runs in Reach Cube, and 4200 runs in ImageNav. Entropy of the Representation Norm Distribution. Given N representations produced by encoding N images per model, where each representation contains P spatial tokens, we discretize the distribution of token norms over all N ×P tokens by using a histogram. We normalize the count of each bin in histogram by the total number of tokens to obtain the probabilities of each bin. We then calculate the Shannon entropy, given byi p i log(p i ). We find that the distilled models have higher entropy than the regular models, so we divide them into two distinct groups. Results are plotted as model performance vs entropy on the MuJoCo tasks. We attach the full version of plot presented on the left of Figure 7 here in Figure 11 , including two plots corresponding to each category of models and one plot for all models.At the top of Figure 12 , we find that both CLIP and VC-1 have high-norm outlier tokens. To better visualize the values of normal tokens, we use the median of norm values to clip the values. Specifically, we clip the norm values to range [0, 2 * median] and visualize the clipped norm values on the bottom of Figure 12 . We find that the high-norm tokens are still not task-relevant.In Figure 13 , we attach the feature norm map of all other models. Among those, we find that MVP [19] , which performs well on CortexBench, also produces features without outlier tokens. Feature norms from Depth-Anything [1] and SAM [10] , in contrast, have low diversity.We investigated the distributions of the feature norms that could give more information. Plots in Figure 14 we show the histogram of the feature norm entropy of 4 selected models: ViT-B, DINOv2-B, VC-1-B, and Theia-B (refered to as rvfm deit base 4features in the legend).We find that Theia has the most balanced distribution of the feature norm, while the others are longtailed. ViT-B has two peaks in the histogram. VC-1 is extremely centered and heavily long-tailed. DINOv2 is relatively less long-tailed than ViT-B and VC-1, and has a relatively nicer distribution around the median value.PCA Explained-Variance Ratio of Representations. Similar to the entropy analysis, given N×P spatial token representations, we apply PCA to them and extract the explained-variance ratio (EVR) of each latent dimension. We calculate and plot the cumulative sum of EVRs, as well as calculate the Area Under the Curve (AUC) of cumulative sum of the EVR. When comparing Theia-B with ViT-B, DINOv2-B and VC-1-B (Figure 15 ), we find that Theia-B has the lowest AUC and the best MuJoCo performance, while VC-1 has the highest AUC and the worse MuJoCo performance amongst these 4 models. The higher AUC is caused by one or few principle components that have very high EVRs, indicating that these components are capturing the majority of the variance of the feature representations. This means that less information is encoded within such representations. In contrast, the Theia-representation has a low AUC which we believe is due to the rich information that has been encoded within the latent space. However, when extending the scope to encompass all the models we evaluated (Figure 16 left), we find that the AUC of the EVR does not have a strong correlation with robot learning performance.Cosine Similarity of Representations. We also use cosine similarity to analyze the representations from different models by first calculating the mean of all representations and then computing cosine similarity between each representation and this mean representation. Results are shown on the right of Figure 16 , which shows very weak correlation between cosine similarity and performance on CortexBench.In addition to robot learning, we evaluate the Theia-representation on vision tasks to show to how well such abilities are maintained after the distillation process. For example, to evaluate image classification performance we apply linear probing on the Theia-representation to classify images from ImageNet-1k [23] . We use mean pooling of the Theia-representation (i.e. spatial tokens) and the same training schedule as MAE [14] . Results are shown in Table 14 , where we find that Theia outperforms MAE [14] at the same model size, but is not comparable to SOTA results from models like DINOv2 [7] . [14] 67.5% DINOv2 [7] (ViT-L) 84.5%We evaluate Theia-representation on 3D awareness to show the learned diverse vision knowledge in the representation. In particular, we adopt probe3D [74] , an evaluation toolbox to investigate the 3D-awareness of pretrained models. This evaluation includes Monocular Depth Estimation, Surface Normal Estimation, and Multiview Correspondence. We present the evaluation results in Tables 15, 16 and 17, and we summarize them here:• Theia achieves the third best Depth Estimation performance among over 10 types of vision models (only worse than DeiT-III and DINOv2 models) (Table 2 in Appendix of probe3D [74] , NYU dataset).• Theia achieves comparable performance as the best performing model, DINOv2, on Multiview Correspondence on ScanNet (Table 4 , Block3 in Appendix of probe3D [74] ).• Theia has better Surface Normal Estimation performance than CLIP (Table 3 in Appendix of probe3D [74] , NYU dataset).We note that the Theia model is trained only on ImageNet with 224x224 resolution images, while many baseline models in this benchmark are trained on billions of images and higher resolutions. Theia achieves better performance than models with similar training data and resolution. Figure 13 : Feature norm map visualizations of ViT [5] , DINOv2 [7] , MVP [19] , Depth-Anything [1] , SAM [10] , RADIO [53] , and E-RADIO [53] .We evaluate Theia-B and baseline models on a simulated benchmark, Factor-World [75] , to test the generalization ability of the pre-trained visual representations. We use the Door-Open task from Factor-World benchmark environment, for which we collected 100 trajectories across 20 randomized environments. We trained the policies using the same policy head that we use in CortexBench MuJoCo tasks (Table 8 ), and we evaluated the generalization ability when varying each of 6 factors of variations over 20 different runs controlled by one factor at a time. In Table 18 , we report the success rates of the baselines when varying each factor. These results show that Theia-B has the best generalization abilities compared to R3M, E-RADIO, and MVP. 8th Conference on Robot Learning (CoRL 2024), Munich, Germany.For simplicity, we denote DeiT-Tiny and DeiT-Small[42] by ViT-Tiny and ViT-Small. DeiT was introduced by Touvron et al.[42] and the original ViT[5] does not have Tiny and Small models.Due to reproducibility issues, we are not able to evaluate Move Cube, ObjectNav, and MobilePick tasks 


Decision: This paper presents Theia, a distilled vision foundation model for robot learning, combining features from multiple pre-trained models. The reviewers acknowledge the paper's timeliness and extensive evaluation across simulated and real-world environments, showing competitive performance and improved inference efficiency. However, concerns are raised about the novelty of the approach, given its similarity to previous work like RADIO. The evaluation strategy is questioned, with suggestions for more robust generalization tests and diverse robot learning tasks beyond simple policy learning. While the feature norm entropy analysis is interesting, reviewers note it could be further developed.
Following the rebuttal, two of the three reviewers engaged in discussions with the authors, and the majority of their concerns were addressed, with both leaning towards accepting the paper.

Reviews: review: title: title: Official Review
paper_summary: paper_summary: This work presents a simple approach for distilling existing pretrained visual representation models such as CLIP, DINOv2, and SAM into a unified, smaller model. Through straightforward application of knowledge distillation techniques, the resulting family of Theia models (at different parameter counts) show marginally stronger results on policy learning (behavioral cloning and RL) in both simulation and the real-world compared to individual representations and naive agglomeration strategies.
main_review: main_review: This is a straightforward paper that leverages existing knowledge distillation techniques to learn a “multi-purpose” pretrained visual representation that captures the different inductive biases of existing models such as CLIP, DINOv2, SAM, and ViTs pretrained for ImageNet classification; the motivation of this work is that each of these features captures something different, and agglomerating/distilling them into a single model is strictly beneficial for a host of robotics tasks.

The evaluations on various forms of policy learning (via behavioral cloning or reinforcement learning) on top of the learned representations capture several different simulation environments and real-world environments, with different robot embodiments and task objectives. In each of these settings, the results show marginal improvement (in many cases, the gains are not significant) over baselines spanning individual representation backbones, as well as what I’d consider a compelling set of “strong baselines” involving channel-wise stacking of existing backbones, evaluating different feature variants from each model (CLS token or spatial tokens), as well as model scales. Finally, the analysis of the internals of the various representations are straightforward (scaling, register features, etc.).

Unfortunately, there are two weaknesses that I feel need to be addressed:

W1. **Evaluation Strategy & Practical Implications** — All representations in this work are fine-tuned for single task policy learning, and evaluated in-distribution (i.e., evaluation initial conditions like object positions, lighting, etc. are sampled from the same distribution as training). This does not feel sufficient, especially given that in many of the experiments, the gains over baselines are marginal. Following the prior work in evaluating representations for robotics, I think this paper would strictly benefit from targeted robustness and generalization evaluations — for example, those in Xie. et. al (https://arxiv.org/abs/2307.03659). 

Furthermore, the decision on when to freeze/finetune through the visual representation backbone is not clear, and only on looking at the Appendix did I realize that in some cases, frozen representations fail completely (0% success rate on Toy-Microwave). Why is this the case? If fine-tuning through the backbone strictly helps (and seems more in line with what existing policy learning methods like Diffusion Policy, ACT, Behavior Transformers, etc. do), why not evaluate this directly? Given the large body of work in robot learning and pretrained visual representations, it feels like the time for evaluating frozen visual representations in small-scale, individual target task finetuning experiments is over — I really think we need to look to the complex, multi-task settings, with clear robustness and generalization evaluations. Anything less fails to meet the way existing representations are used practically in existing policy learning pipelines, and seriously hurts the impact of work like this!

W2. **Failure to Evaluate “Diversity” of Visual Representations** — A claim made through the beginning of the paper is the importance of combining diverse pretrained features (and the different inductive bias that each of them provide) for improved downstream performance across robot learning tasks. Unfortunately, the rest of the paper only focuses on rather simple, single-task policy learning as the only evaluation task. Robot learning spans more than just this type of policy learning, compressing visual features to inform action prediction — for example, policies that output keypoints such as PerAct/VoxPoser, or many of the policies trained for longer-horizon manipulation (see papers that evaluate on RLBench as an example). Similarly, other, more modular policies output bounding boxes, operate over multiple image frames, and encode temporal history. Compressing robot learning to just “visual features → simple policy head → action prediction” seems like it’s missing a lot of the power of the Theia representations... I would highly encourage the authors to expand the set of evaluation tasks to optimize for this diversity, and actually estimate things like the representation’s ability to understand low-level shape priors as well as higher-level semantic priors.
questions: questions: Q1. Given that you’re learning the feature translators anyway, why is the feature normalization step (Section 3.2) necessary? Doesn’t this by nature change the feature norm distribution across patches (nullifying the analysis at the end of the paper)?

Q2. The paper is currently missing a lot of important details about the evaluation. While many of the task settings are from the VC-1 paper, it would be worth going into more detail about the exact setup here — especially since in the Appendix it seems that some details of the VC-1 evaluation were tweaked for this work. For the real world experiments, how many demonstrations were collected, and with what interface? Who collected them? What does the initial state distribution actually look like? How did you choose the policy architecture? What hyperparameters did you sweep over?
limitations: limitations: 1
review_summary: review_summary: Based on the above weaknesses and lack of important experimental details, I recommend rejection at this time.
title: title: Review for Theia: Distilling Diverse Vision Foundation Models for Robot Learning
paper_summary: paper_summary: This submission introduces Theia, a vision foundation model for robot learning applications, distilled for multiple foundation models for primitive vision tasks. During training Imagenet samples are processed through several VFMs (Clip, Dino,ViT etc) and their resulting spatial features are used for distillation-based training of another ViT-based transformer encoder, exploiting some learnable feature translator modules, to conclude to a unified representation. Additionally, an interesting insight that correlates the entropy of features norms with its effectiveness to robot learning tasks.
main_review: main_review: Strengths:
- The manuscript studies a very interesting and timely problem. 
- The distilled model Theia has been extensively evaluated across simulation and real-world environments, and achieves state-of-the-art accuracy outperforming all examined baselines.
- The proposed model also achieves a significant improvement in inference efficiency, offering an improved speed-accuracy trade-off that is particularly relevant for resource-constrained robot platforms. 

Comments:
- Although the results demonstrate the effectiveness of the proposed approach, and thus its contribution, the novelty of this submission is more limited. As the manuscript already discusses, a similar distillation has been proposed in [50] and the differences outlined in Sec. 2.3 can be considered an extension/ablation study of this work. Similarly, the discussion on the feature norm entropy is also very interesting, but largely inspired by the work of [24], applied on a different task. 
- The contribution of the paper could also be improved if the naive equally-weighted distillation scheme across models, was replaced by some more advanced or adaptive weighting scheme, e.g. taking into consideration the loss magnitude or the insights provided by the feature norm. 
- Some of the arguably unexpected results that arose during the ablation study (e.g. lack of effectiveness from distilling SAM features) have not been further analysed to provide more insights on the underling reasons. 

Presentation:
- The paper is really well-written and easy to follow, adequately describes related work and background information, and features
- The accompanying video is also helpful, and demonstrates the conducted real-world experiments and qualitative results.
questions: questions: - Do the authors have any insights, or qualitative observations, that justify the poor performance of Segment-Anything and Depth-Anything models? Such visual cues can be considered very important for robot manipulation, and it is surprising that their participation in the proposed distillation scheme downgrades accuracy. 

- Would a more naïve fusion method (e.g. feature vector averaging or appending) on the distillation of different models be comparably effective with the proposed feature translators? Experiments indicated that this is not a good option for off-the-shelve VFM fusion, but under the finetuning scheme of the proposed distillation approach, the contribution of the feature translators remains unclear.

- It is unclear whether the authors intend to open-source the weights of the proposed Theia model. Since the SoTA accuracy of the model is in the reviewers opinion the most important contribution of this submission, sharing it with the community would enhance the argument for paper acceptance. 

- How could the insights obtained by the proposed Feature Norm entropy metric be fed back to the training process, to improve its accuracy or data efficiency ?
limitations: limitations: 1
review_summary: review_summary: Overall, the manuscript contributes a new state-of-the-art vision foundation model for robotics, achieving state of the art accuracy on robot learning tasks. The paper is well-written and the results are convincing. The proposed distillation scheme and visual representation evaluation approach, although particularly effective, can be considered incremental to existing approaches.
title: title: Ensemble Distillation Method of VFMs
paper_summary: paper_summary: This paper produces a new VFM that distills knowledge from multiple existing VFMs such as CLIP, SAM, DINO, etc.
main_review: main_review: This paper distills multiple VFMs from CLIP, SAM, DINO, Depth-Anything, ViT for robotic tasks, similar to RADIO. Compared to RADIO, the main differences are the introduction of new VFMs, and the use of CLS tokens (which are minor). On simulation and real robotic tasks, Theia shows competitive performance and efficiency in the inference.

In order to learn from multiple teachers, it applies a feature translator to match the teacher representation and uses cos-L1 loss for the distillation, following RADIO.

The distilled model, called Theia, is shown to scale effectively, achieving good performance with less MACs. Moreover, the authors did some analysis on the use of CLS tokens and visualization of the ensemble representation. Finally, the authors show that the entropy of feature norms might be related to the quality of the visual representations.

Overall this paper is an extension of RADIO. The distillation method has the same cos-L1 loss and the same hyperparameters. The PCA visualization is also the same. Novelties include the removal of CLS token and the analysis of feature norm entropy. However, I'm not sure if this is enough for the acceptance in CoRL.
questions: questions: 1. Line 131-132: how are the teacher representations normalized? Is it batch normalization or layer normalization?
2. Could you explain more on the spatial tokens?
3. Why does the CLS token encodes less information than the spatial tokens?
4. Line 216: as far as I understand ViT is more like an architecture than a special model.
5. The formula for the entropy of the feature norms needs to be written down.
limitations: limitations: 2
review_summary: review_summary: Overall the novelty against RADIO is limited.