Stem-OB: Generalizable Visual Imitation Learning with Stem-Like Convergent Observation through Diffusion Inversion

We thank you a lot for the detailed recognition of our work in the strength section and the thoughtful suggestions on the paper to make it better deliver our main idea!

As the suggestions mainly focus on the writing and structure of the paper, we've made corresponding changes in the revised version, the change of the paper can be summarized as follows:

1. We rephrased and put the content of Section 2.2 to Section 3.1 as suggested, and canceled the sole subsection title of Section 2.1. This makes the preliminary part more focused on introducing the important idea of diffusion inversion.
2. We appreciate the suggestion and rearranged the sections accordingly, and we added a problem definition section which makes it clear that we are dealing with the visual generalization of imitation learning methods through operating on the observation before the training pipeline.
3. We modify line 85 to make it clear that the appropriate inversion step is selected via empirical study, balancing between a strong perturbation and better preservation of semantic information.
4. For line 226, we make it clear that the variation here refers to a visually altered scene of the source image that can include low-level details or high-level structures.
5. We also make other modifications in sections 5.1 and 5.2 along with the introduction to further clean up the writing.

Thank you again for the valuable suggestions, and any further thoughts or questions are more than welcome!

We sincerely thank the reviewer for the thoughtful feedback and for recognizing our contributions by revising the score.

We thank the reviewer for acknowledge our method's clear motivation and "solid experimental validation", and will address the concerns as follows:

Q1: Benchmarking with SSL methods.

A1: Thank you for pointing out this baseline for us! Self-supervised representation(SSL) learning also aims to unify raw observations into a common representation space.

We benchmark our method against R3M [1]. It pre-trains a visual representation using the Ego4D human video dataset including a self-supervised term, which can be can be used as a perception module for downstream policy learning. We then change our visual encoder from resnet18 to the pretrained R3M resnet18 model. We have tried two ways, one is freezing the pretrained model, and the second is finetuning the encoder during training. Due to the time limit, we test R3M on PushCube and PickCube. The results are shown in the table below.

The experiment results show that the R3M pretrained encoder doesn't work well on the SAPIEN environment tasks. Since loading R3M is simple and we use the official parameters, we think this may be due to the heavy supervision imposed by R3M not suitable for diffusion policies. We plan to benchmark other SSL methods like MVP[2] and will report the results once the results are available.

[1] Nair, S., Rajeswaran, A., Kumar, V., Finn, C., & Gupta, A. (2022). R3m: A universal visual representation for robot manipulation. arXiv preprint arXiv:2203.12601.

[2] Xiao, T., Radosavovic, I., Darrell, T., & Malik, J. (2022). Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173.

Q2: Other Generative Models

A2: Thanks for the suggestion! We perform additional benchmarking on another version of stable-diffusion, SD1.4, as we conducted some early experiments on this model other than the SD2.1 in the MimicGen environment. We then conduct additional experiments in the SAPIEN environment for further validation. For MimicGen tasks, we use 1/8 as the inversion steps in DDPM inversion, and 15/50 for SAPIEN tasks. The results of the two environments are listed in the table below.

As we can see in the table, the success rates of these two versions of stable-diffusion are similar. This supports our methods' feasibility regardless of the underlying pretrained diffusion model.

The authors' response has effectively addressed my concerns, so I am increasing my rating to 8.

We thank the detailed and in-depth review provided by reviewer UAwF, we view much of the raised concerns as legitimate and conduct additional experiments to verify them, while there are cases where our results are supposedly misunderstood due to unclear statements.

Q1: Feasibility of theory assumption

A1: We appreciate the reviewer's detailed feedback and the opportunity to clarify our theoretical assumptions. The assumption that semantically similar images are closer in the latent space of pretrained diffusion models is fundamental to our approach and is supported by both prior research and empirical results.

Firstly, as outlined in [1], the concept of attribute loss is inherently linked to the indistinguishability of different images during the inversion process. Attribute loss quantifies the likelihood that a diffusion model fails to distinguish noisy samples drawn from two overlapping distributions, $q(x_t|x_0)$ and $q(y_t|y_0)$. This involves:

1. Performing diffusion inversion on $x_0$ and $y_0$ for $t$ steps to obtain $x_t$ and $y_t$.
2. Reconstructing $x'_0$ and $y'_0$ from $x_t$ and $y_t$.
3. Assessing reconstruction similarity, where semantically similar images are more likely to produce similar reconstructions, resulting in a higher attribute loss.

We adopt this refined definition in our paper (Equations (6) and (7)). By focusing on semantic leaf nodes in Section 4.2, we precisely control and measure the initial similarity between $x_0$ and $y_0$. This aligns with the formulation of attribute loss, which relies on known initial similarities and this does not mean that the distance between latent features in other semantic hierarchies is needed to fulfill the same assumption.

Regarding the observations in Table 1, we acknowledge that the latent representations of 'cup' samples appear closer to some other classes than to 'bowl'. This can be attributed to the variability in visual features captured by the latent space, influenced by factors such as viewpoint, texture, and shape. While 'cup' and 'bowl' can be viewed as semantically similar, the latent space may not always reflect this due to these factors. Our main objective is to demonstrate that intra-class variations (within the same class) result in closer latent representations than inter-class variations.

Importantly, the assumption that semantically similar images are closer in the diffusion model's latent space is supported by recent studies in semantic correspondence. Tang et al. [2] and Zhang et al.[3] demonstrated that pretrained diffusion models' latent spaces can effectively map semantically similar pixels across images using simple nearest-neighbor searches. This suggests that the latent space does capture semantic similarities to a significant extent.

Our user study further validates this assumption by showing that semantically similar images become indistinguishable sooner during the diffusion inversion process. This empirical evidence supports the behavior predicted by attribute loss and underscores the effectiveness of our approach.

We will clarify these points in the revised manuscript and provide additional evidence and references to strengthen our theoretical foundation.

[1] Yue, Z., Wang, J., Sun, Q., Ji, L., Chang, E. I., & Zhang, H. (2024). Exploring diffusion time-steps for unsupervised representation learning. arXiv preprint arXiv:2401.11430.

[2] Tang, L., Jia, M., Wang, Q., Phoo, C. P., & Hariharan, B. (2023). Emergent correspondence from image diffusion. Advances in Neural Information Processing Systems, 36, 1363-1389.

[3] Zhang, J., Herrmann, C., Hur, J., Polania Cabrera, L., Jampani, V., Sun, D., & Yang, M. H. (2024). A tale of two features: Stable diffusion complements dino for zero-shot semantic correspondence. Advances in Neural Information Processing Systems, 36.

Dear reviewer UAwF: Thank you for your reply, we may ask for some clarifications on your comment.

Generally speaking, do you think the current name is not appropriate since it implies our method has the ability to perform generalization in "multi-level semantic hierarchies", which it doesn't? If so, there are a few questions:

1. Can you give us a few examples to better understand the meaning of multi-level semantic hierarchies here? For example, at the base level, there are different types of cups and at the medium level, there are cups and bowls, etc.
2. Do you think the ability of inter-category generalization is the key to proving that our method can generalize across different semantic hierarchies? If so, could an additional experiment about, for example, generalizing to pick up cups from demonstrations of picking up bowls prove such ability?
3. What level of semantic hierarchy do you wish to see for generalization, since we have shown the ability to generalize across different lighting conditions, different object surface patterns, and different normal maps (equivalent to the shadow over the surface)
4. Lastly, could you please give us a hint on which part of the current title you think is the most inappropriate? Maybe a change in the paper title could better deliver our contribution.

We would like to argue that the different lighting conditions and the appearance of the same category of object are the most useful variations in practice, which Stem-Ob has already demonstrated through simulation and real-world experiments. If we are generalizing across categories, the definition of the task becomes ambiguous. For example, putting a lid on a cup is not similar to putting a lid on a bowl, thus making it less useful in real tasks. While Stem-Ob could potentially perform such kind of generalization given more inversion steps, we think the upper level of the semantic tree is the most worthy to generalize since these objects share the same function and can be easily generalized by humans.

Thank you for your comment and they are of great value for us to improve the paper.

We thank reviewer GgNc for acknowledging our novelty and thorough evaluation! We'll address the concerns as follows:

Q1: Structure somewhat difficult to follow.

A1: Thank you for highlighting this point. We would like to clarify that diffusion inversion refers to a process where an original image is mapped to a noisy image while preserving semantic information to allow reconstruction during a separate backward process. The "inversion-altered space" consists of observations injected with inversion noise, which effectively emphasizes high-level structural information while erasing out lower-level details.

During the inversion process, different images gradually lose their fine-grained details and converge towards Gaussian noise. the general idea of why an inversion-altered observation can improve the generalization performance is strongly related to our theoretical intuition: images that differ only in lower-level details become indistinguishable at an earlier inversion step. Consequently, observations of the same object under various visual perturbations converge to a unified "stem observation." Training policies on these unified observations enhance generalization to new environments, even before applying any downstream visual imitation learning (IL) pipeline.

We have revised the relevant part in the preliminary part of the paper to better convey these concepts and improve clarity.

Q2: Trade-off between inversion steps

A2: This observation is indeed accurate, and the underlying reason aligns with our theoretical framework. As diffusion inversion merges different visual variations of the same scene, using too few inversion steps results in incomplete merging, which hinders the generalization ability. Conversely, excessive inversion may impact the semantic information critical to the task. In practice, we find a balanced inversion step of 15/50 to be effective for most real-world tasks and 10/50 for simulation tasks.

Q3: Theoretical analysis ambiguity

A3: As noted in the paper, the concept of attribute loss originates from [1]. Attribute loss quantifies the similarity between two images after a specified number of diffusion inversion steps, without imposing constraints on how these images are related. Accordingly, we adopt a general notation, $x_0$ and $y_0$, to represent two distinct images in Equation (6), consistent with the notation in [1]. A more formal definition of attribute loss is quantifying how likely DM fails to distinguish the noisy samples drawn from two overlapping noisy sample distributions $q(x_t|x_0)$ and $q(y_t|y_0)$ [1].

In our context, the purpose of using attribute loss is to quantify the semantic differences between an image $x_0$ and its variations $\hat{x}_0$ and $\tilde{x}_0$ with diffusion inversion step $t$. This is represented as $loss(x_0, \hat{x}_0, t)$ and $loss(x_0, \tilde{x}_0, t)$. We want to clarify that $\hat{x}_0$ and $\tilde{x}_0$ here are not the inversed data. They are source images without undergoing any inversion process. However, they exhibit some degree of similarity or relation to the original image $x_0$, which can be regarded as a deviation from $x_0$ in specific ways. For example, given a scene image $x_0$ shown in Fig. 4 (a), $\hat{x}_0$ indicates $x_0$ with only fine-grained attribute changes such as changing the texture on some of the objects. In contrast, $\tilde{x}_0$ indicates $x_0$ with coarse-grained modifications such as drastically changing the lighting condition of the scene. In this context, granularity refers to the pixel-level alterations introduced when modifying an attribute of the image. And what attribute loss indicates is $\hat{x}_0$ (fine-grained attribute change) becomes indistinguishable sooner than $\tilde{x}_0$ (coarse-grained attribute change) during diffusion inversion (equation (8)).

[1] Yue, Z., Wang, J., Sun, Q., Ji, L., Chang, E. I., & Zhang, H. (2024). Exploring diffusion time-steps for unsupervised representation learning. arXiv preprint arXiv:2401.11430.

Dear Reviewer GgNC,

As the rebuttal period is coming to an end, we want to kindly check if there are any remaining questions or concerns regarding our responses to your review. We deeply value your feedback and are willing to provide any additional clarifications if needed.

Thank you for your time and thoughtful assessment of our work!