Control-oriented Clustering of Visual Latent Representation

## Abstract

Abstract We initiate a study of the geometry of the visual representation space –the information channel from the vision encoder to the action decoder– in an image-based control pipeline learned from behavior cloning. Inspired by the phenomenon of neural collapse (NC) in image classification (Papyan et al., 2020 ) , we investigate whether a similar law of clustering emerges in the visual representation space. Since image-based control is a regression task without explicitly defined classes, the central piece of the puzzle lies in determining according to what implicit classes the visual features cluster, if such a law exists. Focusing on image-based planar pushing , we posit the most important role of the visual representation in a control task is to convey a goal to the action decoder; for instance, “rotate the object clockwise and push it northeast”. We then classify training samples of expert demonstrations into eight “ control-oriented ” classes –based on (a) the relative pose between the object and the target in the input or (b) the relative pose of the object induced by expert actions in the output– where one class corresponds to one relative pose orthant ( Repo ).
Across four different instantiations of the vision-based control architecture,
we report the prevalent emergence of control-oriented clustering (similar to NC) in the visual representation space according to the eight Repo s. Beyond empirical observation, we show such a law of clustering can be leveraged as an algorithmic tool to improve test-time performance when training a policy with a limited amount of expert demonstrations. Particularly, we pretrain the vision encoder using NC as a regularization to encourage control-oriented clustering of the visual features. Surprisingly, such an NC-pretrained vision encoder, when finetuned end-to-end with the action decoder, boosts the test-time performance by 10 % percent 10 10\% to 35 % percent 35 35\% in the low-data regime. Real-world vision-based planar pushing experiments confirmed the surprising advantage of control-oriented visual representation pretraining. 1 1 1 https://computationalrobotics.seas.harvard.edu/ControlOriented_NC/

## 1 Introduction

We use a toy example to (a) introduce the concept of neural collapse and the task of policy learning from expert demonstrations, and (b) synchronize readers from the respective communities.

#### Minimum-time double integrator

Consider a dynamical system known as the *double integrator*

$$ $\displaystyle\ddot{q}(t)=u(t),$ (1) $$

where $q\in{{\mathbb{R}}}$ is the position, and $u\in\mathbb{U}:=[-1,1]$ is the external control that decides the system’s acceleration. For an example, imagine $q$ as the position of a car and $u$ as how much throttle or braking is applied to the car. Let $x(t):=(q(t),\dot{q}(t))\in{{\mathbb{R}}^{2}}$ be the state of the system. Suppose the system starts at $x(0)=\chi$, and we want to find the optimal state-feedback policy that drives the system to the origin using *minimum time*. Formally, this is an optimal control problem written as

$$ $\displaystyle\min_{u(t)}\ T,\quad\text{ subject to }\ \ x(0)=\chi,\ \ x(T)=0,\ \ u(t)\in\mathbb{U},\forall t,\ \ \text{and~{}}\eqref{eq:double-integrator}.$ (2) $$

Problem ([2](#S1.E2)) admits a closed-form optimal policy (Rao & Bernstein, ):

$$ $\displaystyle u_{\star}=\pi_{\star}(x):=\begin{cases}+1&\text{if }\left(\dot{q}<0\text{ and }q\leq\frac{1}{2}\dot{q}^{2}\right)\text{ or }\left(\dot{q}\geq 0\text{ and }q<-\frac{1}{2}\dot{q}^{2}\right)\\ 0&\text{if }q=0\text{ and }\dot{q}=0\\ -1&\text{otherwise.}\end{cases}$ (3) $$

This optimal policy is *bang-bang*: it applies either full throttle or full brake until reaching the origin.

#### Behavior cloning

We now try to learn the optimal policy from data, using a strategy called *behavior cloning* (Torabi et al., ). We collect $3N$ samples (a.k.a. expert demonstrations) from the optimal policy $\pi_{\star}(x)$; each batch of $N$ samples have “label” $+1$, 0 0, and $-1$, respectively. Then we treat policy learning as a classification problem given input $x$ and output $u$. We design a six-layer MLP ($2\rightarrow 64\rightarrow 64\rightarrow 64\rightarrow 64\rightarrow 3\rightarrow 3$) and train it with the cross-entropy loss. We use $N=5000$. For the $u=0$ class, we repeat $N$ times the sample “$x=0$” to balance the dataset.

#### Geometry of the representation space

As expected, the MLP is able to learn a good policy, but this is not the purpose of our experiment. We instead look at the geometry of the three-dimensional feature space at the penultimate layer. Particularly, let $\{f^{+1}_{i}\}_{i=1}^{N}$, $\{f^{0}_{i}\}_{i=1}^{N}$, $\{f^{-1}_{i}\}_{i=1}^{N}$ be the three sets of feature vectors corresponding to controls $\{+1,0,-1\}$. We compute the per-class and global mean vectors:

$$ $\displaystyle\mu^{c}=\frac{1}{N}\sum_{i=1}^{N}f_{i}^{c},\ \ c=+1,0,-1;\quad\mu=\frac{1}{3}\left(\mu^{+1}+\mu^{0}+\mu^{-1}\right).$ (4) $$

Let $\tilde{\mu}^{c}:=\mu^{c}-\mu$ be the *globally-centered* class mean for each class $c$.
Fig. [1](#S1.F1) plots the class means $\tilde{\mu}^{c}$ and the globally-centered feature vectors $\tilde{f}^{c}_{i}:=f^{c}_{i}-\mu$ in different colors, as training progresses. We observe a clear *clustering* of the features according to their labels. Moreover, the clustering admits a precise geometry: (a) the lengths of globally-centered class means $\|\tilde{\mu}^{c}\|$ tend to be equal to each other for $c=+1,0,-1$, and (b) the angles spanned by pairs of class mean vectors $\angle(\tilde{\mu}^{c_{1}},\tilde{\mu}^{c_{2}})$ also tend to be equal to each other for $c_{1}\neq c_{2}$, shown by the perfect “tripod” in Fig. [1](#S1.F1) epoch 1000.

Figure: Figure 1: Per-class (red, blue, green) globally centered features (points) and mean vectors (black lines with $\star$ endpoints) from the penultimate latent space when cloning the optimal bang-bang policy ([3](#S1.E3)) from expert demonstrations. Numbers in the blue band represent the lengths of the three per-class mean vectors, and numbers in the red band represent the angles spanned by pairs of per-class mean vectors. The lengths and angles tend to be equal to each other as training progresses.
Refer to caption: /html/2410.05063/assets/x1.png

#### Neural collapse

The clustering phenomenon shown in Fig. [1](#S1.F1) was first observed by Papyan et al. () in image classification and dubbed the name *neural collapse* (NC). In particular, NC refers to a set of four manifestations in the representation space (i.e., the penultimate layer):

- (NC1)
Variability collapse: feature vectors of the same class converge to their class mean.
- (NC2)
Simplex ETF: globally-centered class mean vectors converge to a geometric configuration known as Simplex Equiangular Tight Frame (ETF), i.e., mean vectors have the same lengths and form equal angles pairwise (as shown in Fig. [1](#S1.F1)).
- (NC3)
Self-duality: the class means and the last-layer linear classifiers are self-dual.
- (NC4)
Nearest class-center prediction: the network predicts the class whose mean vector has the minimum Euclidean distance to the feature of the test image.

Since its original discovery, NC has attracted significant interests, both empirical (Jiang et al., ; Wu & Papyan, ; Rangamani et al., ) and theoretical (Fang et al., ; Han et al., ); see Appendix [D](#A4) for a detailed discussion of existing literature. The results we show in Fig. [1](#S1.F1) are just another example to reinforce the prevalence of neural collapse, because behavior cloning of the bang-bang optimal controller reduces to a classification problem.

#### Our goal

Does a similar *law of clustering*, in the spirit of NC, happen when cloning image-based control policies? We focus on the task of *vision-based planar pushing*, where an agent needs to push an object to a target position given only image observations, see Fig. [2](#S1.F2) for an example of pushing a T-shaped block. A general image-to-action architecture consists of (a) a *vision encoder* that embeds high-dimensional images as compact visual features (He et al., ; Oquab et al., ), and (b) an *action decoder* that generates control outputs given the latent vectors (Mandlekar et al., ; Chi et al., ). The entire pipeline is trained end-to-end using expect demonstrations.

Figure: Figure 2: Investigation of a law of clustering, similar to NC, in the visual representation space.
Refer to caption: /html/2410.05063/assets/x2.png

The bridge and information channel from vision to control is the visual representation space. We aim to study the geometry of the visual representation space from the lens of neural collapse. Particularly, we seek to answer two fundamental questions.

- (Q1)
Does NC happen in the visual representation space? In other words, do the latent features form into clusters (like Fig. [1](#S1.F1))? If so, according to which “classes” do they cluster?
- (Q2)
Is the extent to which neural collapse happens related to the model’s test-time performance? If so, can we transform neural collapse from a “phenomenon” to an “algorithmic tool”?

These two questions have never been answered before –either empirically or theoretically– because vision-based control is in stark contrast with existing literature in neural collapse for image classification. First, a prerequisite of NC is that the training data need to already be classified.
Planar pushing (and vision-based control in general), however, is a *regression* problem where the output is a continuous control signal. There is no supervision coming from classification whatsoever. Second, theoretical analysis of NC typically assumes a linear classifier from the representation space to the model output (Han et al., ; Fang et al., ; Jiang et al., ). This assumption is strongly violated in vision-based control because from visual representation space to control lies a rather nonlinear and complicated action decoder (see Fig. [2](#S1.F2)).

#### Our contribution

Despite the challenges mentioned, we empirically demonstrate that the answers to both questions are affirmative in vision-based planar pushing. Our contributions are:

- (C1)
Control-oriented clustering  A natural path to study NC for a regression problem is to give each sample a class label and check whether the visual features cluster according to the labels. Then, the nontrivial question becomes what should these classes be? Based on the posit that the visual representation should *convey a goal of control* for the action decoder, we design two “control-oriented” classification strategies. They compute (a) the relative pose between the object and the target in the input image space, or (b) the relative pose change of the object induced by the sequence of expert actions,
and classify the samples into 8 classes, each corresponding to one orthant in the relative pose space (called a Repo). We then demonstrate the prevalent emergence of neural collapse in the visual representation space, across four different architecture designs of the image-based control pipeline.
- (C2)
Visual representation pretraining with Neural Collapse
When the number of expert demonstrations is decreased, the strength of NC decreases, so does the model’s test-time performance. This motivates us to leverage NC as an algorithmic tool to improve the model’s performance under insufficient demonstrations. Indeed, we show that by using the control-oriented NC metrics as the loss function to pretrain the vision encoder, we obtain $10\%$ to $35\%$ boost in the model’s test-time performance. Real-world robotic manipulation experiments confirmed the advantage of NC pretraining.

#### Paper organization

We introduce the problem setup of planar pushing and control-oriented classification in §[2](#S2). We show the prevalence of neural collapse in the visual representation space across four different architectures in §[3](#S3) (that answers [(Q1)](#S1.I2.i1)). We describe our method of NC pretraining and show it improves model performance in §[4](#S4) (which answers [(Q2)](#S1.I2.i2)). We demonstrate real-world robotic experiments in §[5](#S5) and conclude in §[6](#S6). Appendices contain related work and extra results.

## 2 Control-oriented Classification

### 2.1 Vision-based Planar Pushing

Planar pushing (Lynch & Mason, ; Yu et al., ) is a longstanding problem in robotics, and consists of controlling a *pusher* to push a given *object* to follow certain trajectories or to a target position. This problem is fundamental because manipulating objects is crucial to deploy robots in the physical world (Bicchi & Kumar, ). Despite how effortlessly humans perform this task, planar pushing represents one of the most challenging problems for classical model-based control (e.g., optimization-based control (Wensing et al., )) due to the underlying hybrid dynamics and underactuation (i.e., the controller needs to plan where/when the pusher should make/break contact with the object and whether to slide along or stick to the surface of the object, leading to different “modes” in the dynamics (Goyal et al., ; Hogan & Rodriguez, )). Planar pushing is made even harder when the controller is only allowed to use images to make decisions, which is the vision-based planar pushing setup we will focus on.

#### Policy learning from demonstrations

Chi et al. () showed that by collecting a set of expert demonstrations and training a *diffusion policy* (DP) to imitate the expert, vision-based planar pushing can be reliably solved. DP consists of a ResNet to embed images in a latent space, and a diffusion model (DM) whose score function is parameterized by the latent representation to predict control actions (i.e., ResNet+DM in Fig. [2](#S1.F2)). Formally, we are given a collection of $N$ expert demonstrations ${\cal D}=\{D_{i}\}_{i=1}^{N}$ where each $D_{i}$ is a sequence of images and controls

$$ $D_{i}=(I_{0},u_{0},I_{1},u_{1},\dots,I_{l_{i}-1},u_{l_{i}-1},I_{l_{i}}),$ $$

with $I$ the image, $u$ the control, $l_{i}$ the length of $D_{i}$, and in the final image $I_{l_{i}}$ the object is pushed to the target position. From ${\cal D}$ we extract a set of $M$ training samples ${\cal S}=\{s_{t}\}_{t=1}^{M}$ where each sample $s_{t}$ consists of a sequence of $K$ images and a sequence of $H$ controls

$$ $\displaystyle s_{t}=(I_{t-K+1},\dots,I_{t-1},I_{t}\mid u_{t},u_{t+1},\dots,u_{t+H-1}).$ (5) $$

$M$ is usually much larger than $N$. The end-to-end policy is then trained on ${\cal S}$, which takes as input the sequence of images and outputs the sequence of controls. At test time, the policy is executed in a receding horizon fashion (Mayne & Michalska, ) where only the first predicted control $u_{t}$ is executed and then a new sequence of controls is re-predicted to incorporate feedback from the new image observations. In this work, we focus on the more general vision-based control architecture in Fig. [2](#S1.F2) with DP a concrete instantiation but not the only choice (e.g., we can replace ResNet with DINOv2 and diffusion model with LSTM).

#### Probing the visual representation space

Our goal is to probe and understand the visual representation space. Formally, given the training sample ([5](#S2.E5)), denote

$$ $\displaystyle f_{t}=\mathrm{VisionEncoder}(I_{t-K+1},\dots,I_{t-1},I_{t})$ (6) $$

as the visual feature of the sequence of images. We are interested in understanding whether the set of feature vectors ${\cal F}=\{f_{t}\}_{t=1}^{M}$, collectively, exhibit certain geometric properties similar to neural collapse, just like what we have shown in Fig. [1](#S1.F1).
A motivation to study this problem is, as observed in Chi et al. (); Kim et al. (); Team et al. (), *pretrained* ResNet features deliver poor performance when used for control (such as planar pushing) and end-to-end finetuning is necessary to adapt the vision features. This is mysterious and counterintuitive. The widespread success of computer vision has confirmed the ability of neural networks to learn general and rich visual features (Huh et al., ). Why are pretrained visual features insufficient to push a T block? In other words, calling ${\cal F}_{\text{pretrain}}$ the set of pretrained visual features, and ${\cal F}$ the set of finetuned visual features (containing those defined in ([6](#S2.E6))), the end-to-end “control-oriented” finetuning must have “reshaped” the visual representation space, but how?

### 2.2 Classification according to Relative Pose Orthants (Repos)

Figure: Figure 3: Control-oriented classification of training samples. Left: goal-based; Right: action-based.
Refer to caption: /html/2410.05063/assets/x3.png

The essence of neural collapse is a *law of clustering* in the representation space. To study this, we need to assign every feature vector $f_{t}$ a class. In image classification, this class is explicitly given. However, in vision-based control the output is a sequence of controls, what should the “class” be?

A natural choice is to perform $k$-means clustering of the output actions. Unfortunately, not only is this classification not interpretable, it also does not lead to observation of NC (see Appendix [A](#A1)).

We then conjecture that for a control task such as planar pushing, the role of the vision encoder is to convey a “control goal” to the action decoder from image observations. For example, looking at the left image in Fig. [3](#S2.F3), the vision encoder may set the goal of control to “push the T block southwest and rotate it counter-clockwise”.
Building upon this, we design two classification strategies:

- •
Goal-based classification   Given a sample $s_{t}$ as in ([5](#S2.E5)), we look at image $I_{t}$. We compute the relative pose of the target position with respect to the object. As depicted in Fig. [3](#S2.F3) left, this relative pose is a triplet $(x,y,\theta)$ containing a 2D translation and a 1D rotation.
- •
Action-based classification  Given a sample $s_{t}$ as in ([5](#S2.E5)), we look at the images induced by the action sequence $(u_{t},\dots,u_{t+H-1})$, i.e., $(I_{t},I_{t+1},\dots,I_{t+H})$. Importantly, it is possible that the expert did not finish her action plan at time $t+H$, which can be detected if the pusher is still in contact with the object. In those cases, we enlarge $H$ until the pusher breaks contact with the object, which forms a full episode of expert plan. Fig. [3](#S2.F3) right illustrates such an episode, with the object transparency diminishing as time progresses. In a full episode, we compute the relative pose of the object at time $t+H$ with respect to the object at time $t$, which consists of a triplet $(x,y,\theta)$ visualized in Fig. [3](#S2.F3) right.

The two classification strategies are correlated but different. Both of them measure the goal of the expert, but the first strategy measures the *long-term* goal while the second strategy measures the *short-term* goal (in an episode). In both cases, we obtain a relative pose triplet $(x,y,\theta)$ and we divide that 3D space into eight classes based on the signs of $x,y,\theta$. In other words, each class corresponds to one orthant of the space (called a *relative pose orthant*, or in short a Repo), as visualized in Fig. [3](#S2.F3) middle. A nice property is that the resulting classes are semantically interpretable!

###### Remark 1 (Finegrained Repo s) .

What will happen if the relative pose space is divided into a larger number of classes? In Appendix [B](#A2), we divide the relative pose space into $64$ and $216$ classes, and demonstrate that such a law of clustering still holds, albeit to a slightly weaker extent.

###### Remark 1 (Finegrained Repo s) .

## 3 Prevalent Neural Collapse in the Visual Representation Space

Using the control-oriented classification strategies described above, we are ready to study whether a law of clustering similar to NC emerges in the visual representation space. Given the set of visual features ${\cal F}=\{f_{t}\}_{t=1}^{M}$, we assign a class label $c\in[C]$ to each feature vector $f_{t}$ and denote it $f_{t}^{c}$.

#### Neural collapse metrics

We compute three metrics to evaluate [(NC1)](#S1.I1.i1) and [(NC2)](#S1.I1.i2). Define

$$ $\displaystyle\mu^{c}=\frac{1}{M_{c}}\sum_{i=1}^{M_{c}}f_{i}^{c},c=1,\dots,C,\quad\mu=\frac{1}{M}\sum_{c=1}^{C}\sum_{i=1}^{M_{c}}f^{c}_{i},$ (7) $$

as the class mean vectors and the global mean vector, respectively. Note that $M_{c}$ denotes the total number of samples in each class $c$ and $\sum_{c=1}^{C}M_{c}=M$. Then define $\tilde{\mu}^{c}:=\mu^{c}-\mu$ as the globally-centered class means, and $\tilde{f}^{c}_{i}:=f_{i}^{c}-\mu$ as the globally-centered feature vectors. Consistent with Wu & Papyan (), we evaluate [(NC1)](#S1.I1.i1) using the *class-distance normalized variance* (CDNV) metric that depends on the ratio of within-class to between-class variabilities:

$$ $\displaystyle\text{CDNV}_{c,c^{\prime}}:=\frac{\sigma_{c}^{2}+\sigma_{c^{\prime}}^{2}}{2\|\tilde{\mu}^{c}-\tilde{\mu}^{c^{\prime}}\|^{2}},\quad\forall c\neq c^{\prime},$ (8) $$

where $\sigma_{c}^{2}:=\frac{1}{M_{c}-1}\sum_{i=1}^{M_{c}}\|\tilde{f}^{c}_{i}-\tilde{\mu}^{c}\|^{2}$ is the within-class variation. Clearly, [(NC1)](#S1.I1.i1) happens when $\text{CDNV}_{c,c^{\prime}}\rightarrow 0$ for any $c\neq c^{\prime}$. We use a single number CDNV to denote the mean of all $\text{CDNV}_{c,c^{\prime}}$ for $c\neq c^{\prime}$. We evaluate [(NC2)](#S1.I1.i2) using the standard deviation ($\mathrm{STD}$) of the lengths and angles spanned by $\tilde{\mu}^{c}$ (calling $\mathrm{AVE}$ as the shortcut for averaging):

$$ $\displaystyle\text{STDNorm}:=\frac{\mathrm{STD}\left(\{\|\tilde{\mu}^{c}\|\}_{c=1}^{C}\right)}{\mathrm{AVE}\left(\{\|\tilde{\mu}^{c}\|\}_{c=1}^{C}\right)},\quad\text{STDAngle}:=\mathrm{STD}\left(\left\{\left\langle\frac{\tilde{\mu}^{c}}{\|\tilde{\mu}^{c}\|},\frac{\tilde{\mu}^{c^{\prime}}}{\|\tilde{\mu}^{c^{\prime}}\|}\right\rangle\right\}_{c\neq c^{\prime}}\right).$ (9) $$

Clearly, [(NC2)](#S1.I1.i2) happens if and only if both STDNorm and STDAngle become zero. We do not evaluate [(NC3)](#S1.I1.i3) and [(NC4)](#S1.I1.i4) because they require a linear classifier from the representation space to the output, which does not hold in our vision-based control setup.

Figure: Figure 4: Test scores w.r.t. training epoches of four different instantiations of the image-based control pipeline in Fig. [2](#S1.F2). In (a) and (b) we show test scores of three random seeds. In (c) and (d) we show test scores of a single seed because using LSTM as the action decoder leads to poor test-time performance, an observation that is consistent with Chi et al. (). This confirms the dominating advantage of using a diffusion model (DM) as the action decoder for planar pushing.
Refer to caption: /html/2410.05063/assets/x4.png

#### Simulation setup

We collect $N=500$ expert demonstration on a push-T setup shown in Fig. [2](#S1.F2). At each round, the object and target positions are randomly initialized and the same human expert controls the pusher to push the object into alignment with the target position (through a computer interface provided by pymunk (Blomqvist, )). This provides $M=55,480$ training samples. We train four different instantiations of the image-based control pipeline as shown in Fig. [2](#S1.F2). The four trained models are evaluated on a test push-T dataset with 100 tasks. We define our primary evaluation metric as the ratio of the overlapping area between the object and the target position to the total area of the target position.

#### Results

Fig. [4](#S3.F4) shows the evaluation scores of four different models. DINOv2 combined with a diffusion model (DM) attains the best performance around $70\%$. ResNet combined with DM (the original diffusion policy from Chi et al. ()) is slightly worse but quite close. When an LSTM is used as the action decoder to replace DM, the performance significantly drops, an observation that is consistent with Chi et al. () and confirms the advantage of using DM as the action encoder. For this reason, it is not worthwhile training models with LSTM from different random seeds.

Fig. [5](#S3.F5) plots the three NC evaluation metrics w.r.t. training epochs for two different models (both with DM) with both classification strategies described in §[2.2](#S2.SS2). We observe consistent decrease of three NC metrics as training progresses, suggesting the prevalent emergence of a law of clustering that is similar to NC. Notably, the CDNV metric in the case of using a goal-based classification is significantly smaller than that of an action-based classification, suggesting that the goal-based classification may correspond to a stronger level of clustering. Despite the poor test-time performance of trained models using LSTM, similar neural collapse is observed and shown in Appendix [B](#A2).

#### From pretrained to finetuned ResNet features

We now have a plausible answer to why related works observed that pretrained ResNet features deliver poor performance in control. ResNet is pretrained to perform image classification, and according to neural collapse, the pretrained ResNet features are clustered according to the class labels in image classification, such as dogs and cats. However, under our neural collapse observation in Fig. [5](#S3.F5), ResNet features for planar pushing are clustered according to “control-oriented” class labels that are related to the relative pose between the object and the target. Therefore, during finetuning, we conjecture the visual features have “re-clustered” according to the new task. We further verify this conjecture in the next section.

Figure: Figure 5: Prevalent emergence of neural collapse in the visual representation space. (a) Three NC metrics w.r.t. training epochs using ResNet as the vision encoder and diffusion model (DM) as the action decoder. (b) Three NC metrics w.r.t. training epochs using DINOv2 as the vision encoder and DM as the action decoder. All the plots show the mean and standard deviation (in shaded band) over three random seeds. The test scores are shown in Fig. [4](#S3.F4)(a)(b). Similar observations of NC hold when replacing DM with LSTM as the action decoder, shown in Appendix [B](#A2).
Refer to caption: /html/2410.05063/assets/x8.png

### 3.1 Neural Re-Collapse during Finetuning

Figure: Figure 6: Continual training of a (ResNet+DM) image-based control pipeline on five different planar pushing domains with changes in object color, shape, and visual distractors.
Refer to caption: /html/2410.05063/assets/figures/domain-visualizations/domain18_1_sample.png

We collect 500 expert demonstrations in each of the five planar pushing domains shown in Fig. [6](#S3.F6).
We then continuously train a (ResNet+DM) model across five domains, each with 300 epochs. Fig. [7](#S3.F7) shows the performance score and three NC metrics (according to the goal-based classification), where blue curves correspond to data in the target (new) domain and yellow curves correspond to data in the source (old) domain. We make three observations. (i) During retraining on the target domain, performance increases in the target domain but drops in the source domain. (ii) The NC metrics consistently decrease when retraining on the target domain (and increase or plateau on the source domain), suggesting the visual representation space is “re-clustered”. (iii) Contradictory to recent literature that observe *loss of plasticity* in continual learning (Dohare et al., ; Muppidi et al., ), there is no clear evidence the (ResNet+DM) model loses its ability to learn new tasks during continual finetuning (in every new domain, the model reaches around $70\%$ performance). Results using the action-based classification are similar and presented in Appendix [B](#A2).

Figure: Figure 7: Neural re-collapse in the visual representation space during continual learning across different domains. Blue curves represent data in the target (new) domain, and yellow curves represent data in the source (old) domain. All results use the goal-based classification. Results using the action-based classification are similar and shown in Appendix [B](#A2).
Refer to caption: /html/2410.05063/assets/x20.png

## 4 Visual Representation Pretraining with Neural Collapse

Our experiments in §[3](#S3) empirically demonstrated that an image-based control policy trained on sufficient amount of data exhibits an elegant control-oriented law of clustering in its visual representation space. We now study whether such clustering still holds under insufficient data.

#### Weak clustering

We use the same setup as in §[3](#S3), focus on the (ResNet+DM) model, and train it using only $100$ demonstrations. Fig. [8](#S4.F8) shows the test-time performance and three NC metrics according to the goal-based classification. Neural collapse appears to still hold, but to a much weaker extent, particularly when looking at the STDNorm metric.

Figure: Figure 8: Weak neural collapse under $100$ demonstrations. NC metrics are calculated using goal-based classification. See Appendix [B](#A2) for results using action-based classification.
Refer to caption: /html/2410.05063/assets/x36.png

#### Visual pretraining

If we knew a well-trained model exhibits strong clustering in its visual representation space, how about we encourage such a phenomenon even without sufficient data? To test this idea, we pretrain the vision encoder by explicitly minimizing the three NC metrics (([8](#S3.E8)) and ([9](#S3.E9))) prior to training the vision encoder jointly with the action decoder. Table [1](#S4.T1) presents the comparison of test-time performance between the baseline model and models whose vision encoder is pretrained to encourage control-oriented clustering. We observe a substantial improvement by using NC pretraining: the minimum improvement is around $10\%$ on the letter B, and the maximum improvement is around $35\%$ on the letter O (we intentionally carved out a piece to break symmetry). Note that all models are trained using only 100 expert demonstrations!

**Table 1: Pretraining a ResNet vision encoder by minimizing the NC metrics significantly improves test-time model performance on four domains corresponding to four letters in the word “ROBOT”.**
| Pushing tasks |  |  |  |  |
| --- | --- | --- | --- | --- |
| Baseline | $0.202$ | $0.423$ | $0.385$ | $0.241$ |
| NC-pretrained (goal-based) | $0.360$ | $0.720$ | $0.504$ | $0.347$ |
| NC-pretrained (action-based) | $0.378$ | $0.769$ | $0.488$ | $0.351$ |

## 5 Real-world Validation

The improvement shown in Table [1](#S4.T1) is astonishing and motivates us to further validate the effectiveness of NC pretraining in the real world.

Figure: Figure 9: Real-world setup.
Refer to caption: /html/2410.05063/assets/x40.png

#### Setup

Fig. [9](#S5.F9) depicts our real-world setup for vision-based planar pushing. The pusher is attached to the end-effector of a Franka Panda robotic arm to push a T block into the target position on a flat surface. An overhead Intel RealSense camera is used to provide image observations to the pusher. To enable NC pretraining in the real world, we need groundtruth relative pose of the object w.r.t. the target position. To provide reliable estimate of the relative pose, we attach an AprilTag (Olson, ) to the bottom of the T-block, which enables real-time accurate pose estimation. To collect real-world expert demonstrations, we follow Chi et al. () and use a SpaceMouse to teleoperate the robotic arm to push the T block. We collected $100$ demonstrations where the initial position of the T-block is randomized at each run, and this took around 10 hours. We trained two policies: (a) a baseline that is the (ResNet+DM) model trained directly using 100 demonstrations, and (b) an NC-pretrained model whose ResNet is first pretrained by minimizing NC metrics and then jointly trained with DM.

#### Results

We test both models on a set of 10 new push-T tasks where the T block is initialized at positions not seen during training. We intentionally tried our best to initialize the T block at exactly the same position for fair comparison of two models (the reader can check this in Fig. [10](#S5.F10)). The baseline model succeeded 5 out of 10 times, and the NC-pretrained model succeeded 8 out of 10. Fig. [10](#S5.F10) shows three tests comparing the trajectories of the baseline and the NC-pretrained models. The rest of the tests are shown in Appendix [C](#A3) and supplementary videos.

Figure: Figure 10: Three real-world tests. In every test, top row shows trajectories of NC-pretrained model and bottom row shows trajectories of baseline. The first and last column show the initial and final position of the object, respectively. The middle three columns show overlaid trajectories generated by the policies. Full results of 10 tests are shown in Appendix [C](#A3) and supplementary videos.
Refer to caption: /html/2410.05063/assets/figures/real_world_trajectory_plots/NC_encouraged_100demos/eval_demo_3_ver2.png

## 6 Conclusion

We presented an empirical investigation of the geometry of the visual representation space in an end-to-end image-based control policy for planar pushing. We demonstrated that, similar to the phenomenon of neural collapse in image classification, a prevalent control-oriented law of clustering emerges in the visual representation space. Further, this law of clustering can be effectively leveraged to pretrain the vision encoder for improved test-time performance.

Our research opens up lots of future research directions, such as theoretical analysis and justification of control-oriented clustering, extension of such investigation to image-based control policies trained from reinforcement learning, and connections to neuroscience and cognitive science.

## Acknowledgments

We would like to express our sincere appreciation to X.Y. Han for sharing valuable related works on Neural Collapse, and to Cheng Chi and Yilun Du for their insightful discussions on Diffusion Policy. We thank Yifeng Zhu, Shucheng Kang, and Yulin Li for their discussion and assistance with the Franka Panda robot arm setup. Additionally, Haocheng Yin acknowledges the Swiss-European Mobility Programme (SEMP) for their financial support, which partially covered travel expenses from ETH Zürich to Harvard University and living costs during the stay.

## Appendix A k 𝑘 k -Means Clustering of Expert Actions

Using the $500$ expert demonstrations collected as described in §[3](#S3), we attempt to perform a $k$-means clustering of the output actions directly to classify the training samples.

To decide what is the optimal number of clusters for $k$-means clustering, we follow the popular Silhouette analysis (Rousseeuw, ). Fig. [11](#A1.F11) shows the average Silhouette score when varying the number of clusters, and Fig. [12](#A1.F12) shows the Silhouette analysis for selected values of $k$. Based on the Silhouette analysis, we choose the optimal $k=5$.

Fig. [13](#A1.F13) then plots the three NC metrics calculated according to the classes determined from $k$-means clustering. We observe that although the STDAngle and STDNorm metrics decrease, the CDNV metric keeps increasing, providing a negative signal for neural collapse.

Figure: Figure 11: Average Silhouette score of $k$-means clustering of the output action with respect to number of clusters ($k$ value) from 3 to 19. The silhouette score is a metric used to evaluate the quality of clustering in $k$-Means clustering. The optimal $k$ value is 5, having the highest average Silhouette score.
Refer to caption: /html/2410.05063/assets/x41.png

Figure: Figure 12: Silhouette analysis for $k$-Means clustering on sample data with number of clusters = 3, 5, 8, 12, 16 (from left to right). The red dotted vertical line represents the average Silhouette score while scores of samples in different clusters are represented with different colors.
Refer to caption: /html/2410.05063/assets/x42.png

Figure: Figure 13: $k$-Means clustering of the output actions does not lead to observation of neural collapse, even using the optimal $k$ value ($k$ = 5). Data obtained from training a (ResNet+DM) model on $500$ expert demonstrations described in §[3](#S3).
Refer to caption: /html/2410.05063/assets/x47.png

## Appendix B Supplementary Experimental Results

### B.1 Observation of Neural Collapse with LSTM

In the main text, we demonstrated emergence of neural collapse when using diffusion model as the action decoder. Here we provide evidence of similar neural collapse when using LSTM as the action decoder, as shown in Fig. [14](#A2.F14).

Figure: Figure 14: Prevalent emergence of neural collapse in the visual representation space. (a) Three NC metrics w.r.t. training epochs using ResNet as the vision encoder and LSTM as the action decoder. (b) Three NC metrics w.r.t. training epochs using DINOv2 as the vision encoder and LSTM as the action decoder. The test scores are shown in Fig. [4](#S3.F4) (c)(d). Only one seed of results are shown because the test performance is poor when using LSTM as the action decoder. NC results when using diffusion model as the action decoder are shown in Fig. [5](#S3.F5).
Refer to caption: /html/2410.05063/assets/x51.png

### B.2 Neural Re-Collapse

In the main text, we showed the visual features re-collapse when continuously finetuned on different domains, using the goal-based classification.
Here we provide similar neural re-collapse observations when using the action-based classification, as shown in Fig. [15](#A2.F15).

Figure: Figure 15: Neural re-collapse in the visual representation space during continual retraining across different domains. All results use the action-based classification. Results using the goal-based classification are presented in main text Fig. [7](#S3.F7).
Refer to caption: /html/2410.05063/assets/x63.png

### B.3 Finegrained Repo s

In Remark [1](#Thmtheorem1) we stated that when increasing the number of classes used to divide the space of relative pose, neural collapse still holds. Here we provide more details.

Recall that the relative pose is a triplet $(x,y,\theta)$ containing a 2D translation and a 1D rotation. In the $8$-class division, we simply divide each dimension $(x,y,\theta)$ into two bins $(-\infty,0]$ and $[0,+\infty)$, leading to $2^{3}=8$ classes. To divide the space into more classes, we can divide each dimension into more bins. In particular, we try two extra options. First, we divide each dimension into $4$ bins, leading to $64$ classes. Second, we divide each dimension into $6$ bins, leading to $216$ classes. Fig. [16](#A2.F16) compares the NC metrics under $8$, $64$, and $216$ classes using the (ResNet+DM) model and the goal-based classification strategy. We can clearly observe that when the number of classes is increased, neural collapse still emerges. However, it appears that the CDNV metric can “bounce back” when the number of classes is increased. Similar observations hold for the (DINOv2+DM) model and shown in Fig. [17](#A2.F17).

Figure: Figure 16: NC metrics of finegrainded Repos for ResNet + DM model. (a-c) represent the three NC metrics w.r.t. training epochs using ResNet as the vision encoder and diffusion model (DM) as the action decoder, respectively for 8, 64, 216 classes of Repos.
Refer to caption: /html/2410.05063/assets/x79.png

Figure: Figure 17: NC metrics of finegrainded Repos for DINOv2 + DM model. (a-c) represent the three NC metrics w.r.t. training epochs using DINOv2 as the vision encoder and diffusion model (DM) as the action decoder, respectively for 8, 64, 216 classes of Repos.
Refer to caption: /html/2410.05063/assets/x88.png

### B.4 Weak Clustering

In Fig. [8](#S4.F8) we showed weak clustering when training using $100$ demonstrations according to the goal-based classification. Here Fig. [18](#A2.F18) shows weak clustering according to the action-based classification.

Figure: Figure 18: Weak neural collapse under $100$ demonstrations. NC metrics are calculated using action-based classification.
Refer to caption: /html/2410.05063/assets/x97.png

## Appendix C Real-World Experimental Results

Here we provide full results for real-world experiments.

- •
100 demonstrations  In the main text we showed results for 3 tests. Fig. [21](#A3.F21) and [22](#A3.F22) show the results for all 10 tests. As seen, the NC-pretrained policy succeeded 8 out of 10, but the baseline policy succeeded 5 out of 10.
- •
50 demonstrations  In addition, we also trained two policies using only 50 demonstrations. Fig. [19](#A3.F19) and [20](#A3.F20) show the results for all 10 tests. The NC-pretrained policy succeeded 4 out of 10, but the baseline policy succeeded 2 out of 10.

For video recordings of the test results, please consult the supplementary material.

Figure: Figure 19: Real-world evaluation for test 0 to test 4. In every test, top row shows trajectories of NC-pretrained model and bottom row shows trajectories of baseline. Both models are trained under 50 training demonstrations. The first and last column show the initial and final position of the object, respectively. The middle three columns show overlaid trajectories generated by the policies.
Refer to caption: /html/2410.05063/assets/figures/real_world_trajectory_plots/NC_encouraged_50demos/eval_demo_0.png

Figure: Figure 20: Real-world evaluation for test 5 to test 9. In every test, top row shows trajectories of NC-pretrained model and bottom row shows trajectories of baseline. Both models are trained under 50 training demonstrations. The first and last column show the initial and final position of the object, respectively. The middle three columns show overlaid trajectories generated by the policies.
Refer to caption: /html/2410.05063/assets/figures/real_world_trajectory_plots/NC_encouraged_50demos/eval_demo_5.png

Figure: Figure 21: Real-world evaluation for test 0 to test 4. In every test, top row shows trajectories of NC-pretrained model and bottom row shows trajectories of baseline. Both models are trained under 100 training demonstrations. The first and last column show the initial and final position of the object, respectively. The middle three columns show overlaid trajectories generated by the policies.
Refer to caption: /html/2410.05063/assets/figures/real_world_trajectory_plots/NC_encouraged_100demos/eval_demo_0.png

Figure: Figure 22: Real-world evaluation for test 5 to test 9. In every test, top row shows trajectories of NC-pretrained model and bottom row shows trajectories of baseline. Both models are trained under 100 training demonstrations. The first and last column show the initial and final position of the object, respectively. The middle three columns show overlaid trajectories generated by the policies.
Refer to caption: /html/2410.05063/assets/figures/real_world_trajectory_plots/NC_encouraged_100demos/eval_demo_5.png

## Appendix D Related Work

#### Neural Collapse

The phenomenon of neural collapse (NC) is firstly introduced in Papyan et al. (), about an elegant geometric structure of the last-layer feature and classifier for a well-trained model in classification tasks. While the initial NC is observed for model trained with cross entrophy (CE) loss, Han et al. () demonstrated NC also holds when using the Mean Squared Error (MSE) loss for training. With the simplified assumption that only considers the last-layer optimization, neural collapse is proved to be the phenomenon occurring when training achieves global optimality, trained with CE loss function (E & Wojtowytsch, ; Zhu et al., ), and MSE loss function (Zhou et al., ; Tirer & Bruna, ). Later, NC is studied in more diverse situations, e.g., NC in imbalanced training (Fang et al., ; Hong & Ling, ; Dang et al., ), NC through intermediate layers (Rangamani et al., ; Parker et al., ), NC in transfer learning (Galanti et al., ; ), and relationship between NC and optimization methods (Xu et al., ; Rangamani & Banburski-Fahey, ). Since NC is a critical indicator of model robustness, it is also studied as a tool to improve the model (Bonifazi et al., ; Liu et al., ; Yang et al., ). Recently, neural collapse is studied in more complex applications, like large language model training (Wu & Papyan, ; Jiang et al., ). These works study neural collapse in the classification tasks and gain insightful results about the relationship between NC and training deep neural networks. In this work, we study NC of the visual representation space in an image-based control pipeline and show that control-oriented classification plays a key role in identifying NC.

#### Behavior Cloning

Behavior Cloning (BC) is a fundamental approach in Imitation Learning, where an agent learns to mimic expert demonstrations (Argall et al., ; Hussein et al., ). It typically involves collecting expert demonstrations and training a supervised learning model to reproduce the demonstrated behavior. While conceptually simple, BC has proven effective in many scenarios (Muller et al., ), particularly when combined with data augmentation techniques (Laskey et al., ; Peng et al., ) or iterative refinement (Ross & Bagnell, ; Ho & Ermon, ). This method has been widely applied across various robotic tasks including arm manipulation (Sharma et al., ; Rajeswaran et al., ), autonomous driving (Bojarski, ; Codevilla et al., ), and game playing (Silver et al., ; Vinyals et al., ). However, BC faces challenges such as compounding errors (Ross et al., ) and struggles with long-horizon tasks (Sun et al., ). Recent advancements have focused on addressing these limitations through improved data collection strategies (Mandlekar et al., ), robust loss functions  (Reddy et al., ; Brown et al., ), and integration with other learning paradigms (Rajeswaran et al., ; Zhu et al., ). In this work, we study BC from a representation perspective, and demonstrated that representation regularization, using neural collapse, can effectively improve the performance of BC in the low-data regime.

#### Diffusion Policy

The advent of diffusion models in imitation learning, exemplified by Diffusion Policy (Chi et al., ), has introduced a groundbreaking approach to behavior cloning for robot control. This method leverages the principles of denoising diffusion probabilistic models (Ho et al., ) to learn a policy that progressively refines random noise into expert-like actions. By doing so, it offers a powerful framework for modeling complex, multi-modal action distributions, which is particularly beneficial in robotics tasks involving high-dimensional action spaces (Janner et al., ). Diffusion policy has demonstrated impressive results in challenging robotics tasks, effectively addressing longstanding issues in traditional behavior cloning approaches such as distribution shift and long-horizon planning (Wang et al., ). As the field progresses, diffusion policy stands as a promising direction for enhancing the capabilities and robustness of imitation learning in complex robotic systems. In this work, we show that probing the geometry of the visual representation space helps understand and improve diffusion policy.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Can you validate these findings on more visually complex domains?
soundness: 3
strengths: To my knowledge, this is the first paper to study neural collapse in control. I appreciated some of the unconventional writing choices. The introduction had a nice unification of optimal control and behavior cloning and made the relevance of neural collapse to control very apparent. The experimental approach was really unique, creative, and insightful. The continual learning findings were especially interesting from my perspective.
summary: This paper studies neural collapse for control tasks. To facilitate this study, the authors learn policies with behavior cloning end-to-end. To probe the learned representations, they define a heuristic classification of features learned within the encoder called REPOs that categorize the position and orientation of the object in the image relative to the target object (this is specific to planar push tasks). They find that pretraining the vision encoder using neural collapse as a regularization technique improves performance on a simulated and real-world push-t task. They also find that, when the model is continually fine-tuned on different planar push tasks, the visual representation space undergoes re-clustering instead of saturating on early training tasks.
weaknesses: The paper could benefit from a more concise abstract and a more broad introduction that provides stronger motivation for the problem and clearer articulation of the paper's contributions. I like the unconventional approach to the abstract, but it would be more clear to the audience if the toy experiment appeared as a preliminary to the method section. 

I took issue with the discussion of resnet features throughout the paper. 
- First in response to L213: pre-trained resents are still useful compared to randomly initialized models, and second, that fine-tuning is often required for many transfer learning tasks, so I don't find it surprising that we need to fine-tune imagenet trained resnets for control domains.
- Because the image input of push-t is so limited it's unfair to make broad claims about the utility of resnet features. The benefit of resnet features will be more apparent in more visually challenging settings (e.g., cluttered table-top manipulation). 

Overall, as a study of visual representations the findings need to be verified on more visually complex domains.


Typo: "Expect" instead of "expert" on line 118.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: - In Table 1, is the performance improvement from access to privileged information (ground truth state), or from the NC regularization itself? It would be good to see a baseline where the orthant class label is provided as input to the policy alongside the image embedding, but no NC regularization is used.
- Could you clarify what constitutes neural collapse when we look at CDNV, STD Angle, STD Norm? In line 293-304, it is stated that these metrics should tend to zero as an indicator of neural collapse. However, in the figures below, none of the metrics seem to tend to zero.
- How sensitive is the observation of neural collapse to the choice of clustering? The hypothesis that the visual representation conveys a goal to the action decoder seems arbitrary. For example, one could also posit that the visual representation encodes the current environment state (block and EE pose). What happens if we cluster on environment state? Do we still observe neural collapse? What if we assign class labels uniformly at random? It would be convincing to see these comparisons as a justification for this choice.
- If the core claim is that neural collapse happens for visual encoders trained end-to-end in control tasks, it could perhaps be clearer and more convincing to apply this first to discrete control, and show that neural collapse in visual representation occurs across various discrete control environments like CartPole, etc.
soundness: 3
strengths: - The study of neural collapse in the visuomotor control setting is novel and well-motivated.
- The minimum-time double integrator example is very clear.
- Sim and real experiments show that neural collapse regularization objectives improve downstream task performance.
- Fine-tuning experiments show that neural collapse happens when transferring across domains.
summary: This paper investigates whether the phenomenon of neural collapse, previously observed in image classification and other tasks, appears in visual representations trained for continuous control. The paper considers the setup of training a visual encoder end-to-end jointly with a behavioral cloning policy on a planar pushing task. To study it as a classification problem, the frames are binned by relative poses to future states. The paper shows that during end-to-end training, the visual representation undergoes neural collapse. Finally, the neural collapse metrics are used as a regularization objective to improve downstream behavioral cloning performance.
weaknesses: - It seems unclear if the downstream improvement is from the NC regularization itself, or from access to ground truth state (by proxy of the class labels).
- The criteria for determining whether neural collapse has occurred seems unclear in this setting.
- Experiments are mostly on the planar pushing task.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 7. Will the results shown for planar pushing generalize for other robotic manipulation tasks such as Lifting Cube or Stacking cubes? There are several simulators such as Robosuite that provide with demonstrations of such task -- can REPO be extended to these non-planar tasks in a trivial manner or is it a future avenue of research?

8. Some results in Sec 3.1 (finetuning experiments), are not new. Observation (i) is stated as "During retraining on the target domain, performance increases in the target domain but drops in the source domain." -- Isn't this true for deep learning networks in general? 

9. Why aren't DinoV2 features also used for Table 1 (i.e pretraining DinoV2 to minimize NC metrics) and using that for control policies?

10. What is the Baseline model architecture that is used in Table 1?
soundness: 3
strengths: 1. I like the interpretable way of converting a regression task (in the context of behavior cloning) of action prediction to (x, y, $\theta$), which can be discretized.

2. It is interesting to see that the authors have included experiments on a real robot for the task of planar pushing and especially the boost that pretraining to minimize NC metrics gives compared to the baseline.
summary: This paper investigates whether visual features in image-based robotic control systems exhibit neural collapse (NC) - a clustering phenomenon previously observed in image classification tasks where features group into distinct classes. The authors find that on robotic control task of planar pushing, the visual features do indeed cluster into eight "control-oriented" classes based on relative pose between objects and targets. Most importantly, they show this clustering behavior can be leveraged to improve robot performance by up to 35% in low-data scenarios by pretraining the visual encoder to encourage such clustering, demonstrating benefits in both simulated and real-world robotic pushing tasks.
weaknesses: 3. The introduction can have more motivation of what NC is and why and where is this phenomenon observed rather than elaborating too much on the toy example. For instance, in the **Our goal** paragraph (L112), the authors raise the question "Does a similar law of clustering, in the spirit of NC, happen when cloning image-based control policies?" However, it is not very clear at this point on why this clustering is important to begin with? What happens if my control policies are working well and aren't clustered in the spirit of NC? Why should the reader care. Specifically it'd be good if authors can add the following:
	- Explicitly state potential benefits of finding NC in control policies (e.g. improved generalization, more interpretable features)
	- Expand on why clustering in the visual representation space might be important for control policies, even if they are already working well

4. In addition to this the related works is pushed to the appendix giving the reader a hard time to contextualize the current work. I'd highly suggest rewriting the related works as well as the motivation/introduction to address #1.

5. Fig 5 is really hard to read. Can't the plots for the 4 methods (ResNet/DinoV2 + LSTM/DM) be combined into a single plot for each of the 3 metrics (CDNV, STD Angle, STD Norm), resulting in 3 plots? 

6. Can the authors visualize (x, y $\theta$) for the 4 methods during testing policy rollouts in a 3D space as that was the main driving point for proposing REPO (L260)?

----
**Rationale for voting**: As of now, because of lack of clarity and potential re-writing on the motivation of this work (points #3, 4) and lack of clear experimental results that interpret the pretrained models with minimizing NC metrics (point #6) and organization of the results (point #5) -- I currently believe that this paper needs more work and hence am voting for a reject. However, my final decision would be based on other reviewers' comments and authors' rebuttal.
---
**Post rebuttal update**: Based on the authors rebuttal, I have increased my score to 6 as they have addressed all my major concerns regarding their paper.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: Why does the clustering phenomenon emerge for the planar pushing task with a continuous control signal? The toy example has discrete states in the closed-form optimal policy, so one can see why the clustering phenomenon might arise. The planar pushing task presents a different scenario where this emergence seems less intuitive.
soundness: 4
strengths: 1. Novel and interesting results: The paper presents intriguing and slightly surprising results, showing that control-oriented clustering in visual representations emerges even though the planar pushing task involves continuous control signals. Additionally, it demonstrates that pretraining the visual encoder improves test-time performance when training with limited expert demonstrations.
2. Excellent presentation and writing: The presentation of this paper is clear and engaging. The authors begin with a motivating toy example that introduces the phenomenon of NC, then proceed to answer two key questions: (a) Does NC occur in the visual representation space? and (b) Is the extent of NC related to model performance? The arguments are supported with convincing evidence, and the writing is very easy to follow.
3. Real-World Validation: The paper includes real-world experiments on vision-based planar pushing to show the effectiveness of NC pretraining

Overall, this paper presents novel and interesting results, with clear presentation and strong supporting evidence. This paper also opens up lots of future research directions in visual representation learning for robotics control, which will greatly benefit the community. This paper is clearly above the acceptance bar for ICLR.
summary: This paper investigates the geometry of the visual representation space in an image-based control pipeline trained through behavior cloning, focusing specifically on planar pushing tasks. It explores whether a phenomenon similar to neural collapse (NC), originally observed in image classification, can emerge in a regression context such as control tasks. The authors introduce "control-oriented" classes based on the relative poses of objects and targets, or relative poses of objects at different time steps, leading to the following key findings:
1. Control-oriented clustering emerges in the visual representation space, with latent features forming clusters according to control-relevant classes.
2. Using NC as a pretraining strategy for the vision encoder improves test-time performance, particularly when training data is limited.
3. Real-world experiments validate the effectiveness of NC pretraining in enhancing control performance.
weaknesses: 1. Scalability of the proposed approach for more complex tasks: Leveraging NC for control requires manually specified class definitions. In the planar pushing task, these classes are defined using the ground-truth relative pose of the target position with respect to the object, or the relative pose of the object at time t+H with respect to the object at time t. These classes are easy to define because the T-pushing task involves a single object, and its pose can be easily estimated using an AprilTag attached at the bottom of the T. However, the same class definition and pre-training strategy may not apply to more complex real-world tasks. In real-world robotic applications with diverse objects, accurate pose estimation can be challenging, particularly for deformable objects where defining a canonical pose is difficult. While accurate pose estimation is not the focus of this paper, it would be beneficial to discuss the potential limitations of the proposed class definitions in scaling to more complex tasks, or consider possible alternative class definitions.
2. Limited pretraining details and ablation studies: The paper lacks details about the visual pre-training, such as training datasets and training schedules. Moreover, there are no ablation studies provided on the relative importance of the two NC metrics in pretraining. Without these details, reproducing the results from the paper is challenging. Including such information would improve the reproducibility and robustness of the findings.

Meta-review: The reviewers are generally enthusiastic about this work, highlighting several strengths and providing positive ratings. Notable strengths include being the first paper to study neural collapse in control, employing a unique, creative, and insightful experimental approach, and conducting real robot experiments. Additionally, the presentation of this paper is clear and engaging. The paper received 8,8,8,6 ratings. The AC checked the paper, the reviews and the author responses and the discussions. The AC is in agreement with the reviewers and recommends acceptance.

Decision: Accept (Spotlight)