EC-Diffuser: Multi-Object Manipulation via Entity-Centric Behavior Generation

We appreciate the reviewer’s acknowledgement of our method’s efficacy and generalization capabilities and of the clarity and reproducibility of our work.

Novelty

To the best of our knowledge, we are the first to use diffusion models on top of unsupervised object-centric image representations for decision-making, as well as the first to demonstrate such zero-shot compositional generalization capabilities in the behavioral cloning setting.

A major challenge in BC is capturing multi-modal action distributions often present in offline data. This multi-modality is present even in single object manipulation, but increases combinatorially with the number of objects. Diffusion models have demonstrated success in capturing diverse action distributions in the BC setting but without entity-centric structure, they fail catastrophically when scaling the number of objects (as our experiments reveal). One aspect of our work’s novelty is identifying and demonstrating that the above combination works well in the context of multi-object manipulation. To achieve this, we developed a novel Transformer architecture that includes conditioning on actions by incorporating them as a particle state as well as using Adaptive Layer Normalization throughout the denoising diffusion process. In addition, due to the unique structure of the DLP representation, our design choices include equivariance considerations, such as removing positional embeddings.

We would like to respectfully address a few misunderstandings that may be reflected in the reviewer's comments:

Dependance on DLP – A key aspect in our method is the use of a good object-centric image representation, and one of our goals in this work is to highlight the advantages of leveraging such representations. However, the components of our method that build upon the DLP representation are not inherently tied to the object-centric factorization it provides. To demonstrate this point we conduct experiments showcasing our method’s compatibility with alternative object-centric representations, including Slot-Attention as well as ground-truth state representations. A full analysis and results can be found in Appendix D.2 of the new revision. Our method without any object-centric representation nor entity-centric architecture (the unstructured ablation in Table 4), no longer possesses any key aspects of our method and is essentially Diffuser using a Transformer instead of a U-Net architecture. This ablation is comparable to VQ-BeT and thus it is not surprising it achieves equivalent performance.

Our Transformer Architecture – We develop a novel Transformer architecture for our purposes, as described above. We do not use nor adapt VQ-BeT’s architecture.

Choice of Diffuser – Our method builds upon the Diffuser framework. The best performing baseline is an integration of the core aspects of our proposed method in an action-only generating diffusion model, which is equivalent to Diffusion Policy. It was not clear that the entity-centric Diffuser would perform better than the entity-centric Diffusion Policy in our setting. Part of our contribution is identifying that co-generating states is a key aspect to performance and our analysis of why this is the case (Section 5.1, end of second paragraph and Section 5.4).

Real-world Evaluation

We kindly refer the reviewer to our general response regarding real world applicability. In addition, we provide answers to the reviewer’s questions:

Can this model be evaluated on real-world datasets such as NuScenes (as VQ-Bet does)?

The nuScenes self-driving environment assumes access to an object-centric observation decomposition. Furthermore, it assumes access to object classes and object tracking information (see Figure 11 in the last page of the Appendix of the VQ-BeT paper). In contrast, we do not make these assumptions and propose a method that acquires a representation from images in an unsupervised manner and can leverage it for sequential decision-making without additional assumptions on the structure of the decomposition other than that it is object-centric. While autonomous driving is an interesting domain to incorporate entity-centric structure, the variations in assumptions and the fact it is not an object manipulation environment (which we clearly convey is the focus of our work) are the main reasons we do not include such experiments.

Can EC-Diffuser be adapted for use as GPT-driver?

To our understanding, GPT-Driver is a language-model-based approach for motion planning for autonomous driving. We kindly ask the reviewer to clarify the term “adapted for use as GPT-Driver”. Our work uses a Diffusion process, and not language modeling, and we do not rely on LLMs for training or inference.

We do believe EC-Diffuser can be adapted to the autonomous driving domain, and this is an interesting avenue for future work.

We thank the reviewer for acknowledging the novelty and contribution of our work and for their genuine interest and attention to detail, providing insightful suggestions for strengthening our paper.

Alternative Object-centric Encoders

To the reviewer’s suggestion, we have conducted experiments with alternative object-centric representations, including Slot Attention as well as ground-truth state representations. A full analysis and results can be found in Appendix D.2 of the new revision. In addition, we summarize our findings here for the reviewer’s convenience:

Ground-truth State - To shed light on the generality of our approach, we experiment on the PushCube environments using object-centric set-based state representations by extracting the object location information from the simulator and appending a one-hot identifier to distinguish each object (a similar baseline was considered in ECRL). This representation is meant to emulate a “perfect” entity-level factorization of images. Our experimental results show that using this representation, the EC-Diffuser achieves slightly better performance than using the DLP representation, as expected. We believe that these results highlight two important aspects: (1) The benefit of entity-centric structure, which motivates such factorizations of images. (2) EC-Diffuser’s ability to handle different types of object-centric representations (the state-based [position, 1-hot] is very different from the DLP representation).

Slot Attention - We replace the DLP encoder with a Slot Attention encoder and experiment on the PushCube environments as well as FrankaKitchen. Performance on PushCube using slot-attention-based representations is better than non-object-centric baselines but worse than learning using DLP representations. On the FrankaKitchen environment, EC-Diffuser using Slot Attention surpases DLP, achieving a new SOTA of $3.340$ compared to the previous $3.031$ obtained by EC-Diffuser using DLP. Judging by the slot decompositions of the different environments (see Figure 13 and 14 in the Appendix of the new revision), it seems that Slot Attention occasionally has trouble with individuating nearby objects and represents them in a single slot. We believe this is a major factor in the performance drop in PushCube with multiple cubes. In FrankaKitchen this does not affect performance, we believe this is due to the fact that most objects in the scene are static (and usually allocated together in the same slot, and the robot in another slot, leading to a foreground-background decomposition) thus not necessarily requiring an accurate decomposition for good downstream control. We emphasize that the object-level state complexity in PushCube and PushT is much higher than in FrankaKitchen, making the former significantly more challenging.

In conclusion, these experiments demonstrate the generality of our method with respect to compatibility with different object-centric representations. We further motivate the specific choice of DLP in our general response.

More Generalization Experiments

To the reviewer’s suggestion, we evaluate our PushCube generalization policy (trained on 3 cubes in different colors chosen from 6 options) on the following:

• Replacing cubes with star-shaped objects
• Replacing cubes with rectangular cuboids
• Replacing cubes with T-shaped blocks
• Introducing unseen cube colors

The results can be found in Appendix D.3, Table 14. We see that EC-Diffuser coupled with DLP is able to generalize zero-shot with little to no drop in performance to new colors as well as new shapes (star, rectangular cuboid). When replacing cubes with T-shaped blocks there is a significant drop in performance although success rate is better than random, suggesting some zero-shot generalization capabilities in this case as well. We see that our policy handles new objects well in cases where they behave similarly in terms of physical dynamics and less when they are significantly different, which is expected.

We did not experiment with new camera views but we do not expect the policy to map observations to the correct actions due to the frame-of-reference shift, requiring re-grounding of observations to actions. Policy adaptation via fine-tuning techniques could help deal with novel views and possibly with novel objects with significantly different dynamics.

Presentation Suggestions

We thank the reviewer for their suggestions. We have fixed the rendering issue with Figure 5 and have indicated the training setup in Tables 2 and 3 in the column header instead of using bold text.

We thank the reviewer for their acknowledgment of the efficacy of our method compared to competitive baselines and for their questions and interest in our work.

Dependance on Object-centric Image Representation

Our method strongly depends on a good object-centric image representation, and one of our goals in this work is to highlight the advantages of leveraging such representations. However, the components of our method that build upon the DLP representation are not inherently tied to the object-centric factorization it provides. In principle, DLP can be replaced with other object-centric representations of images such as slot-based representations. We refer the reviewer to the Slot Attention experiment in Appendix D.2, Table 12 and 13 of the new revision, where we demonstrate the performance of our method coupled with Slot-Attention instead of DLP.

The simplest way to define a good object-centric representation is one that effectively facilitates downstream decision-making. As our experimental results demonstrate, DLP not only fulfills this criterion but also enables zero-shot generalization in certain scenarios. We elaborate on why we believe DLP is better suited than slot-based models for the types of tasks we address in the general response under Choice of Object-centric Representation.

Simulated Experimental Environments

We would kindly like to point out that the simulated environments we use in our work are either equally or more challenging than any simulated environment baseline methods have tested on in terms of multi-object complexity, goal space complexity, occlusion conditions and simulation visual quality.

Most of the simulated environments in VQ-BeT are either not image-based or not related to the focus of this work which is object manipulation. The environments that involve object manipulation from pixels are FrankKitchen (which we experiment on), and PushT. Notice that VQ-BeT outperforms baselines on goal-conditioned ‘Image PushT’ but obtains relatively poor results (0.1 out of 1, see Table 1 in the VQ-BeT paper). We believe that VQ-BeT completely fails in multi-object manipulation from images because this is a very hard problem to solve directly from raw pixel observations without incorporating any prior knowledge of the world in the form of (entity-centric) structure. This is precisely what we aim to demonstrate in our experiments and why we chose to deploy an excellent SOTA algorithm such as VQ-BeT on environments with high multi-object complexity, proving this point.

In the following, we provide details about each environment in VQ-BeT in comparison to the related environment in our work:

BlockPush and UR3 BlockPush: These environments are solved in the VQ-BeT paper from simulation state and not from images. In addition, the start and goal locations of the cubes are small fixed regions on the table and the goal variability comes from which cube is pushed to which one of the 2 goals. In contrast, the PushCube IsaacGym environment is perceived from images, and start and goal configurations are sampled uniformly at random on the table surface at the beginning of each episode.

PushT: The simulated environment introduced in the Diffusion-Policy paper that was also used in VQ-BeT is a 2D environment with a single T-block and a point robot, where the goal is visually present on the manipulation surface. This type of goal-specification is less practical than providing a goal image (would require drawing the goal on the table before each episode) and makes it easier for a standard CNN to track progress in reaching the goal based on how much of the goal ‘shadow’ is occluded. In contrast, the PushT IsaacGym environment is 3D, contains multiple T-blocks, has a full 7-DOF robot arm which requires handling occlusion, goals are specified by separate images requiring that the agent learn the relationship between entities across images.

Multimodal Ant: This is a locomotion environment where the observation space is the simulation state and not images. Our work focuses on object manipulation from pixels, therefore we did not experiment with this environment. It could be interesting to see if an entity-centric structure provides value in learning locomotion tasks from pixels, although we believe the advantage over unstructured image representations in this case would be less apparent.

Dear Reviewer,

Now that the authors have posted their rebuttal, please take a moment and check whether your concerns were addressed. At your earliest convenience, please post a response and update your review, at a minimum acknowledging that you have read your rebuttal.

Thank you, --Your AC

We thank the reviewer for acknowledging our method’s generalization capabilities, and for their in-depth questions. We provide detailed answers below, along with new results from additional experiments.

Novelty

To the best of our knowledge, we are the first to use diffusion models on top of unsupervised object-centric image representations for decision-making, as well as the first to demonstrate such zero-shot compositional generalization capabilities in the behavioral cloning (BC) setting.

A major challenge in BC is capturing multi-modal action distributions often present in offline data. This multi-modality is present even in single object manipulation, but increases combinatorially with the number of objects. Diffusion models have demonstrated success in capturing diverse action distributions in the BC setting but without entity-centric structure, they fail catastrophically when scaling the number of objects (as our experiments reveal). One aspect of our work’s novelty is identifying the potential of this combination and demonstrating that it can solve challenging multi-object tasks from image observation data in a manner that facilitates zero-shot compositionally generalizing behaviors.

To achieve this, we developed a novel Transformer architecture that includes conditioning on actions by incorporating them as a particle state as well as using Adaptive Layer Normalization throughout the denoising diffusion process. In addition, due to the unique structure of the DLP representation, our design choices include equivariance considerations, such as removing positional embeddings. Finally, we demonstrate SOTA results on challenging simulated robotic manipulation environments compared to strong baselines. We compare with prior works that do not use object-centric representations (VQ-BeT, Diffuser), works that use an object-centric representation but do not generate states (Diffusion Policy) or do not use diffusion (EIT+BC) to showcase the advantages of our specific design choices.

While our work builds on some aspects of ECRL, applying them in our setting is not trivial. In this work we show that a naive adaptation of ECRL to the behavioral cloning setting is not able to learn multi-object manipulation (see results for EIT+BC baseline in Table 1), demonstrating that our approach of using diffusion models and our novel architecture accounts for the distinct challenges that arise when learning from offline demonstration data containing diverse behaviors.

Real Robots

We kindly refer the reviewer to the general response regarding real-world applicability.

Questions

Is the DLP pretraining done separately for each task in IsaacGym environment?

DLP pre-training is done separately for each environment (e.g., PushCube or PushT) but not for each task (i.e. different number of objects). For the IsaacGym environments we train a total of 2 DLP models, one for PushCube and one for PushT, each trained on 2 views of an environment with 6 and 3 objects respectively. DLP generalizes to images with a different number of objects than in training. We have added additional details about DLP pretraining in Appendix B under each environment in the new revision.

Further, if PushCube and PushT image observations were used together for pretraining DLP, can it handle the composition where both objects are present together?

If we understand the reviewer correctly, the question is whether DLP has compositional generalization capabilities. DLP is a local representation that encodes scenes as sets of entities, with each entity representing a small region of the image. During training, DLP learns to encode distinct objects from various scenes in its latent space. At inference time, it can compose these objects into new configurations, even if they were not seen together during training. Our generalization results (see Figure 6) demonstrate that DLP can handle scenes with a different number of objects than it was trained on, which requires composing objects it encountered individually during training. This is analogous to the scenario described in the question. Based on this evidence, we expect that DLP would successfully handle scenes containing both T-blocks and cubes, even if it was trained exclusively on images containing only one of these objects at a time.

Thank you to the authors for the detailed rebuttal that addressed many of my concerns.

Regarding the DLP pretraining, my question was whether training the DLP model with image data containing both cubes and T-blocks would improve the success rate—essentially, whether DLP's object-centric representations benefit from additional data. The results in Table 14 demonstrate some generalization to other shapes, which partially addresses my concern. However, for entirely novel objects like T-blocks, the performance drops, so an analysis on whether better object-centric representations could enhance these results will be useful.

My other question pertains to classifier-free guidance, not classifier guidance, which I understand requires training a separate model and increases computational demands. With classifier-free guidance, null-conditioning (e.g., using zero vectors for current state) is applied during training, and a guidance scale hyperparameter is used during inference. Typically, such guidance in image generation improves the quality while reducing the diversity. I wanted to understand whether this approach helps in this setting.