W HAT M AKES C ERTAIN P RE -T RAINED V ISUAL R EPRESENTATIONS B ETTER FOR R OBOTIC L EARNING ?
A BSTRACT
Deep learning for robotics is data-intensive, but collecting high-quality robotics data at scale is prohibitively expensive.
One approach to mitigate this is to leverage visual representations pre-trained on relatively abundant non-robotic datasets.
So far, existing works have focused on proposing pre-training strategies and assessing them via ablation studies, giving high-level knowledge of how pre-training design choices affect downstream performance.
However, the significant gap in data and objective between the two stages motivates a more detailed understanding of what properties of better pre-trained visual representations enable their comparative advantage.
In this work, we empirically analyze the representations of robotic manipulation data from several standard benchmarks under a variety of pre-trained models, correlating key metrics of the representations with closedloop task performance after behavior cloning.
We find evidence that suggests our proposed metrics have substantive predictive power for downstream robotic learning.
I NTRODUCTION
A longstanding goal of the application of machine learning to robotics has been to enable robots to handle the myriad, messy conditions of the physical world.
Although broad strides towards humanlike generalization have recently occurred in other domains of machine learning such as computer vision and natural language processing (Brown et al., 2020; Radford et al., 2021), such progress remains elusive for robotics.
This is in no small part due to the confluence of two observations: that scaling up high-quality robotics data collection (whether in sim or real) is prohibitively expensive (Dasari et al., 2019) and, as Moravec’s paradox suggests, that the requisite scale for a watershed moment in robotics may be far beyond that required for “higher-level” tasks (Moravec, 1988).
Motivated by this there has been a flurry of recent work (Shah & Kumar, 2021; Parisi et al., 2022;
Nair et al., 2022b) studying to what extent datasets that were designed for, e.g., semantic image understanding (Deng et al., 2009; Radford et al., 2021) or egocentric video analysis (Grauman et al., 2022) rather than robotics can nonetheless be leveraged to improve the outcomes of robotic learning pipelines.
A prominent example of this is deploying self-supervised representations obtained via pre-training on this data (He et al., 2020;
Radford et al., 2021; Nair et al., 2022b) on downstream imitation learning and reinforcement learning robotics tasks (Rajeswaran et al., 2017; Gupta et al., 2019; Yu et al., 2020).
However, the significant gap in data and objective between the pre-training and downstream stages makes the entire endeavor rather heuristic, and indeed the aforementioned works mainly focus on proposing pre-training strategies and assessing them via their downstream robotics performance.
While comparison with regard to (proxy) stakeholder metrics such as success rate are of course the gold standard, we lack a fine-grained understanding of what properties are are enabling the comparative advantage of better representations.
This is worth addressing: better representations are somehow ameliorating challenging aspects of the robotic learning problem, including (for the setting we consider) compounding covariate shift and regression to an output space (actions)
that the pre-training is entirely agnostic to.
Being able to attribute improved downstream performance to specific properties of these representations could help us distinguish which challenges may continue to be addressed by further scaling robotics-agnostic visual pre-training, and which need separate solutions.
In this work, we focus on finding properties of pre-trained visual representations that matter for downstream learning of robotic manipulation tasks.
We leverage tools from the neural network representational similarity literature and the intrinsic dimensionality literature to facilitate our empirical analysis of 5 popular pre-trained models on 12 downstream robotic manipulation tasks, each instantiated with 3 camera angles.
We follow Parisi et al.
and Nair et al.
(2022b) in focusing on imitation learning for the robotic learning stage.
We find that:
• Representational similarity between pairs of demonstration trajectories for a task under a pre-trained model correlates with closed-loop task success with behavior cloning under that pre-trained model (Section 3.1).
• Representational similarity between a task’s human-defined low-dimensional state space and a pre-trained model’s state embedding correlates with closed-loop task success with behavior cloning under that pre-trained model (Section 3.2).
• For each pre-trained model, either nonlinear xor linear estimates of representational intrinsic dimensionality correlates with success (Section 4).
P RELIMINARIES
In this section, we establish the scope of our investigation, introduce artifacts we use, and ground the process in notation.
This is done for the pre-trained visual representations (Section 2.1), the downstream robot learning algorithm and environments (Section 2.2), and the techniques we borrow from the representational similarity analysis (Section 2.3) and intrinsic dimensionality literatures (Section 2.4).
2. P RE -T RAINED M ODELS
Let I = { 0 , . .
, 255 } C × H × W denote a space of images.
We consider pre-trained models of the form f : I → R d z that take in an image I and produce a dense vector z
= f ( I ) .
Since we are interested in pre-training on large-scale non-robotics data, we consider the following works from the literature: R3M (Nair et al., 2022b), CLIP (Radford et al., 2021), and ImageNet (He et al., 2016).
Since R3M is generally the best-performing model, we also use two of its ablations.
These works all provide pre-trained models with the ResNet-50 (He et al., 2016) architecture 3 , a standard choice for transfer learning in computer vision.
Indeed, ResNet variants remain among the most popular architectures for vision-based robotic learning (Jang et al., 2022; Sundaresan et al., 2022).
See Table 1 for the dataset and primary objective terms for each model
[Table]
[Table]
[Table]
2. D OWNSTREAM R OBOT L EARNING
For downstream robotic learning, we consider the imitation learning setting.
For some task T , a collection of N demonstration trajectories D T = ( τ i )
Ni =1 is provided, each consisting of observationaction tuples τ =
(( o t , a t ))
Tt =1 , o ∈ O , a ∈ A .
This is used to produce a closed-loop policy π :
O → A that imitates the demonstrations and is successful in generalizing in the environ- ment to accomplish the task.
For our analysis, we will additionally leverage the availability of low-dimensional state information given by physical simulation: a tuple ( s, o ) involves processing information in state s ∈ R d s to render observations.
In the tasks we consider, an observation o =
( I, p ) , I ∈ I , p ∈ R d p is itself a tuple comprising of an image and low-dimensional robot proprioception.
A frozen pre-trained visual representation is used to produce an embedding of the image
z
= f ( I ) .
This is concatenated with the proprioceptive observation vector before being further processed into actions by a parametric function approximator g ϕ : R d z + d p → A .
For convenience, we will overload f : O → R d z + d p to include the concatenation; disambiguation should be clear from context.
The policy π = g ϕ ◦ f is thus a composition of feature extraction and action generation.
We follow Nair et al.
(2022b) in using robotic manipulation tasks sourced from three commonly used benchmarks: 5 from Franka Kitchen (Gupta et al., 2019), 5 from MetaWorld (Yu et al., 2020), and 2 from Adroit (Rajeswaran et al., 2017).
The tasks differ in state space, environmental assets, visual appearance, robot end-effector and action space, and functional behavior.
We use the behavior cloning data and results from Nair et al.
(2022b).
2. R EPRESENTATIONAL S IMILARITY A NALYSIS
Given a dataset D =
( x i ) ni
=1 , x ∈ X , let A ∈ R n × p be the result of applying function
f A : X → R p to
D .
We call A the representation of D under f A .
Let B ∈ R n × q be the representation of D under f
B .
A rich literature of techniques exists for quantitatively comparing high-dimensional neural network representations.
We adopt the convention of doing so in terms of distances (in the colloquial sense; they don’t necessarily satisfy the triangle inequality).
One family stems from the idea of first computing pairwise similarity between points in each representation, before comparing these similarity structures.
This bypasses issues such as neuron correspondence.
As an instance of this strategy, we consider centered kernel alignment (CKA) as popularized by Kornblith et al.
Because the authors find that linear and radial basis function CKA perform similarly on realistically trained neural networks, we follow in using linear CKA discrepancy for simplicity:
[Equation]
where ∥·∥ F is the Frobenius norm.
Ding et al.
(2021) find that CKA is undersensitive, specifically to removing all but the largest principal components of a dataset, and propose the use of the orthogonal Procrustes discrepancy (Sch¨onemann, 1966) as a more sensitive alternative:
[Equation]
where ∥·∥ ∗ is the nuclear norm.
This quantity is the solution to the orthogonal Procrustes problem
[Equation]
of aligning the representations under an orthogonal transformation.
The above is written for p = q , but unequal representation dimensionalities can be matched by the use of “dummy” dimensions.
Per standard practice, we preprocess each representation by subtracting the mean vector and dividing by the resulting matrix’s Frobenius norm.
2. I NTRINSIC D IMENSIONALITY
The manifold hypothesis posits that natural data exhibits low-dimensional structure embedded within conventional representations of high ambient dimensionality.
Much work has argued that the generalization performance of machine learning models, including neural networks, is far better explained by intrinsic rather than ambient dimensionality (Nakada & Imaizumi, 2020; Pope et al., 2021).
Various definitions and estimators of intrinsic dimensionality (ID) have been proposed, each with its own characteristics and interpretation.
We select two estimators that provide complementary information on the representations.
PCA- ˆ ID@ η 2 is the minimal number of principal components needed to achieve a cumulative proportional explained variance of η 2 ∈ [0 , 1] .
By definition, PCA- ˆ ID@ ηis the dimensionality of a linear subspace of the representation.
To provide a nonlinear sense of ID, we consider the two nearest neighbors estimator (TwoNN- ˆ ID; Facco et al., 2017).
This estimator assumes only that the dataset has locally uniform density, specifically in the range of the secondnearest neighbor of each datapoint.
This is used to derive the following relationship between ID, the ratio µ between a datapoint’s distances to its second-nearest and nearest neighbors, and the ratio variable’s cumulative distribution function F ( µ ) :
[Equation]
In practice, the estimate is determined by fitting ((log( µ i ) , − log(1 − F emp ( µ i )))
n i =1 with a straight line through the origin, where the empirical cumulant F emp is determined by sorting the values of ( µ i ) ni =1 in ascending order through a permutation σ , then defining F emp ( µ i ) =
σ ( i ) /n .
R EPRESENTATIONAL C ONSISTENCY
One long-identified challenge of offline imitation learning is the phenomenon of compounding covariate shift , a vicious cycle of imprudent actions and out-of-distribution observations.
We hypothesize that better pre-trained visual representations mitigate this issue by representing observations in a more consistent manner.
We view this as a helpful (if not necessary) condition for the policy to generalize well, a stepping stone towards “stabilization” properties that can e [Figure] g. bring worst-case regret bounds from quadratic down to linear in horizon (Ross & Bagnell, 2010).
3. I NTRAMODEL T RAJECTORY C ONSISTENCY
[Figure]
To evaluate this hypothesis [Figure] we need concrete notions of representational consistency.
We first choose to consider the expected representational similarity between pairs of demonstration observation trajectories for the same task under the same pre-trained model.
This involves making the assumption that for every task there is correspondence amongst demonstration observations recorded at the same time step.
Concretely, for a task T and a pre-trained representation f , we evaluate intramodel trajectory consistency
[Equation]
where d is a measure of representational dissimilarity.
We use d Procrustes and d linear CKA as introduced in Section 2.3.
In Figure 1, we plot c intramodel ( d, T , f ) against the success rate on task T of a policy trained via behavior cloning with frozen pre-trained representation f .
Different points involve a different task, camera angle, or pre-trained representation.
For c intramodel ( d Procrustes ) , we observe a moderate positive correlation with strong statistical significance: r (175) = 0 .
49 , p = 3 .
3 × 10 − 12 (two-sided).
It is also worth noting that the correlation is stronger in the region where trajectory consistency and success rate are both high.
For c intramodel ( d linear CKA ) , we observe undersensitivity, echoing observations from prior work (Ding et al., 2021).
Due to this, we do not view the linear CKA analysis as providing evidence for or against the hypothesis and omit it from the rest of the analysis in this section.
[Figure]
[Figure]
As our data is heterogeneous [Figure] it is important to determine whether this correlation is an artifact of aggregation `a la Simpson’s paradox.
Figure 2 presents the above results stratified by pre-trained model.
We observe relatively consistent moderate positive correlations for every pre-trained model with statistical significance.
Figure 3 stratifies by downstream task domain.
We again see consistent positive correlations for every domain with statistical significance.
Interestingly, there is substantial variation in the strength of the correlations amongst domains.
Figure 10 further stratifies the data by downstream task.
The small sample size per task results in the correlation analysis for most tasks having dubious individual statistical significance.
Nevertheless, out of the 12 tasks, 9 suggest moderate or weak positive correlation, and 3 suggest weak negative correlation.
Overall, the aggregate, by-model, and by-domain correlation analyses support the hypothesis.
Indeed, because the representations are used for downstream learning, we can say that higher downstream behavior cloning performance is partially explained by intramodel trajectory consistency.
3. M ODEL -S TATE C ONSISTENCY
In a simulated robotics context, there is in some sense a reference representation: the implicitly defined inverse model from observations to low-dimensional state.
As this representation is humanengineered to be sufficient yet compact for a given task, we can obtain another notion of representational consistency: [Equation]
[Figure]
[Equation]
that is, for a given task T and pre-trained visual model f
we measure the negative representational distance d between the representation of the task’s observations under f and the corresponding low-dimensional states.
We leverage this correspondence to concretize the notion of consistency, in the sense that variation in the representation should be commensurate to variation in the state.
The assumption of this correspondence is much less objectionable compared to the time-locked correspondence between demonstration trajectories assumed for intramodel trajectory consistency.
[Figure]
[Figure]
R EPRESENTATIONAL I NTRINSIC D IMENSIONALITY
While we have demonstrated some encouraging evidence for our representation consistency hypothesis, the consistency metrics we use in the previous section are degenerate in the following sense: [Figure]
[Figure]
trivial representations such as the zero function achieve perfect consistency
[Figure]
yet obviously would not facilitate downstream learning.
To complement the previous analysis
[Figure]
we measure the intrinsic dimensionality (ID) of the same downstream demonstration data under pre-trained visual representations; trivial representations are detectable by ID estimates near zero.
We consider the following two ID estimates.
Trajectory ID estimates the average ID along a demonstration trajectory:
[Equation]
and task ID estimates the ID of all demonstration data for a task:
[Equation]
To obtain a more meaningful quantity that facilitates comparison between tasks of varying visual and functional complexity, we normalize both quantities by the task ID using low-dimensional state to obtain relative ID:
[Equation]
For the base ID estimators ˆ ID use both TwoNN- ˆ ID and PCA- ˆ ID@ 0 .
As TwoNN- ˆ ID can be sensitive to outliers (Facco et al., 2017), when using it we discard the top 0% , 5% , or 10% of the regression data, chosen by resulting R 2 .
[Figure]
[Equation]
[Figure]
[Figure]
In Figures 8 and
[Figure]
we plot relative task ID with TwoNN- ˆ ID and relative trajectory ID with
PCA- ˆ ID@ @0 .
95 , respectively, against downstream task performance after behavior cloning.
We remark that there is a bifurcation between i) R3M and its variants and ii) CLIP and ImageNet for both analyses.
Furthermore, the two analyses are complementary: for any pre-trained model, either the nonlinear or the linear estimate of ID yields significant moderate correlation with downstream performance, but not both.
Altogether, this suggests that the representations of R3M and its variants are better described by a nonlinear manifold, and that those of CLIP and ImageNet are better described by a linear subspace.
This may be related to the fact that CLIP and ImageNet are trained under “higher-level“ semantic objectives compared to R3M. The positive correlations between the suitable ID estimate for each model and downstream success suggests that some of the failure cases may be due to undersensitivity to task-critical features in the raw observations.
This makes sense considering the significant gap between the pre-training data for each model and downstream domains.
R ELATED W ORK
Most of the pertinent prior work has already been introduced to set context for this work in Sections 1, and 2.
While Zhang et al.
(2022) also empirically analyze self-supervised visual representations and present correlations between representation properties and downstream decision-making, our work is distinct in several respects.
They consider a single domain (Atari) for both pre-training and downstream learning, whereas we consider a variety of pre-training domains for downstream robotic manipulation.
Their focus is on reinforcement learning, whereas ours is on imitation learning.
Their empirical analysis of visual representations is based on linear probing, whereas ours does not involve training any additional parameters.
Seo et al.
(2022) learn visual world models from action-free robotic manipulation videos and assess transfer to robotic manipulation and locomotion tasks in held-out benchmarks.
The breadth of data in their pre-training is relatively limited and their transfer process involves stacking an actionconditional module before model-based reinforcement learning, a pipeline much more complex than simple behavior cloning.
More broadly, we are inspired by the subfield of BERTology (Rogers et al., 2020), in which creative hypotheses about the inner workings of BERT (Devlin et al., 2018) are proposed and assessed; the work of Huh et al. (2016), which investigates the importance of dataset composition for transfer in computer vision; and Naitzat et al.
(2020), who investigate the role of geometry and topology in the success of high-dimensional neural network representations .
D ISCUSSION
In this work, we investigated properties of pre-trained visual representations that explain downstream robotic manipulation performance after imitation learning under the representations.
The representational consistency and representational intrinsic dimensionality metrics we proposed appear to correlate moderately and significantly with downstream success rate.
R3M and its variants yielded consistent behavior for all of our analyses.
This suggests that the success of R3M lies not within any variant’s particular cocktail of objective terms, but rather the temporal contrastive learning on a diverse dataset that they all partake in.
Since the representational metrics proposed in this work moderately correlate with downstream success, and since the direction of causality is established by the use of the representations for downstream learning, our analysis could potentially be leveraged as a starting point for predicting downstream learning outcomes and diagnosing failure cases.
There are several limitations to the analysis conducted in this work.
Though we took care to interpret the analyses at multiple resolutions, it remains possible that the primary source of the correlations we observe is a confounder not accounted for.
The fact that we used models, data, environments, and results from prior work hopefully mitigates concerns regarding inapplicability and cherry-picking, but it could also be fruitful to consider going beyond artifacts from prior work, e.g. designing new probing tasks aimed at further elucidating properties of these representations relevant to robotics.
On that note: this paper is work in progress.
We plan to continue by generating data representing various kinds of distribution shifts in the downstream task environments and investigating whether the tools presented above can be used to engineer a transfer procedure that results in a robust policy.
We thank the reviewers in advance for any comments and feedback about this work.
A CKNOWLEDGMENTS
We gratefully acknowledge the developers and maintainers of software packages that enabled this research: NumPy (Harris et al., 2020), SciPy (Virtanen et al., 2020), PyTorch (Paszke et al., 2019), matplotlib (Hunter, 2007), and seaborn (Waskom, 2021).
A R EPRESENTATIONAL C ONSISTENCY : A DDITIONAL R ESULTS
[Figure]