Control-oriented Clustering of Visual Latent Representation
Authors: Han Qi
Haocheng Yin
Heng Yang
Venue: ICLR 2025 Spotlight (2025)
Keywords: neural collapse
learning from demonstration
vision-based learning control
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We discover a control-oriented law of clustering (similar to neural collapse) in the visual representation space of an end-to-end image-based control pipeline trained from behavior cloning.

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, we empirically demonstrate the prevalent emergence of a similar *law of clustering* in the visual representation space. Specifically, 

- In discrete image-based control (e.g., Lunar Lander), the visual representations cluster according to the natural discrete action labels;

- In continuous image-based control (e.g., Planar Pushing and Block Stacking), the clustering emerges according to ``control-oriented'' classes that are 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. Each of the classes corresponds to one relative pose orthant (REPO).

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 limited 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% to 35%. Real-world vision-based planar pushing experiments confirmed the surprising advantage of control-oriented visual representation pretraining.

Published as a conference paper at ICLR 2025
CONTROL-ORIENTED CLUSTERING OF VISUAL LATENT REPRESENTATION
Han Qi" Haocheng Yin and Heng Yang!
School of Engineering and Applied Sciences, Harvard University *Department of Computer Science, ETH Zürich
ABSTRACT
We initiate study of the geometry of the visual representation space -th infor- mation channel from the vision encoder the action decoder-in 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 empirically demonstrate the prevalent emergence of a similar law of clustering in the visual representation space. Specifically,
In discrete image-based control (e.g., Lunar Lander), the visual representa- tions cluster according to the natural discrete action labels;
In continuous image-baset control (e.g.. Planar Pushing and Block Stack- ing), the clustering emerges according to "conttol-orienteed" classes that are based (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. Each of the classes corresponds one relative pose orthant (REPO).
Beyond empirical observation, we show such law of clustering can be leveraged as an algorithmic tool to improve test-time performance when training a policy with limited expert demonstrations. Particularly, pretrain the vision encoder using NC as a regularization to encourage control-oriented clustering of the vi- sual features. Surprisingly, such NC-pretrained vision encoder, when finetuned end-to-end with the action decoder, boosts the test-time performance by 10% to 35%. Real-world vision-based planar pushing experiments confirmed the surpris- ing advantage of control-oriented visual representation pretraining.
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 dynamical system known as the double integrator
q(t) u(t),
(1)
where q R is the position, and 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 the system starts at (0) = X. and we want to find the optimal state-feedback policy that drives the system the origin using minimum time. Formally, this is an optimal control problem written
the car. Let xt)) = (q(t),
E TR2 be the state of the system. Suppose
min T, subject to z00) X. x(T) 0, u(t) E U, Vt, and (1).
u(t)
Problem (2) admits closed-form optimal policy (Rao & Bernstein, 2001):
+1 if (q and q S 42) or (ğ and g 42))
u, = (x)
q = and =
otherwise.
This optimal policy is bang-bang: it applies either full throttle or full brake until reaching the origin.
"Equal contribution
I Work done during visit at the Harvard Computational Robotics Lab 'htpps: / httpass//ccccepee httpptattttttcccccccccccc/cc...t lonal robot seas. ha rvaaad edu edd//ontral /oo Oriented_Nc

---

Published as a conference paper at ICLR 2025
Epoch
Epoch
Epoch
Epoch 1000
00.3 10.21
9.2 5.6
1
27
177 7 176'
167 33" 161'
151' 82" 127"
135' 110"
115"
128' 114" 119"
Figure 1: Per-class (red, blue, green) globally centered features (points) and mean vectors (black lines with * endpoints) from the penultimate latent space when eloning the optimal bang-bang pol- icy (3) 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.
Behavior cloning We now try to learn the optimal policy from data, using strategy called be- havior cloning (Torabi al., 2018). We collect 3.V samples (a.k.a. expert demonstrations) from the optimal policy (x)): each batch of N samples have "label" +1, 0, and 11, respectively. Then we treat policy learning as a classification problem given input x and output u. We design six-layer MLP (2 + 64 + 64 64 64 - 3 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 The MIP 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 {ff" y (1 Ji {f, ( be the three sets of feature vectors corresponding to controls {+1,0,- -13. We compute the per-class and global mean vectors:
1 (nt 1 + ne + +1²²)
M c= +1,0, p=
Let 1tth := == μi be the globally-centered class mean for each class c. Fig. plots the class means A and the globally-centers feature vectors fr := f in different colors, as training progresses. We observe clear clustering of the features according to their labels. Moreover, the clustering admits precise geometry: (a) the lengths of globally-centreed class means || tend to be equal to each other for e +1, O, 11, and (b) the angles spanned by pairs of class mean vectors _(ji, 77x) also tend to be equal to each other for C1 * shown by the perfect "tripod" in Fig. epoch 1000.
Neural collapse The clustering phenomenon shown in Fig. was first observed by Papyan et al. (2020) in image classification and dubbed the name neural collapse (NC). In particular, NC refers to set of four manifestations in the representation space (i.e., the penultimate layer):
(NCI) 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 Simplex Equiangular Tight Frame (ETF), i.e., mean vectors have the same lengths and form equal angles pairwise (as shown in Fig. 1).
(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 the feature of the test image.
Since its original discovery, NC has attracted significant interests, both empirical (Jiang al., 2023; Wu & Papyan, 2024; Rangamani et al., 2023) and theoretical (Fang et al., 2021; Han et al., 2021); Appendix I for detailed discussion of existing literature. The results show in Fig. are just another example to reinforce the prevalence of neural collapse, because behavior cloning of the bang-bang optimal controller reduces to classification problem.
Does similar law of clustering, in the spirit of NC, happen when cloning image- based control policies? Two motivations underlie this question. First, understanding the structure of learned representation has been fundamental pursuit towards improving the interpretability of deep
Our goal
2

---

Published as a conference paper at ICLR 2025
learning models. While the discovery of NC has deepened understanding of visual representa- tion learning for classification, such am understanding is missing when using visual representation for decision-makingg Second, just related work has shown the benefits of NC for generalization and robustness (Bonifazi et al., 2024; Liu et al., 2023; Yang et al., 2022), we seek to uncover new algorithmic tools that improve model performance by "shaping" the latent representation space.
Towards this goal, we consider a general image-to-action architecture consisting of (a) a vision en- coder that embeds high-dimensional images compact visual features (He et al., 2016; Oquab al., 2023), and (b) an action decoder that generates control outputs given the latent vectors (Mandlekar al., 2021: Chi al., 2023). The entire pipeline is trained end-to-end using expert demonstra- tions. We study the instantiation of this urchitecture in three tasks: Lunar Lander from OpenAI Gym (Brockman, 2016), Planar Pushing that is popular in roboties (Chi et al., 2023), and Block Stacking from MimieGen (Mandlekar et al., 2023). The first is a discrete control task, while the second and the third are continuous control tasks. Fig. 2 overviews the architecture and the tasks.
Input: Images
Discrete Control
Continuous Control
Lander
Object
rippor
Oljjett
Target Posltion
Targer
Poeition
Vision Encoder
Visual Representation Space
Action Decoder
MUP
ResNNet
OR
OR
LSTM
OR
DINOV2
Difiusion Mociel
Output: Control
Pushing
Stacking
Figure Investigation of law of clustering, similar to NC, in the visual representation space.
The bridge and information channel from vision to control is the visual representation space. We aim study the geometry of the visual representation space from the lens of neural collapse. Particularly, we seek to answer two fundamental questions.
(Q1) Does a clustering phenomenon similar NC happen in the visual representation space? 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, transform neural collapse from "phenomenon" an "algorithmic tool"?
These two questions have never been answered before either empirically or theoretically be- cause vision-based control is stark contrast with existing literature in neural collapse for image classification. First, prerequisite of NC is that the training data need to already be classified. Con- tinuous vision-basee control tasks (like planar pushing and block stacking), however, are regression problems where the output is continuous control signal. There is supervision coming from classification whatsoever. Second, theoretical analysis of NC typically assumes a linear classifier from the representation space to the model output (Han al., 2021; Fang al., 2021; Jiang al.,

---

Published as a conference paper at ICLR 2025
2023). This assumption is strongly violated in vision-based control because from visual representa- tion space to control lies a rather nonlinear and complicated action decoder (see Fig. 2).
Our contribution Despite the challenges mentioned, empirically demonstrate that the answers to both questions are affirmative in iision-based control tasks. Our contributions are:
(CI) Control-oriented clustering We first show that, on the discrete-control task Lunar Lander, visual latent features cluster according to the discrete action labels. This forms the natural vision-based generalization of the double integrator example shown in Fig. 1. We then primarily focus continuous-control tasks. A natural path to study NC for regression problem 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 REPO). We then demonstrate the prevalent emergence of neural collapse in the visual representation space, in both Planar Pushing and Block Stacking.
(C2)
Control-oriented visual representation pretraining When the number of expert demon- strations is decreased, the strength of NC decreases, so does the model's s test-time perfor- mance. This motivates us to leverage NC as an algorithmic tool to improve the model's 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 Planar Pushing experiments confirmed the advantage of NC pretraining.
Paper organization We first show observation of neural collapse the discrete-control task Lunar Lander in $2. We then introduce the problem setup of continuous vision-based control tasks, control- oriented classification, and the prevalence neural collapse the visual representation space for continuous-comi control tasks in $3. We describe our method of NC pretraining and show it improves model performance in $4. We demonstrate real-world robotic experiments $5 and conclude in $6.
NEURAL COLLAPSE IN DISCRETE VISION-BASED CONTROL
As transition from our bang-bang policy in $1 to more challenging continuous vision-based control tasks, we first study whether the visual representation clusters according to the given discrete action labels in a discrete vision-based control task, Lunar Lander (Brockman, 2016). nutshell, Lunar Lander is task where the lander aims to reach a given target position by deciding between four discrete actions (see Fig. 2). We first train a performant state-based policy using reinforcement learning (RL) and then use the RL expert to collect vision-based demonstrations for behavior cloning (BC). The BC pipeline uses ResNet18 as the vision encoder and an MLP as the action decoder, supervised by cross-entropy loss with the expert demonstrations, Appendix A provides more details.
Since the action space for Lunar Lander discrete with 4 actions, we directly study the clustering phenomenon of the latent features according to these 4 actions. Suppose the training dataset contains M samples of images (input) and actions (output), and denote the set visual features as F {ff MM where each feature fi is computed by passing the input images through the ResNet18 encoder. We assign the action label E [4] to each feature fe and denote it ff-
Neural collapse metrics
We compute three metrics evaluate (NC1) and (NC2). Define
Me
C
Me
ff C,
f
(5)
M
M
as the class mean vectors and the global mean vector, respectively. Note that M, denotes the total number of samples in each class and Ee= M = M. Then define ji == μt- 1 as the globally- centered class means, and je : =f -1= the globally-centered feature vectors. Consistent with Wu
4

---

Published as a conference paper at ICLR 2025
& Papyan (2024), we evaluate (NC1) using the class-distance normalized variance (CDNV) metric that depends on the ratio of within-clase to between-class variabilities:
oe + oz CDNV o. 2
Vc
p 1
where of M,-I Ei 12 is the within-clas variation. Clearly, (NCI) happens when CDNV cue + for any c é. We use a single number CDNV to denote the mean of CDNV for e. & d. We evaluate (NC2) using the standard deviation (STD) of the lengths and angles spanned by je' (calling AVE as the shortcut for averaging):
(K
STD ((1))))))) STDNorm = AVE ((1))))))))
je ||
STDAngle == STD
Clearly, (NC2) happens if and only if both STDNorm and STDAngle become zero. We do not evaluate (NC3) and (NC4) because they require a linear classifier from the representation space to the output, which does not hold in vision-based control setup.
Results Fig. plots the three NC evaluation metrics w.r.t. training epochs for Lunar Lander. We observe consistent decrease of three NC metries training progresses and approaching the end of training, suggesting strong control-oriented clustering in the visual representation space.
60
0.6
50
0.5
40
Angle
0.4
Norm
0.20
30
0.3
0.15
20
0.0.
0.2
0.05
200
300
500
0.0
100
700
300
500
600
400
500
Epoch
Epoch
Figure 3: Emergence of neural collapse in the visual representation space for Lunar Lander. From left to right, it shows three NC metrics w.r.t. training epochs using ResNet18 as the vision encoder and an MLP the action decoder, with four discrete actions classification labels.
3 NEURAL COLLAPSE IN CONTINUOUS VISION-BASER CONTROL
We now consider continuous vision-based control tasks, where an obvious set of class labels does not exist anymore. We focus two tasks: Planar Pushing and Block Stacking.
Planar pushing (Lynch & Mason, 1996; et al., 2016) is longstanding problem robotics, and consists of controlling pusher to push given object follow certain trajectories position. This problem fundamental because manipulating objects is crucial to deploy robots in the physical world (Bicchi & Kumar, 2000). Despite how effortlessly humans perform this task, planar pushing represents one of the most challenging problems for classical model-based control (e.g., optimization-baseed control (Wensing al., 2023)) due to the underlying hybrid dynamics and underactuation (i.c., the controller needs to plan where/when the pusher should make/break contact with the object and whether to slide along stick to the surface of the object, leading to different "modes" in the dynamics (Goyal et al., 1991; Hogan & Rodriguez, 2020)).
target
Block stacking (Mandlekar al., 2023; Zhu al., 2020) classical manipulation task, aiming to stack one block onto another target block. This involves lifting up a block, moving the block above the target, rotating the block to match the angle the target, and putting the block down.
For both tasks, we describe how to train basend control policy from expert demonstrations.
Policy learning from demonstrations We are given collection of N expert demonstrations D (D.) where each D; is sequence of images and controls
D; (To, uo I, un...... ... J1-11 W1 11 111))
5

---

Published as a conference paper at ICLR 2025
with / the image, the control, 4 the length of D,. and in the final image It the object is moved to the target position. From D we extract a set of M training samples { where each sample 8, consists of a sequence of K images and sequence of H controls
5t-
4, We, ur+1r Mc+H-1).
M is usually much larger than N. The end-to-end policy is then trained on S, which takes input the sequence of images and outputs the sequence of controls. At test time, the policy is executed in receding horizon fashion (Mayne & Michalska, 1988) where only the first predicted control u, is exccuted and then a new sequence of controls re-predicted to incorporate feedback from the new image observations. For every training sample (8), denote
ft VicooenEn (T--deerrr(Kt)) 1,,1, Tt)
(9)
as the visual feature of the sequence of images.
3.1
CLASSIFICATION ACCORDING TO RELATIVE POSE ORTIAANTS (REPOS)
Goal of Planar Pushing
Relative Pose Orthants
Goal Block Stacking
"Move southeast, Rotate clockwise"
"Push southwest,
Rolate counter clockwise"
Figure 4: Control-oriented classification for continuous vision-based control tasks. Left: planar pushing; Right: block stacking.
The essence of neural collapse is law of clustering in the representation space. To study this, we need assign every feature vector ft class. In image classification and discrete vision-based control (e.g.. Lunar Lander), this class is explicitly given. However, in vision-basse continuous 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, also does not lead to observation of NC (see Appendix B).
We then conjecture that for control tasks such as planar pushing and block stacking. 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. 4 for Planar Pushing, the vision encoder may set the goal of control to "push the T block southwest and rotate it counter-clockwise" Similarly, in the right image in Fig. 4 for Block Stacking, the vision encoder may convey "move the block southeast and rotate clockwise". Building upon this intuition, we design two strategies to classify each training sample: that is goal-based and the other that is action-based. Since both strategies lead to similar observations of NC, we only present goal-based classification in the main text and refer the interested reader to Appendix for action-based classification.
Goal-based classification Given sample 8, as in (8), we look at image J,- We compute the relative pose of the target position with respect to the object. As depicted in Fig. 4, this relative pose is triplet (Ir y. 0) containing 2D translation and ID rotation for both planar pushing and block stacking. We divide the 3D space that the relative pose triplet (x, y. 0) lives into eight classes based the signs of x. y, 6. 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. middle. A nice property that the resulting classes semantically interpretable!
Remark 1 (Finegrained REPOS). What will happen the relative pose space is divided into a larger number of classes? In Appendix G.4, we divide the relative pose space into 64 and 216 classes, and demonstrate that such law of clustering still holds, albeit to a slightly weaker extent.

---

Published as a conference paper at ICLR 2025
3.7
2.33
0.13
222
0.10
0-4
1.2
0.02
0.3
0.3
0:0
0.08
52
200
150
200
250
322
52
150
155
200
250
320
232
300
73E 30012201001700
220
250
202 220
302
Epach
Epoch
tpoch
Eppch
(a) ResNett DM
(b) DINOV2 DM
(c) ReRNet + LSTM
(d) DINOv2 + LSTM
Figure 5: Test scores w.r.t. training epoches of four different instantiations of the image-based control pipeline for planar pushing. In (a) and (b) we show test scores of three random seeds. In (c) and (d) we show test scores of single seed because using LSTM as the action decoder leads to poor test-time performance, an observation that is consistent with Chi et al. (2023).
3.2 PREVALENT EMERGENCE OF NEURAL COLLAPSE
Using the -riente classification strategy described above, we are ready to study whether law of clustering similar to NC emerges the visual representation space for continuous control tasks. Given the set of visual features F { J we assign a class label E |C'] to each feature vector Se and denote fe. We used the same neural collapse metries introduced in $2.
3.2.1 PLANAR PUSHING
Simulation setup We collect N 500 expert demonstration on a push-T setup shown in Fig. 2. At each round, the object and target positions are randomly initialized and the same human expert controls the pusher push the object into alignment with the target position (through computer interface provided by pymunk (Blomqvist, 2024)). This provides M 55, 480 training samples. We train four different instantiations of the image-based control pipeline: using ResNet or DINOv2 as the vision encoder, and Diffusion Model (DM) or LSTM as the action decoder. The four trained models are evaluated on a test push-T dataset with 100 tasks. We define our evaluation metric as the ratio of the overlapping area between the object and the target to the total area of the target.
Results Fig. 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. (2023)) is slightly worse but quite close. When an LSTM is used the action decoder replace DM, the performance significantly drops, observation that is consistent with Chi et al. (2023) 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. plots the three NC evaluation metrics w.r.t. training epochs for two different models (both with DM) with goal-based classification strategy described in $3.1. We observe consistent decrease of three NC metrics as training progresses, suggesting the prevalent emergence of a law of cluster- ing that similar to NC. Results with action based classification are similar and provided in Ap- pendix G.2. Despite the poor test-time performance of trained models using LSTM, similar neural collapse is observed and shown in Appendix G.
From pretrained to finetuned ResNet features As observed in Chi et al. (2023); Kim et al. (2024); Team et al. (2024), pretrained ResNet features deliver poor performance when used for con- trol (such planar pushing) and end-to-end finetuning is necessary to adapt the vision features (see Appendix D for concrete evidence). This is intriguing: why are pretrained visual features insuffi- cient for control? We now have plausible answer. ResNet is pretrained for image classification, and according to NC, the pretrained ResNet features are clustered according to the class labels in image classification, such as dogs and cats. However, under our NC observation in Fig. 6, ResNet features for planar pushing are clustered according to "control-oriented" classes that are related to the rela- tive pose between the object and the target. Therefore, during finetuning, we conjecture the visual features have "re-clustered" according to the new task. We verify this conjecture in Appendix E.
3.2.2 BLOOC STACKING
Simulation setup Designed by Mandlekar et al. (2023), the block stacking is implemented as one of the manipulation tasks in MimicGien dataset using robosuite framework (Zhu et al., 2020)
7

---

Published as a conference paper at ICLR 2025
35
30
0.67
0.22
CDNV
25
Angle
0.33
20
0.39
0.09
0.50
0.07
0.54
150
250
150 Epoch
200
250
300
150
200
250
300
Epoch
Epoch
(a) ResNet DM, goal- baaldbased classification
0.30
13
0.62
12
0.61
80060
CDWV
10
0.59
0.00
0.50
0.557
0.07
0.55
0.06
0.54
200
Epoch
Epoch
(b) DINOv2 DM, goal-based classification
Figure 6: Prevalent emergence of neural collapse the visual representation space for planar push- ing. (a) Three NC metries w.r.t. training epochs using ResNet as the vision encoder and diffusion model (DM) as the action decoder, with goal-based classification label. (b) Three NC metrics w.r.t. training epochs using DINOv2 as the vision encoder and DM the action decoder. All the plots show the mean and standard deviation (in shaded band) over three random seeds. The test scores are shown Fig. 5(a)(b). Similar observations of NC hold when replacing DM with LSTM as the action decoder, shown in Appendix G.
backended MuJoCo (Todorov al., 2012). To train the behavior cloning pipeline, used the dataset "core stack_d0" provided by MimicGen which contains N = 1000 demos as our expert demonstrations. This provides M 107, 590 training samples. We used ResNet18 the vision encoder and diffusion model as the action decoder.
Results Fig. plots the three NC evaluation metrics w.r.t training epochs for Resnet+DM with the goal-based classification strategies described in $3.1. We observe consistent decrease of three NC metrics as training progresses, suggesting the prevalent emergence of law of clustering that is similar NC in the block stacking manipulation task.
110
0.62
0.61
90
8
0.60
0.10
CDNV
0.59
0.58
70
0..7
0.50
0.07
60
0.55
0.06
200
250
150 Epoch
200
250
200
250
Epoch
Epoch
Figure 7: Prevalent emergence of neural collapse in the visual representation space for Block Stack- ing. From left to right, it shows three NC metrics w.r.t. training epochs using ResNet the vision encoder and DM as the action decoder, with goal-based classification label.
VISUAL REPRESENTATION PRETRAINING WITH NEURAL COLLAPSE
Our experiments in $3.2 empirically demonstrated that image-based control policy trained on sufficient amount data exhibits an elegant control-oriented law clustering in its visual repre- sentation space. We study whether such clustering still holds under insufficient data.
8

---

Published as a conference paper at ICLR 2025
C
B
Pushing tasks
p
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
Pretraining 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".
Table I:
Weak clustering We use the same setup in $3.2.1, focus the (ResNet+DM) model, and train it using only 100 demonstrations. Fig. shows the test- test-time performance and three NC metrics according to the goal-based classification. Neural collapse appears to still hold, but much weaker extent, particularly when looking at the STDNorm metric.
0.20
0.32
0.22
3.50
3.300
COMM
2.44
3.50
E6.37
2.15
6.30
3.40
6..35
5.34
0.32
0.30
0.32
150
100
100
Figure 8: Weak neural collapse under 100 demonstrations of planar pushing. NC metrics are caleu- lated using goal-based classification. See Appendix G for results using action-bass classification.
Visual pretraining If knew a well-trained model exhibits strong clustering in its visual repre- sentation space, how about we encourage such phenomenon even without sufficient data? To test this idea, we pretrain the vision encoder by explicitly minimizing the three NC metries ((6) and (7)) prior to training the vision encoder jointly with the action decoder. Appendix F presents more de- tails. Table 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 substantial improvement by using NC pretraining: the minimum improvement is around 10% on the letter B, and the maximum improvement is around 35% the letter O (we intentionally carved out a piece to break symmetry). Note that all models are trained using only 100 expert demonstrations!
REAL-WORLD VALIDATION
The improvement shown in Table is astonishing and motivates us to further validate the effective- ness of NC pretraining in the real world.
Setup Fig. 9 depicts our real-world setup for vision- based planar pushing. The pusher is attached to the end- effector Franka Panda robotic arm to push T block into the target position a flat surface. An overhead Intel RealSense camera is used to provide image observations 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 rela- tive pose, we attached an AprilTag (Olson, 2011) to the bottom of the T-block, which enables real-time accurate pose estimation. To collect real-world expert demonstra- tions, we follow Chi et al. (2023) 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 (MesNet+DM)) model trained directly using 100
Franka Penda
Intel TeellSenee Cumeraa
Pusher
ApriTag
Object
Tanget postion
Figure 9: Real-world setup.

---

Published as a conference paper at ICLR 2025
demonstrations, and (b) an NC-peetrained model whose ResNet is first pretrained by minimizing NC metrics and then jointly trained with DM.
Visual complexities Real-world robotic manipulation needs to handle visually complex scenes, Therefore, we created another visually challenging setup where 11 distracting objects and 2 well- known paintings are placed the workspace (see Fig. 10(b)). For this challenging setup, we col- lected 200 expert demonstrations because 100 demonstrations were insufficient to train robust poli- cies. Using this dataset, we trained two policies with the same network setup as described before.
Results We tested both models on 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 block at exactly the same position for fair comparison of two models (the reader can check this in Fig. 10). For the clean background setup, the baseline model succeeded out of 10 times, and the NC- pretrained model succeeded out of 10. For the visually complex background setup, the baseline model succeeded out of 10, and the NC-pretrained model succeeded 9 out of 10. Fig. 10 shows two tests comparing the trajectories of the baseline and the NC-pretrained models under different background setups. The rest of the tests are shown in Appendix H and supplementary videos.
T
Test 8 on clean background
(b) Test on visually complex background Figure 10: Two real-world world tests. In every test, top shows trajectories of NC-pretrained model and bottom row shows trajectories 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 H and supplementary videos.
CONCLUSION
We presented an empirical investigation of the geometry of the visual representation space in an end-to-end image-based control policy for both discrete and continuous environments. We demon- strated that, similar to the phenomenon of neural collapse in image classification, 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.
Limitations and future work Our research opens lots of future research directions, such as the- oretical analysis of control-oriented clustering, extension of such investigation to policies trained from reinforcement learning, and connections neuroscience. the observation of NC and pretraining using NC require groundtruth relative poses during training, which is challenging to scale to complex tasks with multiple or deformable objects.
major limitation of our work is
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We would like to express our sincere appreciation to X.Y. Han for sharing valuable related works on Neural Collapse, and 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 Franka Panda robot arm setup. Additionally, Haocheng Yin acknowledges the Buropen 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.
We are additionally grateful to the anonymous reviewers for their constructive feedback and valuable suggestions that have helped improve the clarity and technical depth of this paper. Their detailed comments have significantly strengthened experimental analysis and theoretical discussions.
REFERENCES
Brenna D Argall, Sonia Chernova, Manuela Veloso, and Brett Browning. A survey of robot learning from demonstration. Robotics and Autonomous Systems, 57(5):469 483, 2009.
Antonio Bicchi and Vijay Kumar. Robotic grasping and contact: A review. In Proceedings 2000 ICRA. Millennium Conference. IEEE International Conference on Robotics and Automation. Sym- posia Proceedings (Cat. No. 00CH37065), volume 1, pp. 348-353. IEEE, 2000.
Victor Blomqvist. Pymunk: A 2d physics library for games and simulations. http://www. pymunk. org, 2024. Version 6.8.1.
Mariusz Bojarski. End end learning for self-driving cars. arXiv preprint 1604.006600003110 2016.
Guglielmo Bonifazi, lason Chalas, Gian Hess, and Jakub Lucki. Can we understand plasticity through neural collapse? arXiv preprint arXiv:2404.02719, 2024.
G Brockman. OpenAl Gym. arXiv preprint arXiv:1606.015400 2016.
Daniel Brown, Wonjoon Goo, Prabhat Nagarajan, and Scott Niekum. Extrapolating beyond sub- optimal demonstrations via inverse reinforcement learning from observations. In International Conference on Machine Learning, volume 97, pp. 783-792. PMLR, 2019.
Cheng Chi, Siyuan Feng. Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shu- ran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023.
Felipe Codevilla, Matthias Müller, Antonio López, Vladlen Koltun, and Alexey Dosovitskiy. End- to-end driving via conditional imitation learning. 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 4693-4700. IEEE, 2018.
Hien Dang, Tho Tran, Tan Nguyen, and Nhat Ho. Neural collapse for cross-entropy class- imbalanced learning with unconstrained ReLU feature model. arXiv preprint arXiv:2401.02058, 2024.
Shibhansh Dohare, Fernando Hernandez-Garcia, Qingfeng Lan, Parash Rahman, A Rupam Mah- mood, and Richard Sutton. Loss of plasticity in deep continual learning. Nature, 632(8026): 2024.
Weinan E and Stephan Wojtowytsch. On the emergence of simplex symmetry in the final and penultimate layers neural network classifiers. arXiv preprint arXiv:2012.05420, 2020.
Cong Fang, Hangfeng He. Qi Long, and Weijie J Su. Exploring deep neural networks via layer- peeled model: Minority collapse in imbalanced training. Proceedings of the National Academy of Sciences, 118(4))e2103091111 2021.
Tomer Galanti, András György, and Marcus Hutter. On the role neural collapse in transfer learn- ing. arXiv preprint 152/12 2021.
11

---

Published as a conference paper at ICLR 2025
Tomer Galanti, András György, and Marcus Hutter. Improved generalization bounds for transfer learning via neural collapse. In First Workshop on Pre-training: Perspectives, Pirfalls, and Paths Forward at ICML 2022, 2022.
Suresh Goyal, Andy Ruina, and Jim Papadopoulos. Planar sliding with dry friction part 1. limit surface and moment function. Wear, 143(2):307-330, 1991.
XY Han, Vardan Papyan, and David L Donoho. Neural collapse under mse loss: Proximity and dynamics on the central path. arXiv preprini urXiv:2106.0220 020000, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 770-778, 2016.
Stefan Hinterstoisser, Vincent Lepetit, Slobodan Ilie, Stefan Holzer, Gary Bradski, Kurt Konolige, and Nassir Navab. Model based training, detection and pose estimation of texture-less 3d ob- jects in heavily cluttered scenes. In Computer Vision-ACCV 2012: 11th Asian Conference on Computer Vision, Daejeon, Korea, November 5-9, 2012, Revised Selected Papers. Part I 11, pp- 548-562. Springer, 2013.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in Neural Information Processing Systems, 29, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in
Neural Information Processing Systems, 33:6840-6851, 2020.
François Robert Hogan and Alberto Rodriguez. Feedback control of the pusher-slider system: A story of hybrid and underactuated contact dynamics. In Algorithmic Foundations of Robotics XII: Proceedings of the Twelfth Workshop on the Algorithmic Foundations of Robotics, pp. 800-815. Springer, 2020.
Wanli Hong and Shuyang Ling. Neural collapse for unconstrained feature model under cross- entropy loss with imbalanced data. arXiv preprint arXiv:2309.09725, 2023.
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Hlyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR), 50(2):1-35, 2017.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 2022.
Jiachen Jiang, Jinxin Zhou, Peng Wang, Qing Qu, Dustin Mixon, Chong You, and Zhihui Zhu. Generalized neural collapse for large number of classes. arXiv preprint arXiv:2310.05351, 2023.
Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan Foster, Grace Lam, Pannag Sanketi, al. OpenVLA: An open-source iisuo--lagguage-action model. arXiv preprint 2rXiv:2406.092466 2024.
Michael Laskey, Jonathan Lee, Roy Fox, Anca Dragan, and Ken Goldberg. DART: Noise injection for robust imitation learning. In Proceedings of the 1st Annual Conference on Robot Learning, volume 78, pp. 143-156. PMLR, 2017.
Xuantong Liu, Jianfeng Zhang, Tianyang Hu, He Cao, Yuan Yao, and Lujia Pan. Inducing neural collapse in deep longetailed learning. In International Conference Artificial Intelligence and Statistics, pp. 1 11534-1154 PMLR, 2023.
Kevin M Lynch and Matthw T Mason. Stable pushing: Mechanics, controllability, and planning. The International Journal of Roboties Research. 15(6):533-556 1996.
Ajay Mandlekar, Yuke Zhu, Animesh Garg, Jonathan Booher, Max Spero, Albert Tung, Julian Gao, John Emmons, Anchit Gupta, Emre Orbay, et al. Roboturk: A crowdsoureing platform for robotic skill learning through imitation. Conference on Robot Learning, volume 87, pp. 879-893 PMLR, 2018.
12

---

Published as a conference paper at ICLR 2025
Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang, Rohun Kulkamn,, Li Fei- Fei, Silvio Savarese. Yuke Zhu, and Roberto Martín-Martín What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXtv:2108..3229: 2021.
Ajay Mandlekar, Soroush Nasiriany, Bowen Wen, Iretiayo Akinola, Yashraj Narang, Linxi Fan, Yuke Zhu, and Dieter Fox. MimicGen: A data generation system for scalable robot learning using human demonstrations, In 7th Annual Conference on Robot Learning. 2023.
David Q Mayne and Hannah Michalska. Receding horizon control of nonlinear systems. In Pro- ceedings of the 27th IEEE Conference on Decision and Control, pp. 464 465. IEEE, 1988.
Urs Muller, Jan Ben, Eric Cosatto, Beat Flepp, and Yann Cun. Off-road obstacle avoidance through end-to-end learning. Advances in Neural Information Processing Systems, 18, 2005.
Aneesh Muppidi, Zhiyu Zhang, and Heng Yang. Fast TRAC: A parameter-free optimizer for lifelong reinforcement learning. In Advances in Neural Information Processing Systems (NeurlPS), 2024.
Edwin Olson. AprilTag: A robust and flexible visual fiducial system. In 2011 IEEE International Conference on Robotics and Automation, pp. 3400-3407. IEEE, 2011.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. DINOv2: Learning robust visual features without supervision. arXiv preprini arXiv:2304.07193. 2023.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40): 24662,, 2020.
Liam Parker, Emre Onal, Anton Stengel, and Jake Intrater. Neural collapse in the intermediate hidden layers of classification neural networks. arXiv preprint urXiv:2308.02700 2023.
Xue Bin Peng, Marcin Andrychowicz, Wojeiech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 3803-3810. IEEE, 2018.
Antonin
Raffin.
RL baselines3
Z00.
https://githu/e wtt./CR-R///DDOO..
rl-baselines3- 200, 2020.
Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. In Proceeuings of Robotics: Science and Systems (RSS), 2018.
Akshay Rangamani and Andrzej Banburski-Fahey. Neural collapse in deep homogeneous classifiers and the role of weight decay. ICASSP 2022 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp.
IEEE, 2022.
Akshay Rangamani, Marius Lindegaard, Tomer Galanti, and Tomaso A Poggio. Feature learning in deep classifiers through intermediate neural collapse. In International Conference on Machine Learning. pp. 2274555. PMLR, 2023.
Venkatesh G Rao and Dennis S Bernstein. Naive control of the double integrator. IEEE Control Systems Magazine, 21(5):86-97 97, 2001.
Siddharth Reddy, Anca D Dragan, and Sergey Levine. SQIL: Imitation learning via reinforcement
learning with sparse rewards. arXiv preprint 19X.:11005.1108, 2019.
Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9, pp- 661-668. JMLR Workshop and Conference Proceedings, 2010.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and struc- tured prediction to no-regret online learning. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, volume 15, pp. 627-635. JMLR Workshop and Conference Proceedings, 2011.
13

---

Published as a conference paper at ICLR 2025
Peter Rousseeuw. Silhouettes: graphical aid to the interpretation and validation of cluster analy- sis. Journal of Computational and Applied Mathematics, 20:53-65, 65, 1987.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1rXiv:1707.06345 2017.
Pratyusha Sharma, Lekha Mohan, Lerrel Pinto, and Abhinav Gupta. Multiple interactions made easy (mime): Large scale demonstrations data for imitation. In Conference on Robot Learning, pp. 906-915. PMLR, 2018.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484
2016.
Wen Sun, Arun Venkatraman, Geoffrey J Gordon, Byron Boots, and J Andrew Bagnell. Deeply ag- grevated: Differentiable imitation learning for sequential prediction. In International Conference on Machine Learning, volume 70, pp. 310933003333 -318. PMLR, 2017.
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et al. Octo: An open-source generalist robot policy. arXiv preprint arXiv:2405. 12211, 2024.
Tom Tirer and Joan Bruna. Extended unconstrained features model for exploring deep neural col- lapse. International Conference on Machine Learning, pp. 21578-21505. PMLR, 2022.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012.
Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In Interna- tional Joint Conference on Artificial Intelligence, 2018.
Oriol Vinyals, Igor Babuschkin, Wojeiech M Czarnecki, Michaël Mathicu, Andrew Dudzik, Juny- oung Chung. David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 375(7782)3350-354 2019.
Zhendong Wang, Jonathan Hunt, and Mingyuan Zhou, Diffusion policies as an expressive policy class for offline reinforcement learning, arXiv preprint 2xXi::2208.061933 2022.
Patrick M Wensing, Michael Posa, Yue Hu, Adrien Escande, Nicolas Mansard, and Andrea Del Prete. Optimization-bann based control for dynamic legged robots. IEEE Transactions on Robotics, 2023.
Robert Wu and Vardan Papyan. Linguistic collapse: Neural collapse in (large) language models. arXiv preprint arXiv:2405. 17767, 2024.
Mengjia Xu, Akshay Rangamani, Qianli Liao, Tomer Galanti, and Tomaso Poggio. Dynamics in deep classifiers trained with the square loss: Normalization, low rank, neural collapse, and gener- alization bounds. Research, 6:0024, 2023.
Yibo Yang, Shixiang Chen, Xiangtai Li, Liang Xie, Zhouchen Lin, and Dacheng Tao. Inducing neural collapse in imbalanced learning: Do we really need a learnable classifier at the end of deep neural network? Advances in Neural Information Processing Systems, 35:37991-38002, 2022.
Kuan-Ting Yu, Maria Bauza, Nima Fazeli, and Alberto Rodriguez. More than million ways be pushed. a high-fidelity experimental dataset of planar pushing. In 2016 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 30-37. IEEE, 2016.
Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In International Conference Machine Learning, pp. 27179 27202. PMLR, 2022.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martín-Martin Abhishek Joshi, Soroush Nasiri- any, and Yifeng Zhu. robosuite: A modular simulation framework and benchmark for robot learning. In arXiv preprint arXiv:2009. /2293, 2020.
14

---

Published as a conference paper at ICLR 2025
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li. Chong You, Jeremias Sulam, and Qing Qu. A ge- ometric analysis of neural collapse with unconstrained features. Advances in Neural Information Processing Systems, 11-29820-299844 2021.
15

---

Published as a conference paper at ICLR 2025
A LUNAR LANDER EXPERIMENT DETAILS
Setup The Lunar Lander task is an environment OpenAl Gym (Brockman, 2016), whose control space discrete with 4 actions:
0: do nothing
1: fire left orientation engine
2: fire main engine
3: fire right orientation engine.
To collect an expert demonstration dataset on the lunar lander task, we trained an expert reinforce- ment learning algorithm Proximal Policy Optimization (PPO) (Schulman al., 2017)- on the LunarLander-v2 gym environment. We applied the default optimal network architecture and train- ing hyperparameters from the popular RL training library zoo3 (Raffin, 2020) to train the PPO. Then we used the trained PPO policy to collect 500 demos as our expert training dataset.
Training details After collecting expert demonstrations, we train an imitation learning model. We use ResNet18 model the vision encoder and use MLP model to map the 512 dimensional latent embedding into 64 dimensional. and predict the next action. The action decoder is structured another MLP network with 6 layers and maps the image latent embeddings to predicted action, then calculating the cross-entropy loss with the ground truth discrete actions. We train the model for 600 epochs.
use sequence
input images as observation input
B k-MEANS CLUSTERING OF EXPERT ACTIONS
Using the 500 expert demonstrations collected as described in $3.2, 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, 1987). Fig. 11 shows the average Silhouette score when varying the number of clusters, and Fig. 12 shows the Silhouette analysis for selected values of k. Based on Silhouette analysis, we choose the optimal
5.
Fig. 13 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 negative signal for neural collapse.
Silbauete Scare
0.40
0.38
0.36
0.32
0.33
10
16
Nurn_clusten clusters
Figure 11: Average Silhouette score of k-mmans clustering of the output action with respect to number of clusters (k value) from 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.
C ACTION-BASED CLASSIFICATION
In addition to the Goal-based classification mentioned in $3.1, we introduce another possible classi- fication method.
Action-based classification Given a sample s, as (8), we look at the images induced by the action sequence (ur,..., -1+ 1), i.e., (Je, Tet1,..., T++t). Importantly, for planar pushing,
16

---

Published as a conference paper at ICLR 2025
Chatte
defidar: edots
The valao
codTcart
Figure 12: Silhouette analysis for k-Means clustering sample data with number of clusters = 3, 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.
3
0.37
3.35
2200
3,74
STD
5.42
0.3
1.50
0.53
Epoch
Epoch
Epoch
Spoch
Figure 13: M-Means clustering of the output actions does lead to observation of neural collapse, even using the optimal value (k 5). Data obtained from training (ResNet+DM) model on 500 expert demonstrations described in $3.2.1.
possible that the expert not finish her action plan at time & + H, which can be detected the pusher is still in contact with the object. In those cases, enlarge H until the pusher breaks contact with the object, which forms a full episode of expert plan. Fig. 14 right illustrates such an episode, with the object transparency diminishing time progresses. In full episode, compute the relative pose of the object at time + with respect to the object at time t, which consists of a triplet (ar, y: 0) visualized in Fig. 14 right (Planar Pushing) and Fig. 15 right (Block Stacking).
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 shori-term goal (in an episode).
Goal (Inout)
Relative Pose Orthants
Action (Output)
"Push southwest, Rotate counter-dlockwlse
Figure 14: Control-oriented classification for Planar Pushing. Left: goal-based; Right: action-based.
D BENEFITS OF IMAGENET-PRETRANIE RESNET
On the planar pushing task, we investigate how much benefits does ImageNet pretraining bring to training a vision-based control pipeline.
Push-T with clean background task with clean background (shown in Fig. 2). Table second column shows the results.
We compare the performance of four pipelines on the push-T
In the table, "Frozen" indicates the ResNet weights are fixed and not finetuned in behavior cloning, while "Finetuned" indicates the weights are end-to-end finetuned from expert demonstrations. "Pre-
17

---

Published as a conference paper at ICLR 2025
Goal (inou)
Relstive Pose Orthants
Action (Duiput)
"Move southeast, Hotate clockwise"
Figure 15: Control-oriented classification for Block Stacking, Left: goal-based; Right: action- based.
trained indicates the ResNet weights are initialized from pretraining on ImageNet, and "Random" indicates random set of initial weights. all cases, the action decoder, i.., the diffusion model (DM) is not frozen and trained end-to-end.
Clearly, the Imagenet-pretrained ResNet does not show benefits either when frozen or finetuned, compared to a set of random weights.
Push-T with complex background We replace the clean background of push-T with messy table-top manipulation background adapted from the well-known LineMod dataset Hinterstoisser et al. (2013). An example image of the generated expert demonstrations can be found in Fig. 19. Table 2 third column shows the results. The pipelines are the same as before.
We make two observations:
When the ResNet is frozen, regardless of whether pretrained from ImageNet randomly initialized, the performance is very low. This shows that, if finetuning is not allowed, pretraining ImageNet does not help.
However, when finetuning is allowed, pretraining on ImageNet significantly increases the performance. This shows the rich priors leamned from ImageNet indeed helps transfer learn- ing for control tasks, when the scenes are visually complex.
Combining the results above, we draw the following conclusions.
In control tasks, if ResNet is not allowed to be finetuned, then neither pretraining Ima- geNet nor random initialization sufficient to learn a good policy for the control task.
When ResNet is allowed to be finetuned:
If the scene is visually simple (such as clean push- T), then pretraining on ImageNet does not show clear advantages over a random initialization.
If the scene is visually complex (with a complex background), then pretraining on ImageNet significantly benefits transfer learning.
Performance on clean push-T (%)
Performance on visually complex push-T (%)
Frozen Pretrained ResNet DM Frozen Random ResNet DM Finetuned Pretrained ResNet DM Finetuned Random ResNet + DM
30.17
32.90
29.95
32.22
62.83
56.10
63,66
46.67
Table 2: Investigation of the benefits of ImageNet Pretraining on planar pushing.
18

---

Published as a conference paper at ICLR 2025
E
NEURAL RE-COLLAPSE DURING FINETUNING
T
Figure 16: 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.
We collect 500 expert demonstrations in each of the five planar pushing domains shown in Fig. We then continuously train (ResNet+DM) model across five domains, each with 300 epochs. Fig. 17 shows the performance score and three NC metrics (according to the goal-based classifica- tion), 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., 2024; Muppidi al., 2024), there is clear evidence the (ResNet+DM) model loses its ability learn new tasks during continual finetuning (in every new domain, the model reaches around 70% performance). Results using the action-baseed classification are similar and presented in Fig. 18.
0.7-
5
5.000
2.53
5.000
CONN
5.070
3.50
5.000
5.055
50
1.0
150
200
250
200
100
150
200
250
300
100
155
250
250
300
100
150
200
255
300
Epoch
Epoch
Epoch
(a) Domain
3.000
0.05
0.00
3.350
0.555
CDNV
2.305
0.30
23.30
5.00
5.07
0.40
3.510
3.500
5.06
Epoch
Epoch
Epoch
Epoch
(h) Domain +
0.125
0.50
25.0
20.00
0.50
3..50
0.055
17.5
2.585
0.00
0.50
15.0
3.300
2.055
200
250
107
21t
250
201
102
157 Spoch
236
207
230
200
251
700
Epoch
Epoch
Epoch
(c) Domain
0.60
3.col
0.200
0.00
5.005
2.000
33.57
23.56
8.0e0
0.59
3.55
0.54
0.005
Epoch
Epoch
(d) Domain +
Figure 17: 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 repre- sent data in the source (old) domain. All results use the goal-based classification. Results using the action-baat based classification are similar and shown in Fig. 18.
19

---

Published as a conference paper at ICLR 2025
5.24
3.75
CONV
6.4
0.25
20
130
150
200
200
150
200
200
100
155
216
250
302
100
135
200
230
300
tpoch
Epach
Epoch
1poch
(a) Domain
0.85
11.
2,300
0.22
0.60
3.215
0.355
CDMV
13.25
0.32
0.50
0.49
20.175
0.40
6.30
0.30
156
700
150
2.325
150
230
205
(b) Domain 2 - 3
0.00
5.00
130
1.20
6.39
5.00
20.60
3.24
5-15
5.34
0.30
2.20
0.50
0.32
150
200
250
150
0.26
250
Epoch
Epach
Epoch
tpoch
(c) Domain
0.60
120
110
0.60
CONV
3,70
Nom
5.37
8
5.28
0.50
3.20
0.50
110
100
150
700
115
100
(d) Domain
Figure 18: Neural re-collapse in the visual representation space during continual retraining across different domains. All results use the action- based classification.
F DETAILS OF VISUAL REPRESENTATION PRETRAINING
Training details In Table 1, for both the baseline and our method, the network architecture is exactly the same We used ResNet the vision encoder and diffusion model (DM) as the action decoder. Baseline trains the entire pipeline end-to-end with expert demonstrations using action noise L.2 loss supervision for 100 epochs, which ensures the loss fully converged.
Our method is different from the baseline only in terms of training strategy. We firstly pretrain the same ResNet using NC regularization (i.e., minimize the NC metrics and encourage the visual features to cluster according to control-oriented labels) for 50 epochs, and then end-to-end train pretrained DM pipeline from expert demonstrations ularized pretraining, because all the data samples cannot be fit into batch due to GPU limitation, we use the largest possible batch size and divide the randomly shuffled dataset into 4 batches in every pretraining epoch. We caleulate the NC regularization for each batch following this NC loss: 0.1 x CDNV + 10 x (STDNorm STDAngle).
another 50 epochs. During NC reg-
We want to emphasize that pipeline with control-oriented pretraining does NOT any extra information compared to the baseline during TEST time. At test time, our pipeline is still a pure vision-based control pipeline.
20

---

Published as a conference paper at ICLR 2025
0.68
0.16
0.67
0.68
0.14
30
65.65
Norm
CDNV
25
0.64
0.12
20
0.63
0.62
15
10
250
150
200
250
300
150
200
250
300
Epoch
Epoch
Epoch
(a-1) Res.Net DM, goal-based classification
260
0.40
Angle
North
0.33
200
0.30
180
0.30
0.25
0.20
0.20
120
200
0.15
Epoch
Epoch
(a-2) ResNet DM, retion-based classification Figure 20: Prevalent emergence of neural collapse in the visual representation space for complex background push-T. (a) Three NC metries w.r.t. training epochs using ResNet as the vision encoder and diffusion model (DM) as the action decoder.
G SUPPLEMENTARY EXPERIMENTAL RESULTS
G.1 SIMULATED PLANAR PUSHING WITH VISUALLY COMPLEX BACKGROUND
Setup Similar to the planar pushing setup in main text, collect N 500 expert demonstration on push-T setup with messy table-top manipulation background adapted from the well-known LineMod dataset (Hinterstoisser et al., 2013). An example of this dataset can be seen in Fig 19. 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. We define primary evaluation metric as the ratio of the overlapping area between the object and the target position to the total of the target position.
Results We used 500 demonstrations to train an end-to-end pipeline us- ing ResNet the vision encoder (pretrained from ImageNet) and Diffusion Model as the action decoder. The test-time performance for a well-trained model is 56.1% Fig. 20 plots the three NC evaluation metrics w.r.t. train- ing epochs with both goal-based and action-based classification. 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, even though task with many visual distractions.
Figure 19: Push-T in LineMod back- ground.
G.2 NEURAL COLLAPSE IN PLANAR PUSHING WITH ACTION-BASED CLASSIFICATION
Fig. 21 plots the three NC evaluation metrics w.r.t training epochs for two different models (both with DM) with action-based classification label described in Appendix C. We also observe con- sistent decrease of three NC metrics as training progresses, suggesting the prevalent emergence of law of clustering that is similar to NC. Notably, the CDNV metric in the based classification as in Fig. 6 is significantly smaller than that of an action- ratioo--bsedd classification, suggesting that the goal-based classification may correspond stronger level of clustering.
of using a goal-
21

---

Published as a conference paper at ICLR 2025
0.45
200
0.40
0.350
Norm
CDNV
160
3.3
0.30
140
0.300
0.275
0.29
0.20
100
80
0.225
200
250
250
300
200
250
Epoch
Epoch
Epoch
(a) ResNet
DM, action-based classification
0.40
180
0.40
160
0.35
0.35
Angle
140
0.30
120
0.25
0.29
100
0.20
0.20
80
0.15
0.15
150
200
200
Epoch
Epoch
(b) DINOV2 DM, action-based classification
Figure 21: Prevalent emergence of neural collapse in the visual representation space for planar push- ing. (a) Three NC metries w.r.t. training epochs using ResNet as the vision encoder and diffusion model (DM) as the action decoder, with oectoon-bssee classification label. (b) Three NC metrics w.r.t. training epochs using DINOv2 as the vision encoder and DM 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. 5(a)(b).
G.3 OBSERVATION OF NEURAL COLLAPSE WITH LSTM
In the main text, we demonstrated emergence 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. 22.
G.4 FINEGRAINED REPOS
In Remark 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 triplet (ar, y, 0) containing a 2D translation and 1D rotation. In the 8-class division, we simply divide each dimension (æ, y, 0) into two bins (-oo, 0] and [0, +oo), leading to 2 8 classes. To divide the space into more classes, into more bins. particular, we try two extra options. First, we divide each dimension into bins, leading to 64 classes. Second, divide each dimension into 6 bins, leading to 216 classes. Fig. 23 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 increased, neural collapse still emerges. However, appears that the CDNV metric can "bounce back" when the number of classes increased. Similar observations hold for the (DINOv2+DM) model and shown in Fig. 24.
divide each dimension
G.5 WEAK CLUSTERING
In Fig. we showed weak clustering when training using 100 demonstrations according to the goal- based classification. Here Fig. 25 shows weak clustering according to the action-based classification.
G.6 VISUAL REPRESENTATION PRETRAINING WITH NC FOR DINOV2
In addition to pretrain a ResNet vision encoder with NC regularization as in Table 1, we also present results of using DINOv2 vision encoder and pretrain it by minimizing the NC metries in Table 3.
22

---

Published as a conference paper at ICLR 2025
19
0.635
0.110
0.924
88
0.633
Norm
CDNV
17
0.632
0.3
0.631
0.095
26
0.630
0.090
0.629
0.085
1509
1500
1750
750 1000 1250
1500
1000
1250 1500 1750
Epoch
Epoch
Epoch
(a-1) ResNet + LSTM, goal-based classification
0.235
0.230
CDNV
230
Ang
9
0.22
225
0.35
0.200
0.34
0.215
0.33
0210
0.000
Epoch
Epoch
(a-2) ResNet + LSTM, action-based classification
140
0.60
0.5
0.666
CDNV
0.4
0.66
0.3
0.2
0.62
0.3
50
100
200
250
30
100
150
200
250
300
100
150 Epoch
250
300
Epoch
Epoch
(b-1) DINOv2 + LSTM, goal- goal-based classification
260
0.5
240
0.4
STD Angle
CDNV
0.3
160
0.2
148
0.1
50
100
150
205
250
330
100
110
200
255
3017
50
100
110
200
200
300
Epoch
Epoch
Epoch
(b-2) DINOv2 + LSTM, action based classification
Figure 22: 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 L.STM 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. 5 (c). Only one seed 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. 6 and Fig. 21.
Clearly, we observe that control-oriented pretraining improves test-time performance, regardless of whether ResNet or DINOv2 are used.
23

---

Published as a conference paper at ICLR 2025
0.23
0.32
Angle
CDNV
20
0.30
15
10
0.08
50
150
250
150 Epoch
200
230
300
150
200
250
Epoch
Epoch
(a) Classes REPOS, ResNet + DM
5.0
0.58
0.56
0.54
9.400
CDNV
0.350
0.50
0.325
0.48
0.300
2.5
0.275
50
100
150 Epoch
700
250
300
50
100
150
200
250
300
50
100
150
700
250
300
Epoch
Epoch
(b) 64
Classes REPOS, ResNet DM
0.550
0.475
2.4
0.525
0.500
0.425
2.0
50.475
0.400
018
05450
1.6
0.425
0.400
0.325
12
50
100
150 Epach
200
250
300
0.375
30
100
150 Epoch
200
250
300
0.300
30
100
150 Epoch
200
250
300
(c) 216 Classes REPOS, ResNet DM
Figure 23: 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 action decoder, respectively for 64, 216 classes of REPOS.
R
C
B
Pushing tasks
R
3
Baseline(w/ DINOv2)
0. 1898
0.3747
0.2463
0.19
Control-oriented pretraining (w/ DINOv2)
0.2110
0.4602
0.2466
0.2468
Table 3: Pretraining DINOY2 vision encoder by minimizing the NC metrics also improves test- time model performance four domains corresponding to four letters in the word "ROBOT".
24

---

Published as a conference paper at ICLR 2025
0.30
15
0.62
12
0.61
11
80060
CDNV
10
0.00
0.07
0.55
0.00
0.54
Epoch
Epoch
Epoch
Classes REPOS, DINOv2 DM
4.25
0.60
0.40
4.00
0.58
9 0.50
3.50
0.50
3.00
2.75
Epoch
Epoch
Epoch
(b) 64 Classes REPOS, DINOv2 DM
0.44
2.4
0.525
0.500
0.475
0.36
2.0
0.450
0.34
0.32
0.30
0.375
0.28
50
100
150 Epoch
200
250
300
50
100
200 250 Epoch
150
300
30
100
150 Epoch
200
250
300
(a) 216 Classes REPOS, DINOv2 + DM
Figure 24: 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.
3.50
0.20
CONY
49
STD
0.30
0.10
200 270
105
100
372
100
275
250
205
100
220
270
700
Epoch
Epoch
Epoch
Spoch
Figure 25: Weak neural collapse under 100 demonstrations, NC metrics calculated using action- based classification.
25

---

Published as a conference paper at ICLR 2025
H REAL-WORLDS EXPERIMENTAL RESULTS
Here we provide full results for real-world experiments.
100 demonstrations on clean background In the main text we showed results for test. Fig. 28 and 29 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 on clean background In addition, we also trained two policies using only 50 demonstrations. Fig. 26 and 27 show the results for all 10 tests. The NC-pretrained policy succeeded out of 10, but the baseline policy succeeded 2 out of 10.
200 demonstrations on visually complex background In the main text we showed results for test. Fig. 30 and show the results for all 10 tests. As seen, the NC-pretrained policy succeeded 9 out of 10, but the baseline policy succeeded 6 out of 10.
For video recordings of the results, please consult the supplementary material.
26

---

Published as a conference paper at ICLR 2025
(a-1) Test 0
(b-1) Test
(c-1)
T
T
(d-1) Test
T
T
(e-1) Test 4
Figure 26: Real-world evaluation on clean background for test to test 4. In every test, top row shows trajectories of NC-pretrained model and bottom 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.
27

---

Published as a conference paper at ICLR 2025
(f-1) Test 5
(g-1) Test
(h-1) Test 7
T
(i-1) Test 8
G-1) Test 9
Figure 27: Real-world evaluation on clean background 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 the object, respectively. The middle three columns show overlaid trajectories generated by the policies.
28

---

Published as a conference paper at ICLR 2025
(a-2) Test
(b-2) Test
(c-2) Test 2
T
T
(d-2) Test 3
T
T
(e-2) Test
Figure 28: Real-world evaluation on clean background for test 0 to test 4. In every test, top row shows trajectories 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 the object, respectively. The middle three columns show overlaid trajectories generated by the policies.
29

---

Published as a conference paper at ICLR 2025
(f-2) Test 5
(g-2) Test
(h-2) Test 7
T
T
(i-2) Test 8
(i-) Test
Figure 29: Real-world evaluation on clean background for test 5 to test 9. In every test, top row shows trajectories 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 the object, respectively. The middle three columns show overlaid trajectories generated by the policies.
30

---

Published as a conference paper at ICLR 2025
(a-3) Test
(b-3) Test
(c-3) Test 2
(d-3) Test
(e-3) Test
Figure 30: Real-world evaluation on visually complex background for test 0 to test 4. In every test, top row shows trajectories of NC-pretrained model and bottom row shows trajectories of base- line. Both models are trained under 200 training demonstrations. The first and last column show the initial and final position the object, respectively. The middle three columns show overlaid trajectories generated by the policies.
31

---

Published as a conference paper at ICLR 2025
(I-3) Test 5
(g-3) Test
(h-3) Test 7
(i-3) Test 8
(i-3 Test 9
Figure 31: Real-world evaluation on visually complex background for test 5 to test 9. In every test, top row shows trajectories of NC-pretrained model and bottom row shows trajectories of base- line. Both models are trained under 200 training demonstrations. The first and last column show the initial and final position the object, respectively. The middle three columns show overlaid trajectories generated by the policies.
32

---

Published as a conference paper at ICLR 2025
1 RELATED WORK
Neural Collapse The phenomenon of neural collapse (NC) is firstly introduced in Papyan et al. (2020), about an elegant geometrie structure of the last-layer feature and classifier for well-trained model in classification tasks. While the initial NC is observed for model trained with cross entrophy (CE) loss, Han et al. (2021) 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 be the phenomenon occurring when training achieves global optimality, trained with CE loss function (E & Wojtowytsch, 2020; Zhu et al., 2021), and MSE loss function (Zhou et al., 2022; Tirer& & Bruna, 2022). Later, NC is studied in more diverse situations, e.g. NC in imbalanced training (Fang al., 2021; Hong & Ling, 2023; Dang et al., 2024). NC through inter- mediate layers (Rangamani et al., 2023; Parker et al., 2023), NC in transfer learning (Galanti et al., 2021; 2022), and relationship between NC and optimization methods (Xu al., 2023; Rangamani & Banburski-Fahey 2022). Since NC is a critical indicator of model robustness, it is also studied tool improve the model (Bonifazi et al., 2024; Liu et al., 2023; Yang et al., 2022). Recently, neural collapse is studied in more complex applications, like large language model training (Wu & Papyan, 2024; Jiang al., 2023). 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 key role in identifying NC.
Behavior Cloning Behavior Cloning (BC) is a fundamental approach in Imitation Learning, where an agent learns mimic expert demonstrations (Argall et al., 2009; Hussein et al., 2017). 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., 2005), particularly when combined with data augmentation tech- niques (Laskey et al., 2017; Peng et al., 2018) or iterative refinement (Ross & Bagnell, 2010; Ho & Ermon, 2016). This method has been widely applied across various robotic tasks including arm manipulation (Sharma al., 2018; Rajeswaran et al., 2018), autonomous driving (Bojarski, 2016; Codevilla et al., 2018), and game playing (Silver et al., 2016; Vinyals et al., 2019). However, BC faces challenges such as compounding errors (Ross al., 2011) and struggles with long-horizor tasks (Sun et al., 2017). Recent advancements have focused on addressing these limitations through improved data collection strategies (Mandlekar al., 2018), robust loss functions (Reddy et al., 2019; Brown et al., 2019), and integration with other learning paradigms (Rajeswaran et al., 2018; Zhu et al., 2020). 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., 2023), has introduced a groundbreaking approach to behavior cloning for robot control. This method leverages the principles of denoising diffusion probabilistic models (Ho al., 2020) to learn 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., 2022). Diffusion policy has demonstrated impressive results in challenging robotics tasks, effec- tively addressing longstanding issues in traditional behavior cloning approaches such distribution shift and long-horizon planning (Wang et al., 2022). As the field progresses, diffusion policy stands as promising direction for enhancing the capabilities and robustness of imitation learning in com- plex robotic systems. In this work, we show that probing the geometry of the visual representation space helps understand and improve diffusion policy.
33


Reviews:
[{"review_id": "s49TP0Lh3o", "reviewer": "Reviewer_93jv", "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.", "questions": "Can you validate these findings on more visually complex domains?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.", "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. \n\nI took issue with the discussion of resnet features throughout the paper. \n- 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.\n- 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). \n\nOverall, as a study of visual representations the findings need to be verified on more visually complex domains.\n\n\nTypo: \"Expect\" instead of \"expert\" on line 118.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wCrOG7JKQ8", "reviewer": "Reviewer_Nat5", "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.", "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.\n- 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.\n- 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.\n- 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The study of neural collapse in the visuomotor control setting is novel and well-motivated.\n- The minimum-time double integrator example is very clear.\n- Sim and real experiments show that neural collapse regularization objectives improve downstream task performance.\n- Fine-tuning experiments show that neural collapse happens when transferring across domains.", "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).\n- The criteria for determining whether neural collapse has occurred seems unclear in this setting.\n- Experiments are mostly on the planar pushing task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qAoZk1NVTS", "reviewer": "Reviewer_FVeR", "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.", "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?\n\n8. 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? \n\n9. Why aren't DinoV2 features also used for Table 1 (i.e pretraining DinoV2 to minimize NC metrics) and using that for control policies?\n\n10. What is the Baseline model architecture that is used in Table 1?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 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.\n\n2. 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.", "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:\n\t- Explicitly state potential benefits of finding NC in control policies (e.g. improved generalization, more interpretable features)\n\t- Expand on why clustering in the visual representation space might be important for control policies, even if they are already working well\n\n4. 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.\n\n5. 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? \n\n6. 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)?\n\n----\n**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.\n---\n**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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "04eBBnfBk8", "reviewer": "Reviewer_NdZb", "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:\n1. Control-oriented clustering emerges in the visual representation space, with latent features forming clusters according to control-relevant classes.\n2. Using NC as a pretraining strategy for the vision encoder improves test-time performance, particularly when training data is limited.\n3. Real-world experiments validate the effectiveness of NC pretraining in enhancing control performance.", "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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 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.\n2. 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.\n3. Real-World Validation: The paper includes real-world experiments on vision-based planar pushing to show the effectiveness of NC pretraining\n\nOverall, 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.", "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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)