Stem-OB: Generalizable Visual Imitation Learning with Stem-Like Convergent Observation through Diffusion Inversion
Authors: Kaizhe Hu
Zihang Rui
Yao He
Yuyao Liu
Pu Hua
Huazhe Xu
Venue: ICLR 2025 Spotlight (2025)
Keywords: Robotics
Imitation Learning
Visual Imitation Learning
Robustness
Diffusion Model
Diffusion Inversion
Primary area: applications to robotics, autonomy, planning

Abstract: Visual imitation learning methods demonstrate strong performance, yet they lack generalization when faced with visual input perturbations like variations in lighting and textures. This limitation hampers their practical application in real-world settings. To address this, we propose ***Stem-OB*** that leverages the inversion process of pretrained image diffusion models to suppress low-level visual differences while maintaining high-level scene structures. This image inversion process is akin to transforming the observation into a shared representation, from which other observations also stem. *Stem-OB* offers a simple yet effective plug-and-play solution that stands in contrast to data augmentation approaches. It demonstrates robustness to various unspecified appearance changes without the need for additional training. We provide theoretical insights and empirical results that validate the efficacy of our approach in simulated and real settings. *Stem-OB* shows an exceptionally significant improvement in real-world robotic tasks, where challenging light and appearance changes are present, with an average increase of **22.2%** in success rates compared to the best baseline. Please refer to [this link](https://stem-ob.github.io/) for more videos and details.

Published as a conference paper at ICLR 2025
Stem-OB: GENERALIZABLE VISUAL IMITATION LEARNING WITH STEM-LIKE CONVERGENT OBSER- VATION THROUGH DIFFUSION INVERSION
Kaizhe Hu Zihang Rui
Yao He
Yuyao Liu-
Pu Hua 123
Huazhe Xu
Tsinghua University Shanghai Qi Zhi Institute Shanghai Al Lab Stanforr University
hukaizhe22@maila tsinghua. edw. en, huazhe sanzhhexx@mail. tsinghua. edu en
$1.33
41.18
Suecco
38.42
9tm-00
SE34
Mis
RD
Avunge Suceess Rate of4 Real World Tisks
92.8
Comparison Stem-CIH with the Roseline
Figure 1: Left: The tree of Stem-OB inversion is composed of different objects progressively in- verted through a diffusion inversion process. Moving downward alone the tree's branches, objects of different textures, appearances, and categories gradually get closer, eventually converging into the same root Gaussian noise, where they are completely indistinguishable Right: Real-world tasks success rate, where Stem-OB showcases significant improvement.
ABSTRACT
Visual imitation learning methods demonstrate strong performance, yet they lack generalization when faced with visual input perturbations like variations light- ing and textures. This limitation hampers their practical application in real-world settings. To address this, we propose Stem-OB that leverages the inversion pro- of pretrained image diffusion models suppress low-level visual differences while maintaining high-level scene structures. This image inversion process is akin to transforming the observation into shared representation, from which other observations also stem. Stem- OB offers simple yet effective plug-and-pl and-play solution that stands in contrast data augmentation approaches. It demonstrates robustness to various unspecified appearance changes without the need for addi- tional training. We provide theoretical insights and empirical results that validate the efficacy of our approach in simulated and real settings. Stem-OB shows ceptionally significant improvement in real-world robotic tasks, where challeng- ing light and appearance changes are present, with average increase of 22.2 % in success rates compared to the best baseline. See our website for more videos.
INTRODUCTION
Visual Imitation Learning (IL), where agent learns to mimic the behavior of the demonstrator by learning a direct mapping from visual observations to low-level actions, has gained popularity
in recent real-world robot tasks (Chi et al. 2023 Zhao et al. 2023 Wang et al., 2023a Chi et al.
Indicates equal contribution.

---

Published as a conference paper at ICLR 2025
2024 Ze et al. 2024). Despite the versatility demonstrated by visual IL. learned policies are often brittle and fail to generalize to unseen environments, even minor perturbations such as altering the lighting conditions or the texture of the object may lead to failure of the learned policy (Xie et al. 2023 Yuan et al., 2024b). The underlying reason is that the high-dimensional visual observation space redundant with virtually infinite variations in appearance that irrelevant to the task and hard to generalize.
As human beings, we can easily manipulate objects that have different appearances. For example, we can pick up a coffee cup regardless of its color, texture, or the lighting condition of the room. This is partially because our visual system is capable of abstracting the high-level semantics of the scene, such as the silhouette of the object, the structure and arrangement of different objects, etc in a hierarchical manner (Hochstein Ahissar 2002), effectively merging scenes with perceptual differences to similar "meta" observations.
Augmentation techniques such as Spectrum Random Masking (SRM) (Huang et al., 2022) and Mixup (Zhang et al., (2188) remove details from observations to encourage the model to focus on structural features; however, they lack the ability to distinguish between low-level and high-level features. It is preferable if we can sweep the photometrical differences while maintaining the high- level structure for the scene. Achieving this requires a semantic understanding of the observations, and naively perturbing the data with Gaussian noise can lead to irreversible information loss.
Pretrained large image diffusion models, such as Stable Diffusion (Rombach et al., 2022; Esser
et al., 2024), embed essential world knowledge for visual understanding. Apart from synthesizing new images from random noise, these models are capable to perform a reverse procedure called inversion (Song et al. (2022), which converts an image back to the space of random noises. A recent study (Yue et al. 2024) indicates that this inversion process selectively eliminates information from the image. Rather than uniformly removing information from different semantic hierarchies, it will push those images with similar structures closer in the early stages of the inversion process. Inver- sion like the reprogramming of a differentiated cell back stem cell, which bears the totipotency to differentiate into any cell type. This characteristic aligns perfectly with our will of enhancing the robustness and generalizability of visual II. algorithms to visual variations. To distill such property into a visual IL policy, we propose an imitation learning pipeline which applies diffusion inversion to the visual observations. We name our method Stem-OB to highlight the similarity between the inversed observation and the stem cell in biology, as illustrated in Figure
To be specific, our method is as simple as inverting the image for reasonable steps before sending them to the downstream visual IL algorithms. The number of steps is chosen empirically to balance removing irrelevant details without erasing essential high-level information. From this perspective, our approach fundamentally distinguishes from generative augmentation methods, which aim to
enrich the training dataset with more unseen objects and appearances (Yu et al., 2023) Mandlekar al. 2023). Moreover, Stem-OB is indifferent many unspecified appearance changes, in contrast to augmentation-based methods that must concentrate on a few selected types of generalization, thereby introducing inevitable inductive biases.
We provide theoretical analysis and a user study to support our claim that Stem-OB can effectively merge scenes with perceptual differences to similar "stem observations". Empirical study demon- strates the effectiveness of our approach in variety of simulated and real-world tasks and range of different perturbations. Stem-OB proves to be particularly effective in real-world tasks where appearance and lighting changes hamper the other baselines, establishing an overall improvement in the success rate of 22.2% %. What's better, no inference time inversion is required for Stem-OB to take effect, making the deployment of our method virtually free of computational cost.
2 RELATED WORKS
2.1 VISUAL IMITATION LEARNING AND GENERALIZATION
Visual Imitation Learning (VIL) is a branch of Imitation learning (IL) that focuses on learning action
mappings from visual observations. It typically follows two approaches: directly imitating expert policies, as in behavior cloning and DAgger (Ross et al. 201T), or inferring a reward function that
aligns the agent's behavior with expert demonstrations, like inverse reinforcement learning (Ng &
2

---

Published as a conference paper at ICLR 2025
Real-World
Simulation
Deployment Unseen Disturbance
II(O)
Training Pollcy
16
Diffusion Inversion
Gt/T
6-7
Figure 2: Overview of Stem-OB: (a). Stem-OB has been evaluated in both real-world and simulated environments. (b). The trained visual IL policies are directly applied to the original observation space o, demonstrating robustness
environmental disturbances. (c). We train the visual IL policy π on the inffusion-inversed latent space 6f/7, where denotes specific inversion step out of total of T', Each composite rectangle in the diffusion inversion process, made up three smaller sections, represents the latent vector of an image, with the smaller section depict the finer attributes (gray). During the inversion process, finer attributes converge earlier than coarser ones.
Russell 2000) and GAIL (Ho & Ermon 2016). The former approach, favored in recent works due to its scalability and practicality in complex, real- real-world tasks, has led to several advancements. Notable methods include Diffusion Policy (Chi et al. 2023), which leverages a diffusion model (Ho
2020 to maximize the likelihood of expert actions, and Action Chunk Transformer (Zhao et al.
ct al.
2023 that uses Transformer [Vaswani] 2017).
To enhance the generalization and robustness of visual imitation learning algorithms, various ap- proaches have been explored (Cetin & Celiktutan| (02T). For instance, [i et al. (2023) leverage trajectory and step level similarity through an estimated reward function, while Wan al. (2023) improve visual robustness by separating task-relevant and irrelevant dynamics models before ap- plying the GAIL framework (Ho & Ermon 2016). Zhang et al. (2023a) mutual information constraints to create compact representations that generalize to unseen environments. However, these approaches are not directly applicable to methods like diffusion policy, which focus imita- tion without reward functions or dynamics models. ZZhng et al. (2023) propose to filter extraneous action subsequence, yet their focus not visual perturbations. Most relevant
setting, eral works in robust visual reinforcement learning have explored adding noise in image or frequency domain improve generalizability (Huang et al. 2022 Lee & Hwang] 2024 Yuan al. 2024a), however, they lack the semantic understanding of the augmentation process.
2.2 INVERSION OF DIFFUSION MODELS AND ITS APPLICATION
Diffusion model inversion aims to recover an initial noise distribution from given image, enabling image reconstruction from the same noise via backward denoising. A common approach is DDIM inversion (Song et al. 2022). which estimates the previous noise using predictions from the current diffusion step, though the approximation will introduce cumulative errors. To address this issue, several methods employ learnable text embeddings (Mokady et al. 2023 Miyake et al. (023). fixed- point iterations (Pan et al.) 2023 Garibi al.) (024), or gradient based techniques (Samuel al. 2024) to refine the result. Another approach, based on the stochastic DDPM scheduler (Ho et al.
2020), reconstructs noisy images at each diffusion step (Huberman- (Hfuherman-SSieeelggae
al. 2024] Brack
al. 2024). In contrast to the DDIM inversion methods, the noise of each step statistically independent, making them ideal for our application since we can obtain the noisy image a certain step without the need to recover other steps, greatly reducing the calculation cost.
Diffusion Inversion is the crucial part of
image editing methods (Meng et al. 2021 Kawar et al. 2023), which typically involves first inverting the diffusion process to recover the noise latent, then denoise the latent with desired editing conditions. Recent works also explore to apply attention control over the denoising process improve the fidelity of the edited image Hertz et al.
3

---

Published as a conference paper at ICLR 2025
2022 T'umanyan al. (202)) and have shown promising application in robot learning tasks Gao et al. )024): Ju et al. p ; Zhu et al. (2024). Beyond that, inversion is also used in tasks like concept extraction Huang al. (023) and personalization Gal et al. (2022). Most recently, [Wang & Chen (2024) proposed the use of diffusion inversion to interpolate between image categories to
improve classification performance. Wang al. (2023h) uses diffusion inversion erase
sub-optimal trajectories from the dataset.
PROBLEM DEFINITION
Given a dataset of observation O and action A pairs, the goal of VII. to learn a policy (T(AOO that maps observations to actions. The policy is typically parameterized by neural network with parameters and is trained to minimize the negative log-likelihood of the actions.
To achieve the goal of generalizing to unseen environments, we seek method to preprocess or transform the observations such that task-irrelevant details are suppressed while preserving the high-level semantic structure that is critical for the task. The problem can be transformed into learning transformation T as the imput to the policy (A)T((O)) where T(O) is the transformed observation emphasizing high-level semanties while removing irrelevant details.
4 PRELIMINARY
We begin by outlining the fundamentals of Diffussio Inversion. A diffusion model operates with two passes: a backward denoising pass, which generates an image from noise, and forward pass, where noise is incrementally added to image until it becomes pure Gaussian noise. This forward process is a Markov chain that starts with øg. and gradually adds noise obtain latent variables 1 $1,22. Dr. Each step in this process is Gaussian transition following the common form
Ft Va + or6, (a of
(1)
where a E 1) represents the scheduler parameter at each step t, while o characterizes the vari- ance of the Gaussian noise e; introduced at each step. In Denoising Diffusion Probabilistic Models (DDPM) (Ho et al. 2020), Eq. by applying the cumulative product ā, = II
vI - Ct. Consequently, equation Eq. (1) can be reformulated as
Kt = V Va + VI- ā,e ~ ~N(a (1.
(2)
Diffusson inversion is similar to the forward process in that they both maps an image to a noise, however, inversion tries to preserve the image's s information and obtain the specific noise that can reconstruct the image during a backward denoising process.
DDPM inversion. We follow the DDPM inversion proposed in Huberman-Spicgelglas al. 2024).
BC Vā + VI ŌgÊE
(3)
The DDPM inversion we consider here differs slightly from Eq. 2 as é, ~ NO, I) are mutu- ally independent distributions, in contrast to the highly correlated E; in Eq. P As mentioned by Huberman-S Spiegellaas et al. (2024), the independence of é; results in a sequence of latent vectors where the structures mg are more strongly imprinted into the noise maps. An error reduction step is conducted in reverse order after the diffusion forward process to improve image reconstruction accuracy during the denoising process:
2, = Et-1 - /i(x))/or 0:-1 = fi(x:)) + ot
DDIM inversion. We follow the DDIM inversion proposed in Song et al. (202)), where at each forward diffusion step
1 āi Lt-1
+
(v=
-
1 Ā
1 co(
(5)
4

---

Published as a conference paper at ICLR 2025
Note that the noise c0( t, C) is now generated by network trained predict the noise based on the previous vector ø;-1 and the text embedding C which, in our case, is Ø.
5 METHOD
In this section, introduce the intuition and implementation of our framework. We first propose the intuition of applying inversion on observations through theoretical analysis based on attribute loss, diffusion-based measurement of image semantic similarity. Then, we conduct an illustra- tive experiment and a user study to validate our intuition. Finally, we explain how to practically implement Stem-OB and incorporate diffusion inversion into a visual imitation learning framework.
5.1 INTUITION DEVIATION BY ATTRIBUTE LOSS
Intuitively, as the diffusion inversion process moves forward, a source image and another variation of it become increasingly indistinguishable. The variation here could be low-level changes like lightning conditions, but also includes semantic changes such as replacing an object. If there are two different variations, alterations will become indistinguishable sooner than the pair with larger and structural changes. We borrow the definition of attribute loss from Yue et al. (2024) quantify the semantic overlapping of the two images at time step / during a inversion process:
want to show that as the inversion step increases, the pair with minor
loss(xo, yo: t) (z)) ((9e88)))
where To and yo are the latent variables of the two images, and OVL is the overlapping coef- ficient quantifying the overlapping area of two probability density functions. For an inversion process where each step follows Gaussian transition, it takes the form = Vā,zo + ae ~ N(Voran, II). The OVL can be further calculated as the overlapping area two Gaussian distri- butions, i.e.,
1 [ err( |1 (y0 - mo>1l 2/2
loss(wo, yo. ()
2v20
where erf is the error function, which strictly increasing. Given source image Vy and its vari- ations Co and To. with To undergoing a larger variation than to. the following conclusion can be easily observed under the same diffusion scheduling:
T So p) T(xo, To. p). s.t.l/:0 æoll ||000 xo
Here, T(20, yo p) = in > 0 loss(xp, Yo. t) > p} represents the earliest step where the loss between æo and yo exceeds the threshold p. and || measures the difference between an image and its variation. Eq. (8) provides a theoretical grounding for our intuition: images with fine- grained attribute changes tend to become indistinguishable sooner than those with coarse-grained modifica- tions under identical diffusion schedules.
We can further derive the attribute loss for DDPM inversion
losspDrM (Zo, Vo. t) = 2 1 - Vl (V mm))l
2 2(1 - (x))
Additionally, we derive that the attribute loss for DDIM inversion exhibits a similar form under certain assumptions. The detailed derivation can be found in Appendix A.2
[1
æa)l
losSpDIM Yo. t) erf
(10)
2
2 2 i= d.
2
(V
-
Because ã E 1) is strictly decreasing, the attribute loss tends to increase the time step in- creases. Furthermore, as discussed in Yue et al. (2024), this attribute loss is equivalent to how likely the DM falsely reconstruct æ, sampled from q( (2, (zp) closer yo instead of æt, and vise versa.
5

---

Published as a conference paper at ICLR 2025
5.2 ILLUSTRATIVE EXPERIMENT
In Sec. 5.1 we made a key assumption that semantically similar images are closer in latent space of the diffusion models, leading the conclusion that such images exhibit higher attribute loss a given inversion step (Eq. (8)) While this assumption is backed by recent study on the zero-shot semantic correspondence ability on diffusion latents (Tang et al. 2023 Zhang et al. 2023b), we fur- ther conducted an illustrative experiment to further support it. To validate this assumption, we used set of images, denoted as I, from the real-world task variations described in Sec. 6. 1. Specifically, we selected the 4 generalization objects in the real-world tasks from 5 categories and calculated the pairwise distance between images on their latent representations, both for intra-category and cross-category image pairs. The results are presented in Tab. 1, where each entry represents the av- erage distance between two categories. It is clear that the diagonal entries indicating intra-category similarity exhibit lower loss compared to cross- cots-ryy images, which justifies our claim.
We then conducted a user study to validate that similar images exhibit higher attribute loss at a given inversion step. We recruited 51 par- ticipants and presented each with image pairs from I after specific steps of diffusion inver- sion. The inversion steps were systematically sampled from 15 to 45, out of a total 50 steps inversion and with intervals of 5. Each partici- pant was asked to determine whether the image pairs depicted the same object, and we recorded the proportion of incorrect responses.
Table 1: Batcooogory image semantic similarity.
Categories
Bowl
Cup
Drawer
Duck
Fancet
Bowl
156.65
172.20
172.81
172.82
167.73
Cup
172.20
154.63
165.43
167.83
167.37
Drawer
172.81
165.43
144.63
161.22
161.57
Duck
172.2
167.83
16182
140.51
147.00
Faucet
167.73
167.37
161.57
147.00
145.34
Proportion of Ingorrect Responses [1%]
Canperies
Inr Cargoey
Chors Caneeryy
The experimental results in Fig. show that at inversion step 4 20, the incorrect response
Humans
0.5
within cotegortes
rate within the same category starts to increase. In contrast, the еттог rate for objects from dif-
0.4
ferent categories only started to rise at the inver-
sion step 12 35. This indicates that inversion makes objects within the same category harder distinguish earlier than affecting the distinc- tion between objects from different categories. Therefore, these results support our claim that similar images exhibit higher attribute loss at given inversion step.
Humans confuse acress categorics
15
25
35
40
Confusion Ratio vs Inversion Steps
Figure User study confusion proportions.
5.3 IMPLEMENTATION OF STEM-OB
Following the theoretical analysis and experimental validation, we propose practical implementa- tion of our framework. Our method applies diffusion inversion to converge different observations during training, and find that the model gains robustness improvement in the original observation space during test time. The detailed implementation is described below.
Training. The training process begins by applying diffusion inversion each observation. We define the following partial diffusion inversion process for an observation o; E O:
6. /T = f(o,,,
(11)
Here, 6./ E Öt/T denotes the observation after t out of T steps of diffusion inversion from the original observation o E O. The function f(.) applies steps of diffusion inversion to using any inversion methods. We select the DDPM inversion method for its efficiency and effectiveness in our experiments, and the selection of and total inversion step discussed in Sec. 6..) The visual imitation learning algorithms are then trained on the inversion-altered space Ót/T Note that we empirically find the error reduction step of DDPM inversion in Eq. is not significant for the performance, so we omit it and approximate the partial inversion process with regard to Eq. (3) only. In this way, avoid the full inversion process reverse-order corrections that involve extensive Diffusion Model inference. This approach significantly reduces preprocessing time with minimal performance impact, achieving average time for preprocessing of 0.2s per image.

---

Published as a conference paper at ICLR 2025
External Light Source
Franka
Intel D435I Camera
(b) Initial Stage
Objects Manipulate
(8) Real-world Setup Overview
(c) Success Stage
Figure 4: Realworld Setup: (a) Overview of the whole setup. (b)(c) These tasks are performed by the robot real-worll environment, from left right: Cup2Plate, Turn on Faucet, Open Drawer, and Duck2Bowl. The figure showcases the initial and final states of the tasks.
Testing. Despite training is entirely conducted on the inversion-altered space, the imitation learning algorithm of choice showcases surprising zero-shot adaptation ability to the original observation space. We adhere to the original observation space O during testing, which means essentially no changes are made to the downstream visual imitation learning algorithms. This approach demon- strates improved robustness to environmental variations without any inference-time overhead, and suitable for any other test-time augmentation techniques.
6 EXPERIMENTS
In this section, present comprehensive of experiments using Diffusion Policy (DP) Chi jtt al. 2023) as our visual imitation learning backhone. Since Stable Diffusion is pretrained on real- world images (Chi al. 2023 Rombach al.) (022), and increasing interest in deploying visual imitation learning in real-world scenarios in the community (Paolo et al. 2024). we focus primarily on evaluating Stem- OB real-world tasks. We extend the testing to simulated environments as
well for further benchmarking. To assess the robustness of our method against visual appearance variations, we design experiments featuring different object textures and lighting conditions. We compare Stem-OB against several baselines to validate its effectiveness.
6.1 EXPERIMENT SETUP
6.1.1 REAL-WORKE TASKS
We conduct real-world experiments using a Franka Emika Panda Arm, and two RealSense D435I cameras positioned different angles for RGB inputs. The setup is shown in Fig. Our experi- ments focus on four tasks, as described below:
Cup Plate (C2P): The robot arm picks up a cup and places it on plate, with variation intro- duced by changing tablecloth patterns.
Duck Bowl (D2B): The robot grasps toy duck and places it into bowl, with variations intro- duced by altering the duck's appearance, Open Drawer (OD): The robot arm grabs drawer handle and pulls it open, with variations intro- duced by modifying the drawer's visual characteristics.
Turn on Faucet (ToF): The robot arm turns on faucet, with variations introduced by altering the appearance of the faucet and bucket.
addition to the visual variations mentioned above, four tasks above involve changes in light- ing conditions, i.e., cool vs. warm light. All the variations only happen during testing time, with the training set contain only a basic setting. The object locations in training set are randomly ini- tialized within specified area, and 100 demonstrations predefined target positions are used. Further details of the environmental variations can be found in Appendix C.1
collected per task. For testing, nine

---

Published as a conference paper at ICLR 2025
6.1.2 SIMULATION TASKS
Our simulation experiments consider different tasks within two frameworks: a photorealistic simu- lation platform SAPIEN 3 (Xiang et al., 2020) and a less realistic framework MimicGen (Mandlekar
ct al. 2023).
SAPIEN 3 provides physical simulation for robots, rigid bodies, and articulated objects. It delivers high-fidelity visual and physical simulations that closely approximate real-world conditions. We leverage the ManiSkill 3 dataset (Gu et al., 2023] Tao al. 2024), collected on SAPIEN 3, for benchmarking. Specifically, we select four tasks from ManiSkill for evaluation:
PushCube:
The robot arm pushes cube target location. PegInsertionSide:
The robot arm inserts a peg into a hole.
PickCube: The robot arm picks up a cube and places on a target location.
StackCube: The robot arm stacks one cube over the other.
During testing, vary the background and lighting conditions in all tasks to generate different visual appearances. 50 episodes are tested for each setting. Details the environmental variations can be found in Appendix C.2
(a) PushCube
(b) PeglnsertionSide
(c) PickCube
(d) StackCube
Figure 5: SAPIEN environments. The figure showcases the visual appearance and task configu-
ration of each setting.
MimicGen is system for generating large diverse datasets from small number of human demon- strations. We utilize the benchmark to evaluate the performance of our approach and conduct experi- ments variety of tasks: MugCleanup. Threading, PtrreePieceAssentie offers numerous task variations, each characterized by different initial distributions. We adopt the default initial distribution (DO) for both training data and test environments. For evaluation, employ a single image as the input to the policy, using 500 samples out of a total of 1000 demos for training. We alter the table texture and object appearances create different test environments. 300 episodes are tested for each setting of all the tasks. Detailed descriptions of the environmental variations can be found Appendix C.3
and Coffee. MimicGen
6.2 BASELINES
We compare Stem-OB with several data augmentation methods aimed at improving the generaliz- ability of visual imitation learning. The baselines include SRM (Huang al.) (2022). which modifies images by adding random noise the frequency domain, Mix-Spectrum (Mix) (Lee & Hwang), which enhances the agent's 'ss focus semantic content by blending original images with randomly selected reference images from the same dataset in the frequency domain, and original images (Org) without any modification. Additionally, as highlighted in (Yuan al. 2024b). Random Overlay (RO) improves generalization in real-world experiments by blending original images with random real-world photos in the image domain. Therefore, include additional study using Random Overlay in our real-world experiments.
6.3
EXPERIMENT RESULTS
6.3.1
REAL-WORLD EXPERIMENTS
In this section, report the average success across all predefined test positions for each task. As shown in Tab. Stem-O8 consistently outperforms the baseline models all settings, demon-
8

---

Published as a conference paper at ICLR 2025
Table Evaluation of real-world experiments. Train.: evaluations in the same settings the training dataset. Gen.: evaluations under different visual perturbations for generalizability analysis. All: evaluations including both Train, and Gen. The tasks are C2P, D2B, OD, and ToF. We report the mean of the success rate (%) over corresponding settings for each task, and the hest results are highlighted in bold.
Task \
C2P
D2B
OD
ToF
Algorithm
Train.
Gen.
All
Train.
Gen.
All
Train.
Gen.
All
Train.
Gen.
All
Stem-OB
89.0
93.3
92.6
78.0
44.4
50.0
100.0
100.0
100.0
100.0
88.9
90.7
Ong
89.0
77.8
79.6
56.0
13.3
20.4
89.0
13.3
25.9
100.0
22.2
35.2
SRM
56.0
73.3
70.4
44.0
8.9
14.8
89.0
15.6
27.8
89.0
31.1
40.7
Mix
89.0
55.6
61.1
44.0
22.2
25.9
33.0
35.6
35.2
100.0
68.9
74.1
RO
56.0
53.3
53.7
44.0
15.6
20.4
89.0
88.9
88.9
78.0
82.2
81.5
Table 3: Evaluations simulated environments. compare the performance of Stem-OB, Org, SRM, and Mix in simulated environments. The left side of the table shows the results on the SAPIEN environment, while the right side is about the MimicGen environment. The results are reported as the mean and standard deviation (only for MimicGen) of the success rate(%) over their own test settings. The specific success rates are listed in Appendix E The performance of Stem-OB is highlighted in orange for better visibility. The best performance of task is highlighted in bold.
Task
SAPIEN
MimieGea
Algorithm
PushCube
PeglavertionSide
PickCube
StockCube
MugClounup
Threading
Three Pieceesssss
Coffee
Stem-OB
99.1
0.4
25.1
27.2
19,5
16.42::
13.1
50.5
Org
61.8
16.4=34
14,8-2
43.3+44
SRM
62.5
0.0
8.4
0.7
15.4250
25.8:44
13.6a25
39. 39.4=29
Mix
97.2
0.6
12.9
8.2
22.3234
18.9
14.2
41.8144
strating superior generalization capability. Under the training conditions, all methods achieve rela- tively high success rates, with our method performing slightly better than the baselines. However, in generalization testing, baseline methods exhibit significant performance drop, while our approach maintains a high success rate, showcasing the superior adaptability of Stem-OB to complex and noisy real-world visual disturbances.
The experiment results demonstrate that previous visual augmentation methods, such as SRM and Mix-Spectrum, struggle to generalize in real-world scenarios, which could be attribute to the com- plexity of real-world environments. Real images contain more redundant information, complicating the frequency domain and potentially leading to the failure of these augmentation methods. The light disturbance introduced in our experiments is a typical example, where wide-range but low-intensity noise is added to the images. Our approach effectively handles real- wallworld noise by extracting high- level structures from the appearances, resulting in better generalization. Even in challenging sce- narios like D2B, where baseline methods mostly fail, our method maintains high success rate.
Interestingly, DP without any image modification (Org) outperforms the baselines in some tasks, such as OD and ToF, suggesting that DP has inherent generalization capabilities. However, this generalizability is inconsistent across tasks. For instance, C2P, we observe that object appearance (cup and plate) had little impact on DP's performance, while the tablecloth pattern significantly affects it. In D2B, the duck's appearance critical, whereas the table cloth variation is more influ- ential. In contrast, our method exhibits consistent generalization across diverse scenarios.
6.4 SIMULATION EXPERIMENTS
Sapien 3 The evaluation results on Sapien 3 are presented on the left side of Tab. 3 Thee complexity of the tasks poses challenges for DP in generalizing across various conditions. However, our method achieves a higher success rate than the bascline methods most of the tasks. In PushCube, both Mix-Spectrum and our method perform well, but our approach is more robust and reach nearly 100% success. These results demonstrate that in more photorealistic simulation environments, our method generalizes more effectively across diverse tasks and conditions.
The evaluation results on the MimicGen benchmark are presented the right side of Tab. 3 At first glance, the results seem unanticipated, as Stem-O1 OB does not perform best in most settings. This can be attributed to the fact that MimicGen environments are less photorealistic, with
MimicGen.

---

Published as a conference paper at ICLR 2025
Table 4: Ablation study on diffusion steps. On the left side of the table, we increase the inversion step with fixed total number of steps. to intensify the effect of the diffusion process. Additionally, we compare the performance of the model with fixed ratio of steps. where the intensity of inversion is approximately the same. The results are reported as mean over 21 kinds of settings. The best performance is highlighted in orange and the second best is highlighted pink.
Task's \
Fixed Total Number of Steps (50)
Fixed Ratio of Steps (30%)
Settings
0/50
5/50
10/50
15/50
20/50
25/50
9/30
15/50
30/100
PushCube
61.8
97.1
91.1
99.1
98.5
96.3
98.1
99.1
98.3
PickCube
8,6
12.0
28.0
25.1
12.7
23.7
25.1
25.6
StackCuhe
4.3
20.7
27.2
19.6
11.6
19.5
19.6
20.2
nearly texture-free and low-resolution images. Consequently, diffusion inversion processing compli- cates the observations, limiting DP from fully leveraging the advantages of Stem-OB. It is important to emphasize that our method is primarily designed for real-world sccenarios, where diverse noise inevitable, in contrast to the controlled MimicGen environment.
6.5 ABLATION
In this section, we test with several design choices of Stem-OB, providing a better understanding of their impact the final performance. We consider the following choices: the number of inversion steps with fixed total steps, the number of total steps with the constant proportion of inversion steps to the total steps, and the selection of inversion methods (DDPM or DDIM inversion). The experiments are conducted on SAPIEN 3. We choose PushCube, PickCube and StackCube, since they have a more significant performance variance
main experiments.
Number of Inversion Steps. We compare the performance of Stem-OB with varying numbers of inversion steps, fixing the total steps to 50 and adjusting the inversion steps from to 25 in increments of 5. The results, shown in Tab. 4 indicate that performance generally increases up to 15 inversion steps before declining. This can be attributed to insufficient removal of low-level appearance features with too few inversion steps, while excessive steps eliminate high-level struc- tural information, hindering task performance. Optimal performance is observed around 10 to 15 inversion steps, depending on the complexity of the tasks.
Number of Total Steps. With the proportion of inversion steps to total steps fixed, we varied the total number of steps 30, 50, and 100. The results, shown the right side of Tab. 4 indicate that performance remains consistent regardless of the total number of steps. This suggests that the proportion of inversion steps is more critical performance than the total number of steps.
Diffusion Inversion Methods. We compare the performance of inversion methods: DDPM inversion and DDIM inver- sion, with results presented in Tab. 5 For DDIM inversion, we use 5/50 inversion steps instead of T5/50 due differences in noise scheduling compared to DDPM inversion. The results show that DDPM inversion outperforms DDIM inversion on the benchmark datasets, supporting our choice of DDPM in- version as the primary diffusion inversion method.
Task PushCube PickCube 25.1
DDPM
DDIM
99.1
81.3
5.4
Table 5: Ablation study different inversion methods. We compare the performance DDPM and DDIM on the tasks of PushCube and PickCube.
CONCLUSION
In this work, we propose Stem-OB, a straightforward preprocessing method for visual IL. By invert- ing the observation in the diffusion latent space for certain steps, we effectively converge different observation variations to the node stems from, Making invariant unseen low-level appearance changes in the observation space. Though only test method the diffusion policy method, our method is general and compatible with any visual IL basclines in theory, and the plug-and-play nature of our method makes it easy integrate. We plan to test method with other visual II. baselines in simulation and real tasks in the future.
The main algorithm of our method is simple applying the opunesourced DDPM inversion method on the dataset before training. We've provided the code for our method in the supplementary material.
Reproducibility:
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work is supported by Tsinghua University Dushi Program.
REFERENCES
Manuel Brack, Felix Friedrich, Katharia Kornmeier, Linoy Tsaban, Patrick Schramowski, Kristian Kersting, and Apolinário Passos. Ledits++: Limitless image editing using text-to-image models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 8861-8870, 2024.
Berk Calli, Aaron Walsman, Arjun Singh, Siddhartha Srinivasa, Pieter Abbeel, and Aaron M. Dol- lar. Benchmarking in manipulation research: Using the yale-cmu-berkeley object and model set. IEEE Robotics & Automation Magazine, 22(3):36-52, 2015. doi: 10.
Edoardo Cetin and Oya Celiktutan. Domain-robust visual imitation learning with mutual informa- tion constraints, 2021. URL
221.. 05079
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric A. Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. ArXiv, abs/2303. 22222333304177 0417 2023. URLhttps. // sennt
semant wtt.e/aawat.cccooooooooooooooooa
Cheng Chi, Zhenjia Xu, Chuer Pan, Eric Cousineau, Benjamin Burchfiel, Siyuan Feng, Russ Tedrake, and Shuran Song. Universal manipulation interface: In-the-wild robot teaching with- out wn-the-wild robots. arXiv preprint 7xXi:22402110329 2024.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual inversion. arXiv preprint arXiv:2208. 01618, 2022.
Jialu Gao, Kaizhe Hu, Guowei Xu, and Huazhe Xu. Can pre-trained text-to-image models generate visual goals for reinforcement learning? Advances in Neural Information Processing Systems, 36, 2024.
Daniel Garibi, Or Patashnik, Andrey Voynov, Hadar Averbuch-Elor, and Daniel Cohen-Or. Renoise: Real image inversion through iterative noising, 2024. URL htttps::/erriiw.rrg//b.
2403. 14602
Jiayuan Gu, Fanbo Xiang, Xuanlin Li, Zhan Ling, Xiqiaing Liu, Tongzhou Mu, Yihe Tang, Stone Tao, Xinyue Wei, Yunchao Yao, Xiaodi Yuan, Pengwei Xie, Zhiao Huang, Rui Chen, and Hao Su. Maniskill2: A unified benchmark for generalizable manipulation skills. In International Conference on Learning Representations, 2023.
Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yacl Pritch, and Daniel Cohen-Or. Prompt-to-prompt image editing with cross attention control. arXiv preprint arXiv:2208. 01626, 2022.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning, 2016. URL https:
// //axxiv... org/ abs/ 11.//666.0000 03346
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv preprint arxiv:2006. 11339, 2020.
Shaul Hochstein and Merav Ahissar. View from the top: Hierarchies and reverse hierarchies in the visual system. Neuron, 36(5):791-804 2002. ISSN 0896-6273. doi: https///oi.org//0. wtt.://eeet..oomccceeccccce 101660089--273300000001177 URL /
atticcle//11/9038527730002

---

Published as a conference paper at ICLR 2025
Yangru Huang, Peixi Peng, Yifan Zhao, Guangyao Chen, and Yonghong Tian. Spectrum random masking for generalization in image-based reinforcement learning, Advances in Neural Informa- fion Processing Systems, 35:20393-20406, 2022.
Ziqi Huang, Tianxing Wu, Yuming Jiang, Kelvin CK Chan, and Ziwei Liu, Reversion: Diffusion- based relation inversion from images. arXiv preprint arXiv:23033 13495, 2023.
Inbar Huberman-Spiegelllaa, Vladimir Kulikov, and Tomer Michaeli. An edit friendly ddpm noise space: Inversion and manipulations. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pullern Recognition. pp. 12469-12478, 2024.
Yuanchen Ju, Kaizhe Hu, Guowei Zhang, Gu Zhang, Mingrun Jiang, and Huazhe Xu. Robo-abe: Affordance generalization beyond categories via semantic correspondence for robot manipulation. In European Conference on Computer Vision, pp. 222-239. Springer, 2025.
Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang. Tali Dekel, Inbar Mosseri, and Michal Irani. Imagie: Text-based real image editing with diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6007-6017, 2023.
Jeong Woon Lee and Hyoseok Hwang. Mix-spectrum for generalization in visual reinforcement
learning. Available at SSRN 4818827.
Jeong Woon Lee and Hyoseok Hwang. Mix-spectrum for generalization in visual reinforcement learning. Available at SSRN 4818827, 2024.
Siyuan Li, Xun Wang, Rongehang Zuo, Kewu Sun, Lingfei Cui, Jishiyu Ding, Peng Liu, and Zhe
Ma. Robust visual imitation learning with inverse dynamics representations, 2023. URL https::
7/ arxiv.org/abs........oo.ooo 223p://bb////....o 14274
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 Robot Learning. 2023.
Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. arXiv preprint arXiv:2/08.0 01073, 2021.
Daiki Miyake, Akihiro Iohara, Yu Saito, and Toshiyuki Tanaka. Negative-prompt inversion: Fast image inversion for editing with text-guided diffusion models, 2023. URL /tttps://arxiv. org/abs 2a0.aaa/2233.68800
Ron Mokady, Amir Hertz, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Null-text inversion for editing real images using guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6038-6047, 2023.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A uni versal visual representation for robot manipulation, 2022. URL wwt.httts//eexii.orrggeeo 2203. 12601
Andrew Y. Ng and Stuart Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning, 2000. URL https: pettps://api.. /pii http://///aap.sssaiitsee..e www.sti was.mmatttccccclar......m .org/ CorpusID:26623320
Zhihong Pan, Riccardo Gherardi, Xiufeng Xie, and Stephen Huang. Effective real image edit- ing with accelerated iterative diffusion inversion, 2023. URL htttps://arxiv/
abs/
2309.044 04907
Giuseppe Paolo, Jonas Billandd,, and Balázs Kégl. A call for embodied ai. arXiv preprint 2.X2::2402.038244, 2024.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference Computer Vision and Patlern Recognition (CVPR), pp. 10684-10695, June 2022.
12

---

Published as a conference paper at ICLR 2025
Stephane Ross, Geoffrey J. Gordon, and J. Andrew Bagnell. A reduction of imitation learning and abs/
structured prediction to no-regret online learning, 2011. URL 1011. 0686]
Dvir Samuel, Barak Meiri, Nir Darshan, Shai Avidan, Gal Chechik, and Rami Ben-Ari. Regularized newton raphson inversion for text-to-image diffusion models, 2024. URL
erg/abs/2312. 12540
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models, 2022. URL 02502
https://arxi..orr//eat
Luming Tang, Menglin Jia, Qianqian Wang, Cheng Perng Phoo, and Bharath Hariharan. Emergent correspondence from image diffusion. Advances in Neural Information Processing Systems, 36: 1363-1389, 2023.
Stone Tao, Fanbo Xiang, Arth Shukla, Chen Bao, Nan Xiao, Rui Chen, Tongzhou Mu, Tse-Kai Chan, Xander Hinrichsen, Xiaodi Yuan, Xinsong Lin, Xuanlin Li, Yuan Gao, Yuzhe Qin, Zhiao Huang, and Hao Su. Maniskill 3 (beta). 2024. GitHub repository.
Narek Tumanyan, Michal Geyer, Shai Bagon, and Tali Dekel. Plug-and-play diffusion features for text-driven maage-to-image translation. In Proceedings of the IEEE/CVF Conference on Com- puler Vision and Patlern Recognition, pp. 1921-1930, 2023.
A Vaswani. Attention all you need. Advances in Neural Information Processing Systems, 2017.
Shenghua Wan, Yucen Wang, Minghao Shao, Ruying Chen, and De-Chuan Zhan. Semail: Eliminat- ing distractors in visual imitation via separated models, 2023. URL lttpps:/axxiv.. abs / 2306 10695
Chen Wang, Linxi Fan, Jiankai Sun, Ruohan Zhang, Li Fei-Fei, Danfei Xu, Yuke Zhu, and An- ima Anandkumar. Mimicplay: Long-horize horizon imitation learning by watching human play. arXiv preprint ar Xrxi:22322. 2023a.
Yanghao Wang and Long Chen. Improving diffusion-based data augmentation with inversion spher- ical interpolation. arXiv preprint arXiv:2408. 16266, 2024.
Yunke Wang, Minjing Dong, Bo Du, and Chang Xu. Imitation learning from purified demonstration. arXiv preprint arXiv:2310.071443 2023b.
Fanbo Xiang, Yuzhe Qin, Kaichun Mo, Yikuan Xia, Hao Zhu, Fangchen Liu, Minghua Liu, Hanxiao Jiang, Yifu Yuan, He Wang, Yi, Angel X. Chang, Leonidas J. Guibas, and Hao Su. SAPIEN: simulated part-based interactive environment. In The IEEE Conference Computer Vision and Pattern Recognition (CVPR), June 2020.
Tete Xiao, Ilija Radosavovic, Trevor Darell, and Jitendra Malik. Masked visual pre- pre-training for motor control, 2022. URL httpst::: /ttp/aaaivv...///////ee////aat./ wwwrxivvw........ /wtg//bb///// abs 2203. 06173
Annie Xie, Lisa Lee, Ted Xiao, and Chelsea Finn. Decomposing the generalization gap in imitation learning for visual robotic manipulation. 2024 IEEE International Conference Robotics and
Automation (ICRA), pp. 3153-3160, 2023. URL https:/ /tap:////pi... weriiiseeantiioo coommattcsccholll /orgg/ org/
CorpusID:25937552
Tianhe Yu, Ted Xiao, Austin Stone, Jonathan Tompson, Anthony Brohan, Wang, Jaspiar Singh, Clayton Tan, Jodilyn Peralta, Brian Ichter, et al. Scaling robot learning with semantically imag- ined experience. arXiv preprint arXiv22302.11550 2023.
Zhecheng Yuan, Tianming Wei, Shuiqi Cheng, Gu Zhang, Yuanpei Chen, and Huazhe Xu. Learning to manipulate anywhere: A visual generalizable framework for reinforcement learning, 2024a.
URL https:://tstiiit org/ abs 2200/ 15815
Zhecheng Yuan, Sizhe Yang, Pu Hua, Can Chang, Kaizhe Hu, and Huazhe Xu. Rl-vigen: A rein- forcement learning benchmark for visual generalization. Advances in Neural Information Pro- cessing Systems, 36, 2024b.
13

---

Published as a conference paper at ICLR 2025
Zhongqi Yue, Jiankun Wang, Qianru Sun, Lei Ji, Eric Chang. Hanwang Zhang, al. Exploring diffusion time-steps for unsupervised representation learning. arXiv preprint arX1v:2401.11430, 2024.
Yanjie Ze, Gu Zhang, Kangning Zhang, Chenyuan Hu, Muhan Wang, and Huazhe Xu. 3d diffusion policy: Generalizable visuomotor policy learning via simple 3d representations. In Proceedings of Robotics: Science and Systems (RSS), 2024.
Haoran Zhang. Yinghong Tian, Liang Yuan, and Yue Lu. Invariant adversarial imitation learning from visual inputs. In ICASSP 2023 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5, 2023a. doi: 10.1 10.1109//CSS 2022. 202.0.00777000000
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empiri- cal risk minimization, 2018. hRtps:ss://aat wttp://wwww.//b/1/11000000000010
Junyi Zhang, Charles Herrmann, Junhwa Hur, Luisa Polania Cabrera, Varun Jampani, Deqing Sun, and Ming-Hsuan Yang. A tale of two features: Stable diffusion complements dino for zero-shot semantic correspondence. 2023b.
Tony 7. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn, Learning fine-grained bimanual ma- nipulation with low-cost hardware, 2023. URL https::/aaw.....//////// /tw.////////.....////// 2304.13705
Ray Chen Zheng, Kaizhe Hu, Zhecheng Yuan, Boyuan Chen, and Huazhe Xu. Extraneousness- aware imitation learning. In 2023 IEFE International Conference on Robotics and Automation (ICRA), pp. 2973-2979, 2023. doi: 10.1 (099/CCRAA449112223..01111
Junzhe Zhu, Yuanchen Ju, Junyi Zhang, Muhan Wang. Zhecheng Yuan, Kaizhe Hu, and Huazhe Xu. Densematcher: Learning 3d semantic correspondence for category-level manipulation from single demo. arXiv preprint arXiv:2412.000 22,00268, 2024.
14

---

Published as a conference paper at ICLR 2025
A
APPENDIX
A.I DERIVATION OF DDPM INVERSION ATTRIBUTE LOSS
The derivation of DDPM inversion is straightforward by simply letting o v1 - ag in Eq.
A.2 DERIVATION OF DDIM INVERSION ATTRIBUTE LOSS
We first rewrite Eq. as linear combination of Bg and noise variable. do so, we expand the recursive equation in Eq. (5) Note that we assume that co (0t-11, t, C) is Gaussian distribution with mean Pt and use E, for simplicity the following derivation.
Ot
3t
It-1 +
It
V à1
Et
a:-1
Ūt-1
04
a1-1
21-2 +
V
1
E1-1
+
+V
(V
-
1 E1
d1-1
a1-2
at-1
Q1-2
0z
01-2
1
1 1-
-
V
1
61-2
0:-2
01-3
01-2
0x-3
1
1
(V=-
+
V
611
Er
at-1
ār-2
ār-1
1
1
=Vaz0
V
āi
āt
1
1
Varao
V a
pe +
V a.
1
āi
āi
ū;-1
(12)
where last equality holds considering the property sum of gaussian dis- tributions, i.e., N(pr,of1) + (122,0ī1)
N(1 + p2, (af + r) And
E - V /a--1 ---1------- bias of model 6g. For two similar images, the biases are approximately equal and thus cancel out when substituted into Eq. 77) Consequently, we can approximately derive the attribute loss for DDIM inversion.
be represents the mean shift resulting from the mean
a.
1 erf(
lossDDIM (Zo. yo. t)
2 2 2
(
i= a
x.-1
(13)
11:0--200 - -xoll
erf
2
2 2
1
ai
(
ac-1
B TRAINING DETAILS
The hyperparameters for Diffusion Policy across all experiments are listed in Tab, 6 We use DDPM inversion in all the basic experiments, with the specific inversion steps for each experiment detailed in Tab. 7
15

---

Published as a conference paper at ICLR 2025
Hyperparameter
value
Hyperparaeeter
value
epoch.every.n.stee.
100
aoim..om.cra.nn
imestegs
100
seq.length
frame_stack
2
ddim.beta_schedule
squa-edcos_cap__V
batch,size
128
utim_cclig.sangle
TRUE
num.epocha
1500
ddim.see...lpha.to
TRUE
Leszninggrrate-nnitiis
0.0001
soim.steps_ofiset
0.1
@cim.preddittion.tyype
epsilon
reqularization
1.2/0.0
Vi =o
obeesvaLion.
2
VisuelCCore.acccccnee class
ResNeL1ECenv
acticn.borizon
8
Vi v.sus_Corre_0ool.0..
SpatialSofttax 32
grecictiio..hoiize.
16
Vi Visus_CCore-goolnn.mm
256
FALSE
anet_dovr_dims.
[256, 512,1024]
Vi sua obs randomi zer: _laas nur.crops
unet keeekkereelssiz
4
CrupRandomize:
anet.n.groups
8
ema.power
0.75
Table 6: Hyperparameters used for Diffusion Policy. We use the same hyperparametens of diffu- sion policy for the experiments.
Task
Step
Task
Step
Task
Step
Cup2Plate
15/50
PushCube
15/50
MugCleanup
10/50
Duck2Bowl
15/50
PegInsertionSide
15/50
Threading
5/50
Drawerr
15/50
PickCube
15/50
ThreePieceAssmbbly
5/50
Turn on Faucet
15/50
StackCube
10/50
Coffee
5/50
Table Tasks and inversion steps used for training the Diffusson Policy. We all use DDPM as the inversion method.
C EXPERIMENT DETAILS
C.I REALWORLD DETAILS
Fig. showws the camera views of the real-world environments. Each task is represented by rows, corresponding to different camera angles. The columns represent task settings, with the first being the default. As lighting conditions vary across tasks, images are grouped in pairs: the odd columns are under warm light and the even columns are under cold light. The low-level object appearance changes along each row.
C.2 SAPIEN ENVIRONMENTS SETTINGS
Each task includes 21 distinct testing environments, as shown in Fig. 7. We alternate between seven unique tablecloth textures and three lighting conditions: white, yellow, and light.
C.3 MIMICGEN ENVIRONMENTS SETTINGS
Fig. demonstrates the four MimicGen environments used in the experiments, and each row in Fig. 8 represents task. In MugCleanup, the robot arm opens drawer, retrieves a mug, and places back. Threading involves the precise insertion of a stick into hole. ThreePleceAssembly requires collecting and assembling three components in a specific order. In Coffee, the robot opens the lid of the coffee machine and places cup inside. Eight different table textures used across all tasks (first eight columns). Additionally, alternating object texture, indicated the last column, creating distinct testing environments.
Threading. and T'hreePiceAAssembll employ
16

---

Published as a conference paper at ICLR 2025
Frest
Figure 6: Realworld Camera Views. Each task is represented by two rows, with each representing a different camera view angle. For each task, we perform experiments on three instances with different object appearances, and two light conditions, displayed along the columns. The first column represents the training setup.
17

---

Published as a conference paper at ICLR 2025
Figure 7: Sapien Testing Settings. Each task consists of 21 distinct settings, arranged in two rows. The images are grouped threes, with each group sharing the same tablecloth texture but differing in lighting conditions: white, yellow, and red light. A total of seven unique tablecloth textures are used. The training environment is the first setting in the first row.
Figure 8: MimicGen Testing Settings. Each row represents a specific task, with eight different table textures across the columns. The training environment the first column. Additionally, Mug- Cleanup, Threading, and ThreePieceAssembly employ alternate object textures in the last column.
18

---

Published as a conference paper at ICLR 2025
(a) MugCleanup
(b) Threading
(c) ThreePicceAssemlly
(d) Coffee
Figure 9: MimicGen environments. These tasks are based MimicGen benchmark. (a) The agent must open the drawer, pick up the mug, and place it into the drawer. (b) This task requires the agent thread string through hole. (c) The agent must assemble three pieces together. (d) The agent required to open the lid, and then place the coffee cup inside it.
C.4 DETAILS OF EXPERIMENT RESULTS
Opuninglloower
Figure 10: Realworld Success Rate. This figure presents the success rates for real-world tasks. The groups are labeled as W1, CI through W3, C3. WI to W3 are conducted under the same warm lighting condition as the training set with different object appearances. CI C3 are performed under identical cold lighting conditions, with varying object appearances. The order of the groups consistent to that shown in Fig.
Fig. 10 Fig. and Fig. 12 illustrate the success rates of Stem-OB and baseline methods in real- world, SAPIEN 3, and MimieGen environments, respectively. These three figures display the per- formance of Stem-OB under each setting with more details.
C.5 DETAILS OF ILLUSTRATIVE EXPERIMENTS
Image Distance Calculation: We use the 2-norm to compute the image distance D(x,, y) in the latent space, where smaller distance indicates greater similarity between the two images.
Intra-Category Distance Calculation: the categry distance is calculated as the mean of the pairwise distances between images within the same category I
N
N
D Dintra
D(a,,
(14)
N(N
1)
1=1 -=
where N is the number of images within I and m; is the i-th image.
19

---

Published as a conference paper at ICLR 2025
Padl dhe Seaceis Rate
Peglasciiaatile Succes Rale
PickCcbe Stacia Rate
StadkC dhe Secccse Ros
Figure 11: Sapien Success Rate. This figure presents the success rates for tasks of the SAPIEN 3 benchmark. The groups are labeled as T1 to T7 and LI to L33 The settings are varied across different tablecloth textures (TI to T7) and 3 distinct lighting conditions (L1 to L3). For TI to T7, the mean and standard deviation of the success rates are calculated over different lighting conditions with 3 lighting conditions combined for one texture. Conversely, for L1 to L3, the mean and standard deviation are computed over various tablecloth textures under a fixed lighting condition, with textures evaluated for each light condition. The order of tablecloth textures and lighting conditions aligns with that presented in Fig. 7.
MagTikanup Stccas
Thxading Stacee Baic
Theee Somernity Succese Rore
Coffee Ras
Figure 12: MimicGen Success Rate. This figure presents the success rates of Stem-OB and base-
lines across four distinct tasks. The mean and standard deviation are computed over 300 episodes. Each group of bars corresponds to one experimental setting. The bar order is consistent with the arrangement shown in Fig. 8.
Figure 13: Illustrative Experiment Objects.
20

---

Published as a conference paper at ICLR 2025
Cross- Category Distance Calculation: the cross-category distance is calculated as the mean of the pairwise distances between images from two different categories Z1 and I2
Dintra
D V2)
(15)
MN
where M and are the size of In and I,, respectively, æ; and y: E Iz.
C.6 DETAILS USER STUDY
Figure 14: Example of User Study Questions. The upper two images represent one example from the first section, and the lower two images represent one example from the second section. Each pair is generated through randomly selected inversion step. To prevent users from making judgments based on image size or orientation, these images have been randomly rotated, cropped, and zoomed.
We conducted user study evaluate the quality of the inversion results through questionnaire- based approach. the first section, consisting 42 questions, participants compare two images per question and asked determine whether the images depict the same object, different objects of the same category, or entirely different objects. Participants are informed that the images have been randomly rotated, cropped, and zoomed, so image size or orientation should not influence their judgment of object similarity. All images are snapshots of items like bowls and faucets that are used the real-world experiment, with each pair inverted by randomly selected inversion step.
The second section, consisting of 56 questions, presented images from real-world tasks, as seen in Fig. 6 In this section, participants evaluate whether the images represent the same task, different scenes from the same task (with variations such as lighting or object appearance), entirely differ- ent tasks, where the target of the task changed (for example, a cup versus a drawer). Participants are informed that scene variations could include changes in environmental conditions, such lighting, or the appearance objects, like a change in a tablecloth or the texture of a drawer, but the under- lying task function remains the same. To avoid bias from image size angle, also randomly cropped, rotated, and zoomed the images. The example images from the two sections is shown in Fig. 14. A total of 51 valid questionnaires collected, with the results presented in Fig. 3.
21

---

Published as a conference paper at ICLR 2025
Figure 15: Realworld Inversion Results. Each task corresponds to two rows, with each repre- senting a distinct camera view. The first column cach row is the original image, while the other columns the DDPM inversion results. Starting from the second column, the diffusion steps increased from 5 to 50, with the incremental step being 5.
@
Figure 16: Simulation Inversion Results. The first four rows correspond to the tasks in SAPIEN 3, while the last four correspond to the tasks in MimicCien. Each row represents task, with the first column displaying the original image and the subsequent columns illustrating the DDPM inversion results. Starting from the second column, the diffusion steps are increased from 5 to 50, with the incremental step being 5.
22

---

Published as a conference paper at ICLR 2025
Table 8: Success rate of SAPIEN tasks. Train.: evaluations in the same settings the training dataset. Gen.: evaluations under different visual perturbations for generalizability analysis. All: evaluations including both Train. and Gen. The tasks are PushCube, PegInsertionSide. PickCube, and StackCube. We report the mean and standard deviation of the success rate (%), and the best results are highlighted in bold. In the following tables, when "Train.", "Gen.", or "All" is mentioned, they carry the same meaning as defined above.
Task
PushCube
PegInsertionSide
PickCube
StackCube
Algorithm
Train.
Gen.
All
Train.
Gen.
All
Train.
Gen.
All
Train.
Gen.
All
Stem-OB
100.0
99.1
99.1
2.0
0.3
0.4
50.0
23.9
25.1
90.0
24.1
27.2
Org
100.0
59.9
61.8
0.0
0.0
0.0
90.0
4.6
8.6
90.0
0.0
4.3
SRM
100.0
60.6
62.5
0.0
0.0
0.0
80.0
4.8
8.4
16.0
0.0
0.7
Mix
100.0
97.1
97.2
2.0
75.0
9.8
12.9
84.0
4.4
8.2
Table 9: Success rate of MimicGen tasks. The tasks are MugCleamu,, Threading, PicreeliceAsemmbly and Coffee. We report the mean and standard deviation of the success rate (%), and the best results are highlighted in bold.
Task ).
MugCleunup
Threading
Three/iec//uumbll
Coffee
Algorithm
Train.
Gen.
Train.
Train.
All
Stem-OB
38.7ans
17.1
19.5as
31.3
14.5gss
16.42n
19.7atx
12:3ga
13.1a
68.0
48.0.>
50.5
Org SRM
33.0azs
14,4a15
16.4s4
41.3-1a
15.2az4
18.1
26.3:4s
13.4ts
148811
61.0zas
40.7=41
43.3
21.0g29
14.7at
15.4
24.3
26.0.aa
25.8
24.7
12,3
13.6
52.3
37,6.21
39.42
39.7
20.1
22.3
32.0.1s
17.3
18.9
17.7
13.8
14.2
55.3gs
39.8.1
41.8
D VISUALIZATION
D.1 INVERSION VISUALIZATION
Fig. 15 and Fig. 16 display the inversion results for real-world and simulation environments, re- spectively. In the real-world environments, each task spans two rows, corresponding the same camera view as in Fig. [6 In the simulation environments, each row represents task. with the first four rows showing SAPIEN environments and the last four showing MimicGen environments. The first column in each row is the original image, while subsequent columns present the inverted results. From the second column onward, the inversion results are generated by DDPM inversion, with 50 total inversion steps and 5 incremental steps from 5 to 50.
E DETAILS OF SUCCESS RATES
We separate the results for training and test variants to provide a clearer picture of the policy's performance. For SAPIEN tasks, we test 50 episodes for each task, while for MimicGen tasks, we test 300 episodes. As is shown in the Tab. g and Tab. 99] Stem-OB performs well in the generalizing settings, while remaining competitive in the training settings.
The standard deviation of the success rates is calculated below. For Mimicgen tasks, consider one single task. There are n kinds of visual generalization variations v; = 1, 1,..., n. for each setting v.. we divide the 300 episodes into groups and calculate the success rate mean μa; and standard deviation o, between each group. The overall mean for each task is the average of the individual means μ n 1n 1P and the overall standard deviation is computed as the square root average
across all settings a = / 1 /
For the SAPIEN tasks, due time constraints, in each task. As a result, we do not have enough data to group the results and compute the mean and standard deviation for each setting. We can use hypothesis testing to demonstrate that Stem-OB significantly outperforms all the baselines.
only able to test 50 episodes for each setting
Suppose the success rate in the given setting for a specific task drawn from unknown geometric distributions with unknown parameters We have observed the success rates from each group over n samples. Let P1 represent Stem-OB success rate, and pe represent the success rate of baseline. We aim to test the null hypothesis Hg Pi & P2, which is one-tailed test. The standard error of
the difference between two independent proportions is given by: SE(D)
pi (1-p))
12 b.)
19
23

---

Published as a conference paper at ICLR 2025
Table 10: Significant count. The data is presented as n/m, where n represents the number of instances where 2 > 1.645, which means Stem- significantly outperform the corresponding baseline, and m represents number of instances where : < -1645,, indicating the contrast.
Task
PushCube
PegInsertionSide
PickCube
StackCube
Org
14/0
0/0
11/1
11/0
SRM
16/0
0/0
11/1
12/0
Mix
4/0
0/0
9/1
11/0
Table 11: Success rate Additional MimicGen tasks. The tasks are Square, StackThree, Stack, and Cof- feePrepuration. We report the mean and standard deviation of the success rate (%), and the best results are highlighted in bold.
Task 1
Square
StackThree
Stack
CoffeePreparation
Algorithm
Train.
Train.
Train.
All
Stem-OB
50,7a17
28.3
31,4
41.3gas
27.6218
29.3
85.7a
70,7224
72.6.1
14,7am1
13.2.1a
13.4gx
Org SRM
53.3g5
22.0g.
25.9
44.0.m
25.7am6
28.0.xs
87.3gm
76.1
18.7
10.121
11.1gs
51.7
22.4.3
26.1
28.0
20.4.1
21.3
75.3
58.7
60.8
15.0
9,5x29
10.2
$3.7
37.5.1
39.5
35.7
18.3
20.5gm
79.7
66.8gsa
68.4g11
17.7
11.9.
12.6gs
where D = pi - p2. The test statistic for this hypothesis test is calculated as the difference D
divided by the standard error SE(D) To perform the hypothesis test at 95% confidence level, the critical value for one-tailed test 1.645. For the null hypothesis be rejected (i.e., to conclude that P1 > p2), the test statistic must exceed this critical value.
On the SAPIEN tasks, calculate the z-value of Stem-OB for each setting against all three base- lines. The number in Tab. i0 presented as n n/m, where n represents the number of instances where 1.645, which means Stem-OB significantly outperform the corresponding baseline, and m rep- resents the number of instances where = < 1.664,, indicating Stem-OB significantly underperform the corresponding baseline.
Tab. 10 shows that Stem-OB significantly outperforms the baselines in most settings across these tasks, with very few or no instances underperformance. Therefore, we can conclude that Stem-OB demonstrates superior generalization ability compared to the baselines, with statistical significance.
F ADDITIONAL MIMICGEN EXPERIMENTS
(a) Square
(b) StackThree
(c) Stack
(d) CoffeePreparation
Figure 17: Additional MimicGen environments. These tasks based MimicGen benchmark. (a) The agent must pick up square nuts and place them onto the correct pole.
(b) The agent is required to stack three cubes in the correct order. (c) The task also needs stacking but only with two cubes. (d) The agent need to place a mug on the machine, open the lid, open the drawer to get coffee pod and insert into the machine.
To further strengthen our evaluation, we add experiments for additional MimicGen tasks. We report results for an additional tasks, including Square, StackThree, Stack, and CoffeePreparation. Fig. 17 shows the images for each task. In Square, the robot arm must pick up square nuts and place them onto the correct pole. In StackThree, the robot arm is asked to stack three cubes in the correct order. Stack is similar Stack but with only two cubes. CoffeePreparation is long-
24

---

Published as a conference paper at ICLR 2025
Figure 18: Additional MimicGen Testing Settings. Each row represents specific task, with eight different table textures across the columns. The training environment is the first column.
Table 12: Success rate RO in SAPTEN. The tasks PushCube, PickCube, and StackCube. We report
the mean and standard deviation the success rate (%), and the best results are highlighted in bold. constructed a simulation dataset by overlaying images from the MimicGen dataset to adapt the RO algorithm to simulation datasets.
Task
PushCube
PickCube
StackCube
Algorithm
Train.
Gen.
Train.
Gen.
Train.
Gen.
Stem-OB
100.0
99.1
50.0
25.1
90.0
27.2
100.0
92.1
92.5
86.0
22.5
25.5
60.0
9.7
horizon task where the robot arm must prepare cup of coffee by placing a mug under the coffee machine, opening the lid of the coffee machine, then retrieving the coffee pod from a drawer and finally inserting it in the machine. The test settings are shown in Fig. [8 Tab. 11 presents the success rates for these tasks. Stem-OB achieves competitive performance across tasks, demonstrating its generalization capabilities.
G RANDOM OVERLAY BASELINE IN SAPIEN ENVIRONMENTS
The original random overlay baseline uses the Places365 dataset, which contains real-world images inconsistent with the simulation environments, thus we didn't 't take into account in simulation tasks. To adapt for simulation, constructed simulation dataset by overlaying images from the MimicGen dataset. Specifically, we randomly selected one of the seven MimicGen environments and picked an arbitrary image from random trajectory overlay onto the original image. Using this adapted RO dataset, we evaluated RO on PushCube, PickCube and StackCube. The results show that Stem-OB still outperforms RO in most tasks as shown in Tab.
H MORE GENERALIZATION EXPERIMENTS
We introduced a new type of visual variance in the SAPIEN environment by adding normal maps to the tabletop and combining this perturbation with other variations, such lighting conditions and table textures. Fig. 19 provides a visualization of these new perturbations. We then evaluated Stem-OB and the baselines on the PushCube and PickCube tasks. The results are shown in Tab. [3] Stem- OB consistently outperforms the baselines, even when encountering previously unseen pertur- bations.
25

---

Published as a conference paper at ICLR 2025
Sapien Testing Settings. Each task consists of 24 distinct settings, arranged in two rows. The first row of each task contains 12 settings, with the left half showing 6 normal only perturbation and the right half showing the normal + texture perturbation. The second row of each task represents 12 normal lighting perturbation settings, with 2 lights combined with normal maps.
Figure 19:
Evaluations in SAPIEN with normal map perturbation. normal only: evaluations with normal map perturbation only. texturt:: evaluations with normal map and texture perturbation, normal +light:: evaluations with normal map and light perturbation. report mean the success rate (%), the best results are highlighted in bold.
Table 13:
Task \
PushCube
PickCube
Algorithm
normal only
normul+texture
normal+light
normal only
normal-texture
normal+light
Stem-OB
98.0
99.3
98.7
30.0
2.3
17.7
Org
74.7
39.0
58.5
2.0
2.0
2.0
SRM
86.0
8.3
75.8
2.3
2.3
2.2
99.0
97.3
22.3
2.0
Figure 20: Cross-category PushCube Testing Settings. Each picture represents test setting random selected object from YCB object sets. The order of the setting is the same as that in the Tab. 14
Table 14: Success rate of Cross-category settings in PushCube. The original cuhe object is replaced with 6 different and random selected object from YCB object sets, such golf ball. We DDPM the inversion model this experiment. We report the mean of the success rate (%), and the best results are highlighted bold.
Settings
15/50
20/50
25/50
org
srm
mix
chef_can
0.30
0.34
0.34
0.30
0.26
0.28
golf
0.30
0.26
0.30
0.24
0.24
0.26
orange
0.34
0.42
0.44
0.32
0.34
0.42
soccer
0.30
0.40
0.40
0.28
0.30
0.34
fish
0.24
0.34
0.34
0.20
0.28
0.34
strawberry
0.20
0.24
0.28
0.16
0.28
0.26
26

---

Published as a conference paper at ICLR 2025
Table 15: Success rate of SSL in SAPIEN. The tasks PushCuue, PickCube. We report the of the success rate (%), the best results are highlighted bold. (w/ freeze) is trained with the frozen encoder, while (w/ freeze) is fine-tuning the encoder during training.
Task
PushCube
PickCube
Algorithm
Train.
Gen.
All
Train.
Gen.
All
Stem-OB
100.0
99.1
99.1
50.0
23.9
25.1
R3M(w/ freeze)
14.0
19.8
19.5
2.0
2.1
2.1
R3M(w/o freeze)
24.0
20.0
20.2
2.0
1.9
1.9
MVP(w/ firezz)
18.0
17.4
17.4
2.0
2.0
2.0
MVP(w/o freeze)
22.0
19.8
19.9
2.0
2.1
2.1
Table 16: Success rate of SD1.4 in SAPIEN. The tasks are Pushiuue, PickCube. We report the mean of the success rate (%), and the best results are highlighted in bold. We perform additional henchmarking another version of stable-diffusion: SD1.4. For SAPTEN environment, the DDPM inverison steps are 15/50.
Task |
PushCube
PickCube
Algorithm
Train.
Gen.
All
Train.
Gen.
All
Stem-OB (sd2.1)
100.0
99.1
99.1
50.0
23.9
25.1
Stem-OB (sd1.4)
100.0
99.7
99.7
36.0
18.3
19.1
While the focus of Stem-Ob's experiments mainly on intra-class generalization, we have also conducted experiments the PushCube task. In this task, the original cube object is replaced with different and random selected object from YCB object sets (Calli et al., (015), such as golf ball, is shown in the Fig. 20 We use DDPM the inversion model in this experiment. Tab. [4 shows that more inversion steps than the previous settings do help in improving the cross-category generalization ability. While most sapien tasks exhibit optimal inversion step of 15 steps, 20 and 25 inversion steps perform better on these cross-category settings. The success rate increases with the number of inversion steps, and Stem-OB outperforms all baselines in most settings. Due to time constraints, we have only tested cross-category generalization on the PushCube tasks so far. However, we believe this premitive experiment shows promising potential of generalization on multiple semantic hierarchies by using different inversion steps.
I COMPARISON WITH SELF-SUPERVISED REPRESENTATION LEARNING
Self-supervised representation learning (SSL) also aims to map raw observations into a unified rep- resentation space. We benchmark our method against R3M Nair et al. (2022) which pre-trains a visual representation using a human video dataset and incorporates a sell-superviedd loss term. This representation can then be used as a perception module for downstream policy learning. We simply replace our visual encoder (ResNet18) with the pretrained R3M ResNet18 model, and then experiment with two approaches: freezing the pretrained model and fine-tuning the encoder during training. We test R3M on PushCube and PickCube. Moreover, we try MVP Xiao et al. (2022). another SSL method, which is trained by masked modeling of natural images. The results are shown in the Tab. [5]
The experiment results show that the R3M pretrained encoder doesn't 't work well on the SAPIEN environment tasks. Since loading R3M is simple and use the official parameters, think this may be due to the heavy supervision imposed by R3M not suitable for diffusion policies.
J GENERATIVE MODELS
We conduct the experiments on other generative models, other than Stable Diffusion 2.1, to find out whether the performance of Stem-OB is model-specific. We perform additional benchmarking on another version of difle-diffusion: SD1.4 in both MimicGen and SAPIEN environments. For MimicGen tasks, 1 / the inversion steps in DDPM inversion, and 15/50 for SAPIEN
27

---

Published as a conference paper at ICLR 2025
Table 17: Success rate of SD144 in MimicGen. The tasks MugCleanpp Threading, ThreePieceAssem- bly, and Coffee. report the mean and standard deviation success rate (%), and the best results are highlighted bold. We perform additional benchmarking on another version of stable-diffusion: SD1.4. For MimicGen environment, the DDPM inverison steps are set to 1/8.
Task
MugCleamup
Threading
ThreePicceAseebbt
Coffee
Algorithm
Train.
Gen.
All
Train.
Gen.
All
Train.
Gen.
All
Train.
Gin..
All
Stem-OR (sd2.1)
19.5.4
31.3g
14.5,xa
19.7
12.3gr
13.12
68.0
Siem-OB (od1.4)
35.3
15.9141
18.0az9
32.0.ax
13.4
15,5.2s
20.7.3
14,0
14.7:24
60.0.A4
37.3
40.2
tasks. The results of the two environments are listed in the Tab, an] Tab. these two versions of stable-diffusion are similar. This supports our methods' feasibility regardless of the underlying pretrained diffusion model.
[7
The success rates of
28


Reviews:
[{"review_id": "Vzchmw7qlc", "reviewer": "Reviewer_GgNC", "summary": "This paper introduces Stem-OB, a method that enhances visual imitation learning by using image inversion from pretrained diffusion models to reduce low-level visual variations while preserving high-level scene structure. Stem-OB creates a shared representation that is robust to various appearance changes without additional training. Empirical results on several benchmarks demonstrate the effectiveness in challenging environments with lighting and appearance changes.", "questions": "Please refer to Weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* This paper focuses on enhancing the robustness of visual imitation learning, addressing a practical and impactful topic. The approach holds significant potential for advancing research in the field of robotics.\n\n* The idea of using diffusion inversion to remove low-level visual variations while preserving high-level scene structures is both novel and intriguing.\n\n* The evaluations are thorough, with experiments conducted in both simulated environments and on real-world robots. The real-world experiments highlight the method's strong generalization capabilities.", "weaknesses": "* I found the structure of this paper somewhat difficult to follow, as certain sections lacked clarity. For instance, in the preliminaries, the definition of diffusion inversion seemed to be a combination of both forward and backward diffusion. But in Figure 2, it seems that proposed method uses the noised observations as the policy input. If the proposed method is trained on the inversion-altered space, why using the noised version of the observation as input can improve the performance? How can the authors guarantee that applying forward diffusion to the image improves generalization?\n\n* There seems to be a trade-off when choosing the inversion step, which can not be either too large or too small. Is there any explanation about this phenomenon? \n\n* The theoretical analysis section was also challenging to understand. The authors discuss a loss between two latent variables,  $x_0$ and $y_0$. What  is the intuition of calculating the loss of two different images? From my understanding, the attribute loss here should refer to the inversed data $\\hat{x}_0$ and its original version $x_0$. Could the authors also clarify the statement, “images with fine-grained attribute changes tend to become indistinguishable sooner than those with coarse-grained modifications under identical diffusion schedules”? \n\n* Section 4.1 reminded me of another study [1], which employs diffusion models to purify noise within noisy demonstrations while preserving the optimal structure. Are there any conceptual similarities between that approach and the proposed method?\n\n* Typos: Line 242: $er$f should be $erf$.\n\n[1] Imitation Learning from Purified Demonstrations, ICML 2024.", "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": "vOFU1QwAxQ", "reviewer": "Reviewer_UAwF", "summary": "This paper proposes to use diffusion inversion as input preprocessing for visual imitation learning methods in order to improve their generalization to visual changes between demonstration and test setup. The proposed approach is evaluated in real-world robotic experiments as well as on 2 simulated datasets.", "questions": "The most significant questions are already listed under weaknesses. In addition, I would be interested to get a clarification on the following points:\n\n- line 272: How many images?\n- line 428: Can the authors explain how they get to this conclusion? The training setting performance of RO in particular looks to me like it was not trained correctly.\n- line 430: Given the result in D2B, I don't think the conclusion can be that there is always a high success rate?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The idea is, as the authors write, 'simple yet effective' and the authors manage to neatly package a quite theoretical idea and a very practically result measured in real-world robotic experiments all in 1 paper\n- While I lack the background to assess the theoretical grounding, the idea itself is well explained\n- The authors combine experiments on simulated datasets with actual real-world tests. Real-world tests are incredibly important for visual methods, yet in case of robotic applications very hard to make reproducible. The combination of both is a commendable experimental design. [**update after discussion phase** the extend and rigour of experiments has even further improved during the discussion phase]\n- Similarly, the authors conduct real-world user studies to investigate their hypothesis that diffusion inversion progresses along semantic hierarchies.", "weaknesses": "- [**update after discussion: this point was fixed**] I find the derivation in Section 4.1 very hard to follow. Only the first sentence of Section 4.2 makes it clear that the whole derivation is based on the assumption that semantically similar images are closer in latent space. That is a HUGE assumption and the following experiments on intra- and inter-class confusion do not show this over multiple levels of semantic hierarchies, but only for the semantic leafs. The introduction (lines 60-65) however argues that much more than leaf classes, abstract semantic hierarchies are important to generalize over perceptual differences.\nIn the empirical study on diffusion inversion, one might argue that out of the 5 empirical examples, \"bowl\" and \"cup\" are most similar. However, Table 1 shows that the latent representation of 'cup' samples is actually on average closer to any of {'drawer', 'duck', 'faucet'} than to 'bowl'. To me this makes the assumption 'semantically similar images are closer in latent space' quite unbelievable. For sure, this assumption and therefore the formulation of semantic overlap as overlap of latent gaussians is anything but 'intuitive' (line 270) in presence of this data.\n- [**update after discussion: this point was fixed**] In relation to the above, I miss a clear definition of what kind of variation should be compensated through diffusion inversion. The abstract and introduction repeatedly claims that diffusion inversion can extract the 'high-level structure' of the scene, suggesting generalization over different object instances, shapes, appearances, relative placements, and robot arms. Section 4.1 considers 'variation', 'fine-grained attribute change', 'coarse-grained modifications', and 'semantic overlap' without defining any of these. The investigation in Section 4.2 is focused on variations within a semantic object class, i.e. a demonstration with 1 cup should be repeated with a different cup. The experiments then consider lighting change and a limited set of object appearance change of sometimes multiple object instances, while locations are fully fixed. The problem is that all of this currently does not fully fit together and it would be good if the authors can define more accurately what kind of variations they expect diffusion inversion to abstract / generalize over, and then design experiments accordingly to show improvement with exactly these variations.\n- There are a couple of odd aspects about the simulation experiments that raise questions about the soundness of the results:\n  - For the benchmarks from ManiSkill and MimicGen, why were not all tasks evaluated? For ManiSkill, plug-charger seems to be specifically excluded and for MimicGen 4 out of 12 tasks were picked without any explanation.  \n[ **update after discussion: now there are 8 out of 12 tasks from MimicGen. It is still a subset, but this indeed makes the results more thorough**]\n  - While I did not quickly find comparable numbers for ManiSkill, the MimicGen paper reports much higher success rates for the investigated tasks. E.g. for Threading, MimicGen reports around 19% sucess from just 10 videos and 98% success from 1000 videos. For the 500 videos used in the experiments here, the success is below 19% for 3 out of 4 variants, including the proposed method. Why are the achieved success rates so low? And can the proposed method actually improve anything in a more state-of-the-art setting? [**update after discussion: This concern partially remains. While the authors show that their success rates for 500 trials are mostly between the baseline MimicGen results of 10 and 1000 trials, there is no case where the authors can demonstrate to improve over these prior results. If one expects the relation between number of trials and success rate to follow a saturation trend rather than a linear trend, their results are still a bit below those of prior works.**]\n  - Why is the RO baseline excluded from the simulation experiments? [**update from discussion: this was resolved**]\n- In all experiments (real and simulated), most of the differences between methods are within the standard deviation, so it is very hard to say if any conclusion can be drawn. Why is the standard deviation so high? [**update from discussion: this has been thoroughly addressed by the authors. The updated results show, similarly to the question about success rates above, that Stem-OB does not push the state-of-the-art further than where it is. But it does more importantly also not appear to be much worse, while the method in some sense is simpler.**]", "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": "izIBU6q0CI", "reviewer": "Reviewer_C5Ru", "summary": "This paper proposes Stem-OB, a method that uses diffusion inversion to transform observations into a shared, robust representation. This approach enhances robustness to various appearance changes, such as shifts in visual patterns and lighting conditions. During inference, Stem-OB approximates the partial inversion process to maintain efficient inference speed.", "questions": "The current method conducts experiments solely on Stable Diffusion. It would be valuable to understand if the conclusions drawn from Stem-OB are applicable to other generative models, such as Flux or SD 3 (which uses flow instead of diffusion). Demonstrating that Stem-OB generalizes across different models would strengthen the claim of robustness and broaden the impact of this approach beyond a single model framework.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well-written, with clear motivation and solid experimental validation, including both real-world and simulated experiments.\n\n2. The central idea resonates well: robotic observations often include excessive low-level details, while effective scene understanding requires capturing high-level structural information. This paper draws on an intriguing insight from Yue et al. (2024): \"Rather than uniformly removing information from different semantic hierarchies, the process brings structurally similar images closer in the early stages of inversion.\" Building on this, Stem-OB leverages this observation to project data into a space that prioritizes inherent structural features.", "weaknesses": "Stem-OB enhances the generalization capabilities of robotic systems by using diffusion inversion to map observations into a shared representation. However, another promising line of research—self-supervised representation learning—also aims to unify raw observations into a common representation space, reducing low-level visual discrepancies. I suggest that the authors consider benchmarking Stem-OB against these approaches, particularly methods that leverage self-supervised learning (SSL) for robotic representations, to offer a more comprehensive comparison.", "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": "jJDbN0CV9s", "reviewer": "Reviewer_pwGy", "summary": "Visual imitation learning is a method for robots to learn tasks through visual human demonstrations. While multiple methods exist for preprocessing visual data that increase the generalizability of VIL, these methods are only partially efficient. To address the generalization issue, this paper proposes to use the inversion process of a pretrained diffusion model as a preprocessing step to increase the generalizability of visual imitation learning methods. The proposed method relies on the intuitive insight that images with only small perturbations between them are closer in latent space (and therefore become indistinguishable under diffusion faster), than data that exhibits large differences. The method is tested on real and synthetic data and shows that the proposed preprocessing significantly increases the models generalization capabilities.", "questions": "See above. I currently give a weak accept to the paper but am inclined to vote for acceptance should my points above be taken into consideration during the rebuttal.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper relies on a central observation from a previous paper (Yue et al. 2024), which found that during the inversion process of a diffusion model, fine-grained (or high-frequency) details of an image are destroyed first before the low-level semantic concepts. This paper builds on this basic concept and proposes a new preprocessing step that uses this theoretical property of diffusion models as a preprocessing step, which, in the process, improves the generalization of the IL algorithm. \n* The technical principle seems well established at this point, but this paper takes these theoretical findings and applies them to a new problem. While the whole approach is very simple from a technical standpoint, the paper re-introduces the necessary background in section 2 before developing the reader's intuition as to why this preprocessing might work in section 4.2 through theoretical explanation and small experiments. I believe that this paper warrants publication because it shows that diffusion model inversion can be useful for additional tasks. \n* The presentation of the paper is clear and easy to follow overall. \n* The results nicely highlight the drawbacks of previous methods and that the proposed method offers much better generalization capabilities. \n* I appreciated the detailed appendix, with theoretical derivations, additional details for reproducibility, and additional experiments.", "weaknesses": "I believe that the paper's technical novelty is small, as the main concept has already been introduced in other papers. That being said, I also believe that the paper is worth publishing, as it adds to the growing body of literature showing the usefulness of diffusion model inversion for various tasks, and the empirical validation is well executed.\n\nI have additional suggestions to improve the paper's readability and make it easier for the reader to understand.\n1. Section 2.2 can be cut (title), and the paragraph can be integrated into Section 3.1. \n2. As a reader, I would really like a short problem formulation that properly introduces the problem the paper addresses. This should properly define the inputs and outputs. I would also suggest rearranging the sections/subsections of the paper to something like this:\n* Introduction\n* Related Work\n* Problem Definition\n* Preliminaries\n* Method\nThis would greatly improve the flow of the paper. \n3. While the paper is generally easy to follow, some sentences are should be simplified or run through a writing program. I have added some examples below:\n* Line 85: \"To be specific, our method is as simple as inverting the image for reasonable steps before\" What is reasonable? Maybe rephrase?\n* Line 226: \"Intuitively, given a source image and its variation, distinguishing between them becomes increasingly difficult\" This sentence does not make a lot of sense (although I can guess what is meant).\nThere are many more of these convoluted sentences. Cleaning up the writing a bit would go a long way to improve the paper.\n\nThere", "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)