EC-Diffuser: Multi-Object Manipulation via Entity-Centric Behavior Generation
Authors: Carl Qi
Dan Haramati
Tal Daniel
Aviv Tamar
Amy Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion
Object-Centric Representation
Robotic Manipulation
Primary area: applications to robotics, autonomy, planning
TL;DR: We propose a behavioral cloning method for multi-object manipulation that combines object-centric representations with diffusion models, enabling zero-shot generalization to novel object compositions.

Abstract: Object manipulation is a common component of everyday tasks, but learning to manipulate objects from high-dimensional observations presents significant challenges. These challenges are heightened in multi-object environments due to the combinatorial complexity of the state space as well as of the desired behaviors. While recent approaches have utilized large-scale offline data to train models from pixel observations, achieving performance gains through scaling, these methods struggle with compositional generalization in unseen object configurations with constrained network and dataset sizes. To address these issues, we propose a novel behavioral cloning (BC) approach that leverages object-centric representations and an entity-centric Transformer with diffusion-based optimization, enabling efficient learning from offline image data. Our method first decomposes observations into Deep Latent Particles (DLP), which are then processed by our entity-centric Transformer that computes attention at the particle level, simultaneously predicting object dynamics and the agent's actions. Combined with the ability of diffusion models to capture multi-modal behavior distributions, this results in substantial performance improvements in multi-object tasks and, more importantly, enables compositional generalization. We present BC agents capable of zero-shot generalization to perform tasks with novel compositions of objects and goals, including larger numbers of objects than seen during training. We provide video rollouts on our webpage: https://sites.google.com/view/ec-diffuser.

Published as a conference paper at ICLR 2025
EC-DIFFUSER: MULTI-OBIECT MANIPULATION VIA ENTITY-CENTRIG BEHAVIOR GENERATION
Carl Qi', Dan Haramati Tal Daniel?, Aviv Tamar?, Amy Zhang 14 1 UT Austin, Technion, Israel Institute of T'echnology. 3 Brown University, 4 Meta Al carlq@utexas edu; dan hwramat i@brown.
ABSTRACT
Object manipulation is a common component of everyday tasks, but learning to manipulate objects from high-dimensional observations presents significant chal- lenges. These challenges are heightened in multi-object environments due to the combinatorial complexity of the state space as well of the desired behaviors. While recent approaches have utilized large-scale offline data to train models from pixel observations, achieving performance gains through scaling, these methods struggle with compositional generalization in unseen object configurations with constrained network and dataset sizes. To address these issues, we propose a novel behavioral cloning (BC) approach that leverages object-centric representations and an entity-centric Transformer with diffusion- based optimization, enabling effi- cient learning from offline image data. Our method first decomposes observations into Deep Latent Particles (DLP), which are then processed by our Transformer that computes attention the particle level, simultaneously predicting object dynamics and the agent's s actions. Combined with the ability of diffusion models capture multi-modal behavior distributions, this results substantial performance improvements in multi-object tasks and, more importantly, enables compositional generalization. We present BC agents capable of zero-shot general- ization to performing tasks with novel compositions of objects and goals, including larger numbers of objects than seen during training. We provide video rollouts on our webpage: https: //s::///siesss google. com wom/viewwect diilee//ec-i-- fuser
INTRODUCTION
Object manipulation is an integral part of our everyday lives. It requires us to reason about multiple objects simultaneously, accounting for their relationships and how they interact. Learning object manipulation is longstanding challenge, especially when learning from high-dimensional obser- vations such as images. Behavioral Cloning (BC) has shown promising results in learning complex manipulation behaviors from expert demonstrations (Chi al.
2023 Lee al., 2024). Recent works (Collaboration et al., 2023] Du et al. 2023) 2024, Zhu et al.
2024) have leveraged vast amount
of offline data paired with large models learn policies from pixel observations. Although scale has proven to be effective in some settings, it is not the most efficient way to deal with problems with combinatorial structure. For instance, despite their impressive generation results, Zhuu al.) (2024) require 2000+ GPU hours to train a model on an object manipulation task. In this work, we incorporate object-centie structure in condicondi BC from pixels produce sample-efficient and generalizing multi-objed manipulation policies.
Multi-object environments pose significant challenges for autonomous agents due to the combinatorial complexity of both the state and goal spaces as well as of the desired behaviors. Assuming n-object environment with m possible single-object goal contigurations, there are mn" total goal configurations and nz! orderings of the objects to manipulate in sequence. When learning from offline data, one cannot expect an agent to encounter all possible combinations of objects and desired tasks during training due either time/compute constraints or lack of such data. We therefore require that agent generalize to novel compositions of objects and/or tasks it has seen during training, ie. require compositional generalization (Lin et al. 2023).

---

Published as a conference paper at ICLR 2025
Reasoning about the different entities becomes increasingly complex when scaling the number of objects in the environment, especially when learning directly from unstructured pixel observations. Object-centric representation learning has shown promise in producing factorized latent representa- tions of images (Locatello et al., 2020 Daniel & Tamar 2022) and videos (Wu et al., 2023 Daniel & Tamar 2024) which can be leveraged for learning visual control tasks that involve several objects and possibly facilitate compositionally generalizing behaviors. Prior works have made progress in using these representations for compositional generalization in control (Chang al. 2023] Zadaianchuk et al. 2021: Haramati et al. 2024). While Haramati et al. (2024) relies on their Transformer based policy to generalize to unseen configurations in an online reinforcement learning setting, we demon- strate that naively taking this approach in the BC setting is insufficient, as the policy fails to capture the multi-modality behaviors as the number of manipulated object increases. This calls for a method that can better leverage object-centric representations when learning from limited offline demonstrations.
We propose a novel diffuusion-baatt BC method that leverages Transformer-base diffusion model with an unsupervised object-centric representation named Deep Latent Particles (DLP) (Daniel & Tamar 2022). We first factorize images into sets of latent entities, referred to particles. We then train a entity-centric Transformer model with diffusion to generate conditoond sequences of particle-states and actions, and use it for Model Predictive Control (MPC). These help deal with the two main challenges in our setting: (1) Multi-modal Behavior Distributions - Diffusion models' ability handle multi-modal distributions aids in capturing the combinatorial nature of multi-object manipulation demonstrations; (2) Combinatorial State Space Our Transformer- based architecture computes attention on the particle level, thus facilitating object-level reasoning which gracefully scales to increasing number of objects, and more importantly, unlocks compositional generalization capabilities. Our Entity-Centree Diffuser (EC-Diffuser) significantly outperforms baselines in manip- ulation tasks involving more than a single object and exhibits zero-shot generalization to entirely new compositions of objects and goals containing more objects than in the data it was trained on.
2 RELATED WORK
Multi-object Manipulation from Pixels: Previous work using single-vector representations of
image observations for control (Levine et al. 2016 Nair et al. 2018 Hafner et al. 2023: Lee et al.)
2024)
fall short compared to methods that leverage object-centric representations (Zadaianchuk et al.
2021
Yoon et al. 2023 Haramati et al., 2024 Ferraro al. 2023 Zhu et al. 2022; Shi al.
2024) in mujeit-bjecct environments. Several works have studied compositional generalization in this
setting. Collaboration et al. (2023); Du et al. 2023 2024) Zhu al. (2024) learn directly from pixel observations with large-scale networks, data and compute and rely on scale for possible generalization and transfer. Other works have proposed approaches that account for the underlying combinatorial structure in object manipulation environments to achieve more systematic compositional generaliza- tion (Zhao et al. 2022 Chang et al. 2023 Haramati et al. 2024), requiring significantly less data and compute. We continue this line of work, dealing with the setting of learning from demonstrations and the various distinct challenges introduces.
Diffusion Models for Decision-Making Diffusion models have been used for decision making in
many recent works (Chi et al. 2023 Janner et al. 2022 Ajay et al. 2023 Reuss et al., 2023 Du
et al. 2024 Zhu et a 2024), thanks to their abilities to handle multi moolti-modaling modality and their robustness
when scaled to larger datasets and tasks. Diffuusion Policies (Chi et al., 2023 Reuss et al. 2023)
use diffusion over the actions conditioned on the observations. Diffuser- based approaches (Janner
jet al. 2021; Zhu et al. (004)) diffuse over both the observations and actions and execute the actions at test time. Finally, Ajay al. (2023) and and et al. (2024) train diffusion over the states and train an inverse model to extract the actions. Compared to these works, we build on top of Diffuser (Janner
et al. 2022) and generate both object-centric factorized states and actions simultaneously.
3 BACKGROUND
In this work, we propose method that leverages Transformer-based diffusion model and object- centric representations for learning policies from demonstrations. In the following, we give a brief overview of the different components.
2

---

Published as a conference paper at ICLR 2025
Nolse Prediction é
Linear
62
Pointwise
Lx
EC-Diffuser Block
Feedforward
Scale, Shift
E
Notsy Latents
Noley Latents
E
COND
Layer
LayerNorm [AdaLN)
000
OOO
Scale
Multi Head Actention
D Eneoder
DLP Eneader
a Eneoder
MLP
Embed
Ye.Pi
Scall, Shift
Layer Norm
MLP
V x Camera
Views
First Observation
Observations fia-1
Goal
Trajectory
Reat-t
Timestep
Diffusion
Input
Conditiooning
Figure 1: EC-Diffuser architecture. Our model learms a conditional latent denoiser that generates sequences of
actions and latent states from trajectory of H image observations across V views 9. - and actions 01.H-1-
10:1 -0.V-1
-0:--1 First, pre-trained DLPv2 encoder transforms each image into set of M latent particles 20:11-1 (where denotes all M particles for brevity). Projected actions then added to the latent set additional particles, forward-diffusion process adds noise the actions and latent states 2, excluding the first and last (goal) latent states, 40 _ and zz - respectively. A Transformer-basee conditional denoiser with L blocks predicts the added noise. Each block employs Adaptive LayerNorm (AdaLN) to modulate intermediate variables using a projection of the conditioning variable & @ cat(t, a-]. which is concatenation of the diffusion timestep t and the action a,,.
Coal-Conditioned Behavioral Cloning: Learning-basee decision-making algorithms that aim to fit distribution of offline demonstration data is commonly referred to Behavioral Cloning (BC). Given dataset of trajectories D = {oj ( containing environment observations o E O and corresponding actions AA the goal of BC is to learn the conditional distribution IP (o.lo:) typically referred to policy O + AA and parameterized by neural network. In Goal-Conditioned (GC) BC, the demonstration trajectories are augmented with goal E g that indicates the task the demonstration was provided for and a goal-conditioned policy æ G The goal can be in the form of e.g, natural language or the last observation in the trajectory g o
learned respectively.
Diffusion: In Denoising Diffusion Probabilistic Models (DDPM (Ho al. ). given data point Xo sampled from real data distribution g(xn). forward diffusion process gradually adds Gaussian noise over timesteps:
g(x:|xx-1) = (xx;; / - 3,1),
where B is noise schedule. DDPM learns a reverse process to denoise the data:
(x)) = 1i t), Mo (x,, t)).
Here, [x(x,, t) represents the learned covariance matrix of the reverse process. In practice, it is often set to a fixed multiple of the identity matrix to simplify the model, i.e., Eg (xx, o²1, where of time-dependent scalar. The model is trained to estimate the mean μa by minimizing a variational lower bound, which is simplified to a loss L of predicting the noise added during the forward process:
E= [[le- fiee-ephe cx(x,, t)]
where Eg is the noise prediction network. To incorporate conditional information, the process be extended to a conditional Diffusion model by conditioning on a variable c, modifying the reverse process px(x;-1/x,, c) and the noise prediction network en Ex(,,,,,
Deep Latent Particles (DLP): DLP is an unsupervised, object-centric image representation method proposed by Daniel & Taman 2022 and later enhanced DLP Daniel & Tamar (2024). DLP
Throughout this work, we DI.v22 but refer simply DI.P for brevity.
3

---

cetimnutttio-qqqivarttt policy network architecture, which we describe in Published as a conference paper at ICLR 2025
M-1 and latent background decomposes image into of M latent foreground partieles 1i=0
particle Sbg. and is trained in an unsupervised manner as a variational autoencoder (VAE (Kingma & Welling 2014)), where the objective is reconstructing the original input image from the latent particles. Each foreground particle consists of multiple attributes 2' [Zps Ra Zd, 21, 2y]² R (6+)); 2p E R represents the position in 2D coordinates (i.e., keypoint); 2, E IR² denotes the scale, specifying the height and width of a bounding box around the particle; 2, E approximates local depth when particles are close together; E R is transparency parameter in the range [0, 1]; and ef E IR encodes visual features from glimpse around the particle, where n is hyper-parameter determining the latent dimension of the visual features. The background is encoded as 2og e R with Tbgg serving as the latent dimension hyper-parameter for the background encoding. We provide an extended background of DLP in Appendix
METHOD
Our goal is to learn a goal-conditioned offline policy from image inputs for tasks involving multiple objects. A common approach is to develop model that can generate sequence of goal-conditioned states and actions, which can then be used for control. The challenge lies effectively learning such model from high-dimensional pixel observations. To achieve this, we propose Entity-Centric Diffuser (EC-Diffuser), a diffusion-bass policy that leverages DLP, an unsupervised object-centric representation for images, and utilizes Transformer-basee architecture to denoise future states and actions. As DLP a single-image representation, encoding each image as an unordered set of particles, it lacks correlation between particles across timesteps and views. This property requires non-trivial design of architecture and optimization objectives. In the following sections, we detail each component of the method. We start by describing the process of encoding images to particles with DLP in Section 4.1 then outline the architecture of our entity-centric Transformer in Section 4.2 Finally, we explain how EC-Diffuser can be applied in goal-conditioned behavior-cloning settings in Section
4.1 OBJECT-CENTRIC REPRESENTATION WITH DLP
We first extract a compact, object-centric representation from pixel observations. Given image observations of state o from V different viewpoints, (4..... ... -V-1). we encode each image with DLPx2 (Daniel & Tamar 2024 encoder ΦDLP. The resulting representation, as described in Section 3
is a set of M latent particles for each view e, denoted by Z = {2 1i-0 where Z = ΦDLP is important to note that there no correspondence between particles from different views (i.e. = and 2 i can represent different objects), nor is there correspondence between particles from n different states (i.e. z and Zg"/ can represent different objects). These properties of the DLP representation require the following section.
Pre-training the DLP: Following Haramati al. (024). first collect a dataset of image ob- servations to train the object-centric partiele representation with DLP. To acquire image data from environments, we employ either a random policy (Haramati al. 2024) or utilize the demonstration data (Lee et al. 2024). For the BC training stage, we freeze the pre-trained DLP encoder and use it solely for extracting representations. Utilizing these low-dimensional particles instead of high- dimensional images significantly reduces memory strain during the BC stage. We provide detailed information about training the DLP models in Appendix C.3
4.2 CEIITY-EENTRIC TRANSFORMER
Equipped with DLP's object-centric representations, the particles, we aim to construct conditional generative model. In this work, we adopt diffusion- approach (DDPM (Ho et al. 2020) to learn the denoising of particles and continuous actions. However, the unique strueture particles an unordered set necessssssssssss permutation-equivariiii denoiser. To address this, we design an entity- centric Transformer- based denoiser architecture, termed EC-Diffuser, which processes a sequence of
We note that the true state of the environment is not observed, and the learned DL.P representation is trained purely from pixels.
4

---

Published as a conference paper at ICLR 2025
PushCube
PushT
FrankaKitchen
Figure 2: Visualization of the simulated oblti-object manipulation environments used in this work.
Image
DLP
Foreground
Background
Figure 3: Visualization of DLP decomposition the PushCube task. From left to right: the original image, DLP position keypoints overlaid the original image. reconstructed image from foreground particles, and reconstructed image from the background particle.
observations encoded as particles with DLP, along with corresponding actions. The architecture is illustrated in Figure
where Z denotes the state particles, a, represents the action as separate token, and is the generation horizon. additional particle to the set, as also Following Haramati et al. (2024), we add the action generate actions. The EC- Diffuser processes these noised particles and predicts the noise added during the diffusion forward process. Notably, omit positional embeddings for individual particles, instead incorporating positional information only differentiate particles from different timesteps and views. Furthermore, the state particles conditioned actions and diffusion timesteps via adaptive layer normalization (AdaLN), which has proven beneficial Transformer based diffusion models (Peebles & Xie 2023). Given action a,, diffusion timestep and intermediate variable AdaLN performs the following modulation in each Transformer block:
Formally, we construct the denoiser inputs as {ZZ,,
Q1, 02 B1, B2,
MLP(cat(t, a,l)),
+ Self Self-Attention(F1 LL((2) + $1),
z 02 ML.P(72 IN((2) + (2).
4.3 ENTITY-CENTRIC DIFFUSER FOR GOAL-CONDITIONED BEHAVIORAL CLONING
We adapt EC-Diffuser to GCBC tasks by incorporating conditioning variables into the diffu- sion process. Formally, the diffusion process operates over future states and actions: X, = with the current state and goal serving conditional variables: Cy = [(2 (") /=))==.,...v--- We define the goal as the last timestep in the demonstration trajectory, i.e. 2 2/, where T is the trajectory length. To train EC-Diffuser, we normalize all input features (DLP's features and actions) to |-1, 1] and employ an l loss both states and actions. One might question the effectiveness using // or 4 losses directly on unordered set inputs. Typically, generating unordered sets like point clouds calls for set-based metrics such as Chamfer distance to compare set similarity. However, our case, the objective is particle- wise denoising: noise is added independently to each particle, and the denoising process neither imposes requires any specific ordering of set elements. Furthermore, we leverage the Transformer's permutation catiiiionnuvvviiiannn.. struc- ture by omitting positional embeddings within the of particles. These factors enable the simple I
{(Z), (Zz,a,)) (-),,a,))=..........
5

---

Published as a conference paper at ICLR 2025
loss function to work effectively with diffusion, aligning with previous works that applied diffusion to point clouds (Vahdat et al. 2022] Melas Kyriazii et al. (023).
[l|e co (x,, t, 92
For control purposes, we execute the first action produced by the model in the environment, i.e. TT(x,, 4, C2) = ao, and perform MPCCstyle control by querying the model at every timestep. In practice, we do not directly use the generated latent states for control. However, we empirically found that denoising these latent states is critical, as we discuss later. Notably, the generated latent states serve valuable purpose for visualization: they can be decoded using the pre-trained DLP decoder to reconstruct images, effectively visualizing the imagined trajectory.
EXPERIMENTS
The experiments in this work are designed to answer the following questions: (I) How do object- centric approaches compare to unstructured baselines learning tasks with combinatorial structure? (II) Does object-centric structure facilitate compositionally generalizing behavioral cloning agents? (III) What aspects contribute to performance and compositional generalization?
To study the above, we evaluate our method on 7 goal-conditioned multi-object manipulation tasks across 3 simulated environments and compare against several competitive BC baselines learning from various image representations.
Environments A visualization of the environments used in this work presented in Figure 2 and visualization of the DLP decomposition for PushCube is shown in Figure 3 PushCube and PushT are both IsaacGym- (Makoviychuk et al. 2021) tabletop manipulation environments introduced in Haramati al. 2024), where Franka Panda arm pushes objects in different colors to a goal configuration specified by an image. In PushCube the objects are cubes and the goals are positions, while in PushT the objects are T-shaped blocks and the goals are orientations. In F FrankaXitcher set of 4 out of 7 possible tasks kitchen environment. We use the goal-conditioned image-based based variant from Lee et al. (2024). Detailed descriptions of these environments as well as the datasets used for training can be found in Appendix These tasks all possess a compositional nature, requiring the agent to manipulate multiple objects to achieve desired goal configuration.
initially introduced in Gupta et al. (202)), the agent is required to complete
Baselines We compare EC-Diffuser's 's performance with the following BC methods: (1) VQ-BeT (Lee) et 2024): a SOTA, non-diffusiinn-bsed method utilizing Transformer architecture. In our experiments, we find that the ResNet18 (He et al. 2016) used in VQ-BeT fails to achieve good performance in PushCube and PushT. Consequently, we use VQ-VAE (Van Den Oord et al. 2017) image representation pre-trained on environment images. (2) Diffuser (Janner et al. 2022):
the original Diffuser trained without guidance and takes flattened VQ-VAE image representations as input. (3) EIT+BC- - a direct adaptation of the HIT policy from Haramati et al. (2024) to the BC setting, learns from DLP image representations. (4) EC Diffusion Policy: inspired by Chi al.) (2023) and modified for the goal-conditioned setting, learns from DLP image representations. Further descriptions and implementation details of each baseline can be found in Appendix C.2
For PushCube and PushT, all results are computed as the mean of 96 randomly initialized con- figurations. In evaluating FrankaKi waaaiichhe.. chen. we adopt the protocol used by VQ-BeT (Lee et al. 2024). sampling 100 goals from the dataset. Standard deviations are calculated across 5 seeds. We provide extended implementation and training details, and report the set of hyper- parameeetsss used in our experiments in Appendix C
5.1 LEARNING FROM DEMONSTRATIONS
In this section we aim to answer the first question comparing object-centric approaches to un- structured baselines in learning tasks with combinatorial structure. Performance metrics for all tasks are reported in Table I In PushCube and PushT, EC-Diffuser significantly outperforms all baselines, with the performance gap widening the number of objects in the environments increases. Notably, it is the only method that achieves better-than mhan-r performance on P Push3T, T, the most challenging task in our suite.

---

Published as a conference paper at ICLR 2025
Env (Metric)
Ohj
VQ-BeT
[Q-BeT
Diffuser
EIT+BC (DLP)
EC Diffusion Policy (DLP)
ECC Difliser (DLP)
0.929
0.032
0.027
0.887
0.031
0.948
0.015
PushCube
2
0.052 = 0.010
0.013 2 0.011
0.146 = 0.125
0.388 1: 0.106
0.917 0.030
(Success Rate f)
0.006
0.001
0.004
0.164
0.668 0.169
0.894 0.025
1.227
0.066
1.522 0.19
0,835
0.081
0.493 + 0.068
0.263 0.022
FushT
1.520
0.056
1,465
0.034
1.214 0.147
0.452
(Avg. Radion Diff.
1,54] 0.045
1.542 + 0.045
1.526
0.047
1,538 + 0.040
0.805 0.256
FrankoKizchen
2.384 1 0.123
0.846 1. 0.101 2,360 1. 0.088
3.046 1. 0.156
3,031 0.087
(Goals Reached 1)
Table 1: Quantitative performance for different methods in the PushCube, Push'I, and Frankak K t en environments for varying number of objects. Methods are trained for 1000 epochs, and the hest performing checkpoints are reported. The best values are in bold. *We obtain a lower score than the one reported in the VQ-BeT paper (2.60) using their released codebase, which does not support fine-tuning the ResNet backbone.
Baselines utilizing unstructured representations, such VQ-BeT and Diffuser, fail entirely when presented with multiple objects. While EIT+BC demonstrates improved performance over these baselines, it struggles behaviors present in multi-object manipulation demonstrations, which the deterministic EIT+ fails to capture. EC Diffusion Policy emerges as the best-performing baseline, differing from our method primarily in that it does not generate states alongside actions. We posit that generating future particle-states serves two purposes: (1) implicitly planning for future object configurations, and (2) acting as an auxiliary abjective-similar in spirit to methods such as those proposed by Yarats et al. (2021) ensuring the model's internal representation aware of all objects and their attributes.
handle more than a single object. This can be attributed the diverse
Our method also outperforms all baselines, including the SOTA method VQ-BeT, on FrankaKi tchen. In experiments with alternative object-centric representations (Appendix (D.2). we found that EC-Diffuser, when paired with Slot Attention (Locatello et al., 2020) achieves state-of- the-art performance with an average of 3.340 goals reached. Intuitively, object-centric representations provide more useful and structured bias compared to learning from unstructured image representa- tions for tasks such as kitchen object manipulation.
To gain further insight into this performance advantage, provide DLP decompositions of envi- ronment images in Figure 3 and Appendix [D.5] These visualizations reveal that particles capture variable aspects of the environment across images-such as the robot, kettle, burner, and parts of hinged while the latent background particle represents the rest. In the context of learning from demonstration images, these captured elements correspond to the controllable parts of the environment, which are of primary interest to our agent, thus facilitating better policy learning. A similar notion of disentanglement has proven beneficial in the RL setting (Gmelin et al., 2023). Finally, EC Diffusion Policy also achieves SOTA performance on FrankaKit chen, suggesting that this task may not be as challenging as our multi-object tasks.
5.2 COMPOSITIONAL GENERALIZATION
The results in Section 5.1 clearly demonstrate that object-centric approaches outperform unstructured methods in multi-object environments when learning from images. As EC-Diffuser is the only method achieving strong performance in manipulating 3 objects, we focus our answer to the second question on whether our method can generalize zero-shot to unseen compositions of objects and/or goals. To address this, consider two generalization settings: (1) PushCube We train our method with 3 objects, their colors randomly chosen at the beginning of each episode from 6 options. We then test it on environments with up objects. visualization for the task is shown Figure 4 (2) PushT We train an agent with 3 objects and 3 goals, and then test on scenarios with up to 4 objects and varying goal compositions. A visualization of this task is shown in Figure 5 For comparison, present results of the best pert-performin baseline EC Diffusion Policy in these generalization settings. Further details on training these models are provided in Appendix C.I
We report the quantitative results for PushCube generalization in Table 2 Our method significantly outperforms the baseline across all configurations of PushCube generalization, maintaining a high success fraction the number of objects increases. As illustrated in Figure 4b), which shows the distribution of per-object goal-reaching success, our agent successfully manipulates up to objects to their goals despite being trained on data containing only 3 objects.
7

---

Published as conference paper at ICLR 2025
Train
Test
Namber
(a) Task Visualization
(b) BC-Diffuser
(c) EC Diffusion Policy
Figure 4: PushCube Generalization Agents trained 3 objects and evaluated on 1 6 objects. See Tahle for quantitative results. Visualization of Push generalization task. (b), (e) Histograms of per-object goal-reaching success (n out 6) in PushCube generalization task with objects.
Number of Objects
2
3 (training)
6
VQ-BeT
0.111 = 0.006
0.081 + 0.002
0.137 = 0.011
0.097 + 0.014
0.084 + 0.008
0.090 +\t 0.015
Diffuser
0.170 = 0.036
0.078
0.018
0.080 = 0.012
0.045
0.015
0.047
0.004
0.030 + 0.011
EIT+BC
0.111 0.006
0.078
0.018
0.084 = 0.012
0.050 +: 0.014
0.048
0.012
0.043 = 0.006
EC Diffusion Policy
0.903 0.030
0.501
0.024
0.385
0.067
0.207
0.024
0.158
0.004
0.122 + 0.010
EC- SC-Diffuser (ours)
0.993 0.006
0,981 0.003
0.886
0.051
0.858 0.002
0.767 0.032
0.711 = 0.070
Table 2: PushCube generalization results. The success fraction (number of successful objects total number of objects) is reported for different numbers of cubes. Agents are trained on objects with colors randomly selected from 6 options. Higher values correspond to better performance The best values are in bold.
We additionally report the quantitative results PushT generalization in Table These results demon- strate only slight performance drop (averaging between 0.1 to 0.26 radians) when manipulating objects with novel goal compositions for up to 4 objects. contrast, EC Diffusion Policy fails to achieve setter-than-randon performance these PushT generalization tasks. Our approach clearly demonstrates zero-shot generalization capabilities in multi-object manipulation tasks. We showcase the rollouts of model the most challenging generalization tasks in Figure additional results on generalization to new objects in Appendix D.4
We provide
5.3 ABLATION STUDIES
In this section we aim to answer the third question what contributes to the performance of our model? We ablate our key design choices on the PushCube environment. The results, presented in Table report the success rate and success fraction for each task. First, we compare our model to version that uses VQ-VAE representations as input, treating them as a single particle without any object-centric structure. This unstructured representation results in significant drop in performance as the number of objects increases, highlighting the importance the object-centric representation (DLP) for multi-object tasks. Next, evaluate model with similar architecture, but trained without diffusion, where DLP and actions are generated in an auto-regressiv manner. For training, we replace the / distance with the -Chhafeerr distance for the latent states particles, and use the distance for the actions. This approach fails to learn even in tasks involving a single object, highlighting the critical role of the diffusion process in effectively co- ro-generating DLP and actions. Finally, we compare to variant of our model that does not generate latent states (as DLP representations). As the number of objects increases, the performance of this ablation rapidly deteriorates, demonstrating that the joint generation of particle states and actions is essential for reasoning about both objects and actions. We provide additional results with alternative object-centric representations in Appendix [D.2]
5.4 DISCUSSION ON STATE GENERATION
To further explore the generalization ability of our model, we present the generated particle states by decoding them into images using the pre-trained DLP decoder. As shown in Figure 7 model can produce high-quality future states that were not present the training data.
Additionally, analyze how well the model maintains the temporal consistency of particles across frames. To do this, we first label the particles according to the object they represent at each timestep in trajectory. We then perform T-test comparing the attention values between particles representing
8

---

Published as conference paper at ICLR 2025
Count
Train
Test
(teciars)
Disance (Rodians]
(a) Task Visualization
(b) EC-Diffuser
(c) EC Diffusion Policy
Figure 5: PushT Generalization agents trained on 3 objects with 3 goals and tested in different combina- tions of objects and goals. See Table 3 for quantitative results. (a) Visualization of generalization task from objects. (b), (c) Histograms orientation difference (in radians) in the generalization task depicted (a).
Objects a of Goals
3/1
3/2
3/3 (training)
4/1
4/2
VQ-BeT
1.556 0.082
0.053
1.515
0.036
1.572 = 0.086
0.005
1.522
Diffuser
1.597 0.080
1.541
0.024
1.545 0.040
1.632
0.068
1.541
0.063
1.527
0.038
EIT+BC
1.515 0.136
1.507
0.071
1.501
0.099
1.554 0.038
1.594
0.038
1.524 0.020
BC Diffusion Policy
1.567 +: 0.027
1.548 0.117
1.499 0.013
1.552 0.071
1.547 0.040
1.557 0.020
EC-Diffuser (ours)
0.781 +: 0.160 0.778 0.115
0.817
0.188
0,945 0,050
0.923 0.119
0.948 0.101
randomly initialized configurations. The average orientation difference (radians) is reported various object and goal compositions. The number of colors matches the number of goal objects. When "# of Goals" less than "# of Objects," multiple objects share the same color and target orientation. Lower values correspond to better performance. The best values are in bold.
Table PushT generalization results, averaged
Method
Cube
Cubes
Cubes
Unstructured Rep.
0.829 + 0.041 /0.82 +: 0.041
0.089 = 0.037 / /0.221 + 0.037
0.008
0.004 0.024
Without Diffusion
0.052 0.024 0.052
0.024
0.004
0.005 /0.04 0.010
0.000
0.000/ 0.051 0.010
Without State Generation
0.941 + 0.020 / 0.941 + 0.020
0.423 : 0.092 /9 + 0.095
0.529 + 0.311 /0,715 + 0.224
EC-Diffuser (ours)
0.948 +: 0.015 / 0.948 +: 0.015
0.917 = 0.030 / 0.948 + 0.023
0.894 + 0.025 0,950 0.016
Table 4:
Quantitative performance for different ablations FushCube tasks. Results are reported as the mean success rate and success fraction (number of object successes total number of objects) for each task.
the same object over time against the attention values between two random particles. The results show that particles representing the same object over time have significantly higher attention values, with T-test yielding p-value of 8.367c-6 < 0.05. This demonstrates that the model implicitly matches objects and enforces object consistency over time, aiding in predicting multi-object dynamics.
CONCLUSION
In this work, introduced Entity-Centric Diffuser, novel liffusion-based behavioral cloning method for multi-object manipulation tasks. By leveraging unsupervised object-centriy represen- tations and Transformer-baseee architecture, EC-Diffuser effectively addresses the challenges of multi-modal behavior distributions and combinatorial state spaces inherent in multi-object environ- ments. We demonstrated significant improvements over existing baselines in manipulation tasks involving multiple objects, and zero-shot generalization to new object compositions and goals, even when faced with more objects than encountered during training. These results highlight the potential of combining object-centric representations with diffusion models for learning complex, generalizing manipulation policies from limited offline demonstrations.
Limitations and Future Work: The performance of EC-Diffuser relies on two core foundations: the quality of the demonstration data and of the object-centric representation. this work, we utilized DLP, which worked well in our environments While DLP provides excellent object-centric decomposition scenes with explicit objects, such as the PushCube and PushT environments, it captures slightly different notions of entities in the trankaKitcher environment. Regarding real-world environments, we believe our general approach as well as our proposed algorithm are applicable real-world object manipulation, and do not fundamental limitation in solving tasks similar to the ones in our simulated suite using our method. We provide preliminary EC-Diffuser results the real-world Language be dataset (Lynch al. 2023) in Appendix D33 That

---

Published as a conference paper at ICLR 2025
First Observation
Goal
PushCube
PushT
FrankaKitchen
Figure 6: Visualization of EC-Diffuser rollouts from cach environment. The final observation (highlighted with green border) demonstrates that our agent successfully completes all tasks.
Observation
Generated Frames
Goal
T
View 1
View
View
View 2
Figure 7: Visualization of generated DLP The DLP states produced by EC-Diffuser decoded into images using the pre-trained DLP decoder. (Top) Generated frames PushCube generalization task: model is trained with 3 cubes and can generate trajectories involving cubes. (Bottom) Generated frames for the PushT generalization task: the model is trained with 3 T-blocks and can generate sequences containing 4. The objects of interest and their respective goals in the generated frames highlighted with orange cirele.
said, acquiring unsupervised object-centric representation can be more challenging in real-world scenes due to higher visual complexity, especially in "in-the-wild" environments. The problem of unsupervised object-centric factorization of natural images is far from being solved, and acquiring such representations is an active field of research (Seitzer et al., 2023 Zadaianchuk et al. 2024). Another limitation for real-world application is the long inference time due the iterative nature of diffusion models. This could be improved by adopting recent approaches focused on reducing the number iterations required produce high-quality samples (Song et al. 20200) Karras al. 2022 Song et al,, (023). Several interesting directions emerge from this work; future research could explore the incorporation of guided diffusion for offline RL, or the application of our approach to planning and world-modeling as well as to new domains such as autonomous driving.
10

---

Published as a conference paper at ICLR 2025
7 REPRODUCIBILITY & ETHICS STATEMENT
We provide extended implementation and training details, and report the set of used in our experiments in Appendix C We provide the project code reproduce the experiments at bttpp:///ithub.co http:////cccoaa///-----ee-e
This research was conducted in simulated environments, with no involvement of human subjects or privacy concerns. The datasets used are publicly available, and the work aims to improve the efficiency of robotic object manipulation tasks. We have adhered to ethical research practices and legal standards, and there are no conflitts of interest external sponsorship influencing this work.
8 ACKNOWLEDGMENTS
CQ and AZ are supported by NSF 2340651, NSF 2402650, DARPA HR001 12490431, and ARO W91 1-9-244-4--01999. This research was partly funded by the European Union (ERC, Bayes-RL, 101041250). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency (ERCEA). Neither the European Union nor the granting authority can be held responsible for them.
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023.
Michael Chang, Alyssa Li Dayan, Franziska Meier, Thomas I.. Griffiths, Sergey Levine, and Amy Zhang. Hierarchical abstraction for combinatorial generalization object rearrangement. In The Eleventh International Conference on Learning Representations, 2023. URL https:/7
openreview net/ forum?id f
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Open X-Fmbodiment Collaboration, Abby O'Neill, Abdul Rehman, Abhiram Maddukuri, Abhishek Gupta, Abhishek Padalkar, Abraham Lee, Acorn Pooley, Agrim Gupta, Ajay Mandlekar, Ajinkya Jain, Albert Tung, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anchit Gupta, Andrew Wang, Anikait Singh, Animesh Garg, Aniruddha Kembhavi, Annie Xie, Anthony Brohan, Antonin Rafin,, Archit Sharma, Arefeh Yavary, Arhan Jain, Ashwin Balakrishna, Ayzaan Wahid, Ben Burgess-Limerick, Beomjoon Kim, Bernhard Schölkopf, Blake Wulfe, Brian Ichter, Cewu Lu, Charles Xu, Charlotte Le, Chelsea Finn, Chen Wang, Chenfeng Xu, Cheng Chi, Chen- guang Huang, Christine Chan, Christopher Agia, Chuer Pan, Chuyuan Fu, Coline Devin, Danfei Xu, Daniel Morton, Danny Driess, Daphne Chen, Deepak Pathak, Dhruv Shah, Dicter Büchler, Dinesh Jayaraman, Dmitry Kalashnikov, Dorsa Sadigh, Edward Johns, Ethan Foster, Fangchen Liu, Federico Ceola, Fei Xia, Feiyu Zhao, Freek Stulp, Gaoyue Zhou, Gaurav S. Sukhatme, Gautam Salhotra, Ge Yan, Gilbert Feng, Giulio Schiavi, Glen Berseth, Gregory Kahn, Guanzhi Wang. Hao Su, Hao-Shu Fang, Haochen Shi, Henghui Bao, Heni Ben Amor, Henrik I Christensen, Hiroki Furuta, Homer Walke, Hongjie Fang. Huy Ha, Igor Mordatch, Ilija Radosavovic, Isabel Leal, Jacky Liang. Jad Abou-Chakra, Jachyung Kim, Jaimyn Drake, Jan Peters, Jan Schneider, Jasmine Hsu, Jeannette Bohg, Jeffrey Bingham, Jeffrey Wu, Jensen Gao, Jiaheng Hu. Jiajun Wu, Jialin Wu, Jiankai Sun, Jianlan Luo, Jiayuan Gu, Jie Tan, Jihoon Oh, Jimmy Wu, Jingpei Lu, Jingyun Yang, Jitendra Malik, João Silvérío, Joey Hejna, Jonathan Booher, Jonathan Tompson, Jonathan Yang, Jordi Salvador, Joseph J. Lim, Junhyek Han, Kaiyuan Wang, Kanishka Rao, Karl Pertsch, Karol Hausman, Keegan Go, Keerthana Gopalakrishnan, Ken Goldberg, Kendra Byrne, Kenneth Oslund, Kento Kawaharazuka, Kevin Black, Kevin Lin, Kevin Zhang, Kiana Ehsani, Kiran Lekkala, Kirsty Ellis, Krishan Rana, Krishnan Srinivasan, Kuan Fang, Kunal Pratap Singh, Kuo-Hao Zeng, Kyle Hatch, Kyle Hsu, Laurent Itti. Lawrence Yunliang Chen, Lerrel Pinto, Li Fei-Fei, Liam Tan, Linxi "Jim" Fan, Lionel Ott, Lisa Lee, Luca Weihs, Magnum Chen, Marion Lepert, Marius Memmel, Masayoshi Tomizuka, Masha Itkina, Mateo Guaman Castro, Max Spero, Maximilian

---

Published as a conference paper at ICLR 2025
Du, Michael Ahn, Michael C. Yip, Mingtong Zhang, Mingyu Ding, Minho Heo, Mohan Kumar Srirama, Mohit Sharma, Moo Jin Kim, Naoaki Kanazawa, Nicklas Hansen, Nicolas Heess, Nikhil J Joshi, Niko Suenderhauf, Ning Liu, Norman Di Palo, Nur Muhammad Mahi Shafiullah, Oier Mees, Oliver Kroemer, Osbert Bastani, Pannag R Sanketi, Patrick "Tree" Miller, Patrick Yin, Paul Wohlhart, Peng Xu, Peter David Fagan, Peter Mitrano, Pierre Sermanet, Pieter Abbeel, Priya Sundaresan, Qiuyu Chen, Quan Vuong, Rafael Rafailov, Ran Tian, Ria Doshi, Roberto Mart 'in- Mart'in, Rohan Baijal, Rosario Scalise, Rose Hendrix, Roy Lin, Runjia Qian, Ruohan Zhang, Russell Mendonca, Rutav Shah, Ryan Hoque, Ryan Julian, Samuel Bustamante, Sean Kirmani, Sergey Levine, Shan Lin, Sherry Moore, Shikhar Bahl, Shivin Dass, Shubham Sonawani, Shuran Song, Sichun Xu, Siddhant Haldar, Siddharth Karamcheti, Simeon Adebola, Simon Guist, Soroush Nasiriany, Stefan Schaal, Stefan Welker, Stephen Tian, Subramanian Ramamoorthy, Sudeep Dasari, Suneel Belkhale, Sungjae Park, Suraj Nair, Suvir Mirchandani, Takayuki Osa, Tanmay Gupta, Tatsuya Harada. Tatsuya Matsushima, Ted Xiao, Thomas Kollar, Tianhe Yu, Tianli Ding, Todor Davehev, Tony Z. Zhao, Travis Armstrong, Trevor Darrell, Trinity Chung. Vidhi Jain, Vincent Vanhoucke, Wei Zhan, Wenxuan Zhou, Wolfram Burgard, Xi Chen, Xiaolong Wang, Xinghao Zhu, Xinyang Geng, Xiyuan Liu, Xu Liangwei, Xuanlin Li, Yao Lu, Yecheng Jason Ma, Yejin Kim, Yevgen Chebotar, Yifan Zhou, Yifeng Zhu, Yilin Wu, Ying Xu, Yixuan Wang, Yonatan Bisk, Yoonyoung Cho, Youngwoon Lee, Yuchen Cui, Yue Cao, Yueh-Hua Wu, Yujin Tang, Yuke Zhu, Yunchu Zhang, Yunfan Jiang, Yumshuang Li, Yunzhu Li, Yusuke Iwasawa, Yutaka Matsuo, Zehan Ma, Zhuo Xu, Zichen Jeff Cui, Zichen Zhang, and Zipeng Lin. Open X- X-Embodiment: Robotic learning datasets and RT-X models. https: arxiv. org abs/ 2310 08864, 2023.
Tal Daniel and Aviv Tamar. Unsupervised image representation learning with deep latent particles. In International Conference on Machine Learning, pp. 4644 4665. PMLR, 2022.
Tal Daniel and Aviv Tamar. Ddlp: Unsupervised object-centric video prediction with deep dynamic latent particles. Transactions Machine Learning Research, 2024.
Yilun Du, Sherry Yang. Pete Florence, Fei Xia, Ayzaan Wahid, Pierre Sermanet, Tianhe Yu, Pieter Abbeel, Joshua B Tenenbaum, Leslie Pack Kaclbling, et al. Video language planning. In The Twelfth International Conference on Learning Representations, 2023.
Yilun Du Sherry Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Josh Tenenbaum, Dale Schuurmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. Advances in Neural Information Processing Systems, 36, 2024.
Stefano Ferraro, Pietro Mazzaglia, Tim Verbelen, and Bart Dhoedt. Focus: Object-centric world models for roboties manipulation. In RSS2023, Robotics: Science and Systems, pp. 1-11, 2023.
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pp. 512-519, IEEE, 2016.
Kevin Gmelin, Shikhar Bahl, Russell Mendonca, and Deepak Pathak. Efficient rl via disentangled environment and agent representations. In International Conference on Machine Learning, pp. 11525-11545.. PMLR, 2023.
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. In Conference on Robot Learning, pp. 1025-1037. PMLR, 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.0410 2023.
Dan Haramati, Tal Daniel, and Aviv Tamar. Entity-centric reinforcement learning for object ma- nipulation from pixels. In The Twelfth International Conference Learning Representations, 2024.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 770-778, 2016.
12

---

Published as a conference paper at ICLR 2025
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-660 6851, 2020.
Max Jaderberg, Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. Advances in neural information processing systems, 28:2017-2025, 2015.
Tomas Jakab, Ankush Gupta, Hakan Bilen, and Andrea Vedaldi. Unsupervised learning of object landmarks through conditional image generation. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 4020-400 4031, 2018.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In Advances in Neural Information Processing Systems, 2021.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in neural information processing systems, 35:26565-26577, 2022.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), ICLR, 2014.
Seungjae Lee, Yibin Wang, Haritheja Etukuru, H Jin Kim, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. Behavior generation with latent actions. Forty-first International Conference on Machine Learning, 2024.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. Journal of Machine Learning Research, 17(39):1-40, 2016.
Baihan Lin, Djallel Bouneffouf, and Irina Rish. A survey compositional generalization in applications. arXiv preprint arXiv:2302.01067 2023.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in neural information processing systems, 13::11525-11536 2020.
Corey Lynch, Ayzaan Wahid, Jonathan Tompson, Tianli Ding, James Betker, Robert Baruch, Travis Armstrong, and Pete Florence. Interactive language: Talking to robots in real time. IEFE Robotics and Automation Letters, 2023.
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, and Gavriel State. Isaac gym: High performance gpu-based physics simulation for robot learning, 2021.
Luke Melas Melas-Kyriadd Kyriazi, Christian Rupprecht, and Andrea Vedaldi. Pc2: Projection-conditioneed point cloud diffusion for single-image 3d reconstruction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12923-12932, 2023.
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. Advances in neural information processing systems, 31, 2018
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEFF/CVF International Conference on Computer Vision, pp.
2023.
Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using score-based diffusion policies. arXiv preprint arXiv:2304 0200220,..222.... 2023.
Maximilian Seitzer, Max Horn, Andrii Zadaianchuk, Dominik Zietlow, Tianjun Xiao, Carl-Johann Simon-Gabrie Tong He, Zheng Zhang, Bernhard Schölkopf, Thomas Brox, et al. Bridging the gap to real-world object-centric learning. In The Eleventh International Conference on Learning Representations, 2023.
13

---

Published as a conference paper at ICLR 2025
Junyao Shi, Jianing Qian, Yecheng Jason Ma, and Dinesh Jayaraman. Composing pre-trained object- centric representations for robotics from" what" and" where" foundation models. arXiv preprint 1344: 2024.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Yang Song. Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01 2023.
Arash Vahdat, Francis Williams, Zan Gojeic, Or Litany, Sanja Fidler, Karsten Kreis, et al. Lion: Latent point diffusion models for 3d shape generation. Advances in Neural Information Processing Systems, 35:10021-10039, 2022.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Ziyi Wu, Nikita Dvomik, Klaus Greff, Thomas Kipf, and Animesh Garg. Slotformer: Unsupervised visual dynamics simulation with object-centric models. In The Eleventh International Conference on Learning Representations, 2023.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the aaai conference on artificial intelligence, pp. 10674-10681, 2021.
Jaesik Yoon, Yi-Fu Wu, Heechul Bae, and Sungjin Ahn. An investigation into pre-training object- centric representations for reinforcement learning. In International Conference on Machine Learning, pp. 40147-40174 PMLR, 2023.
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Self-supervised visual reinforcement learning with object-centric representations. In International Conference on Learning Representa- tions, 2021.
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Object-centric learning for real-world videos by predicting temporal feature similarities. Advances in Neural Information Processing Systems, 36, 2024.
Linfeng Zhao, Lingzhi Kong, Robin Walters, and Lawson LS Wong. Toward compositional general- ization in object-oriented world modeling. In International Conference on Machine Learning, pp. 26841-26864. PMLR, 2022.
Fanggi Zhu, Hongtao Wu, Song Guo, Yuxiao Liu, Chilam Cheang, and Tao Kong. Irasim: Learning interactive real-robot action simulators. arX6v:2406. /2802, 2024.
Yifeng Zhu, Abhishek Joshi, Peter Stone, and Yuke Zhu. Viola: Imitation learning for vision-based manipulation with object proposal priors. 6th Annual Conference on Robot Learning (CoRL), 2022.
14

---

Published as a conference paper at ICLR 2025
Appendix
A EXTENDED DEEP LATENT PARTICLES (DLP) BACKGROUND
In this section, we present an expanded overview of the Deep Latent Particles (DLP) object-centric representation, as introduced by Daniel & Tamar (2022) and Daniel Tamar (2024). DLP an unsupervised, VAF-based model for object-centric image representation. Its core idea is structuring the VAE's latent space set of M particles,
[= 3pl E (Mx()))), Here, 2 E R M encodes visual appearance features, while Ep IR 1g x2 represents particle positions as (z, y) coordinates in pixel-space, i.e., keypoints. Following is a description of the several modifications the standard VAE framework introduced in DLP.
Prior: DLP employs an image-condine conditiond prior p(z|z), with distinct structures for 2f and 2p- P(zplz) comprises Gaussians centered on keypoint proposals, generated by applying CNN to image patches and processed through a spatial-softmax (SSM,akkbb et al. [2018] Finn et al. (0166). The features 2 do not have special prior neural network, and the standard zero-mean unit Gaussian, N(, I), is used.
Encoder: A CNN-based encoder maps the input image to means and log- oorvariiances for Zp (or offsets from them). For 2y. Spatial Transformer Network (STN) (Jaderberg et al. 2015) encodes features from regions ("glimpses") around each keypoint.
Each particle is independently decoded to reconstruct its RGBA glimpse patch (where "A" is the alpha channel of each particle). These glimpses are then composited based on their encoded positions to reconstruct the full image.
Decoder:
Loss: The entire DLP model is trained end-to-end in unsupervised manner by maximizing the ELBO, i.e., minimizing the reconstruction loss and the KL.-divergence between posterior and prior distributions.
KL Loss Term: The posterior keypoints S1 and prior keypoint proposals S2 form unordered sets of Gaussian distributions. As such, the KL term for position latents is replaced with the Chamfer-KL. (
(S, minx,es, (=ppll=) +E,fes, min,,,,,, K (Zppll=())
->.,s,,
DLPv2: Daniel & Tamar (0244) extend the original DLP by incorporating additional particle attributes. DLPv2 provides disentangled latent space structured as a set of M foreground particles: 2 {(pp. &e, Zd, 2e, 2 )
M-1 e RM (6+m), Here, 2p E R2 and 2y E IR" remain unchanged, while new attributes are introduced and described below.
2, E 2 scale, representing the (x. y) dimensions of the particle's bounding box, 2d E R: approxi- mate depth in pixel space, determining particle overlap order when particles are close, and 2 E R[o.1]; transparency.
Additionally, DLPv2 introduces a single abstract background particle, always centered in the image and described by flg latent visual features: Zbg (Mogg xgg E Re In work, following Haramati et al. (2024), we discard the background particle from the latent representation after pre- training the DLP. DLPv2 training follows similar approach to the standard DLP, with modifications to encoding and decoding processes to accommodate the additional attributes.
B ENVIRONMENTS AND DATASETS
In this section we give further details about each environment we used in our experiments including how the demonstration datasets were collected and the metrics we use evaluate performance.
PushCube: SkaacGym-based (Makoviychuk et al. 2021 tabletop manipulation environment in- troduced in Haramati et al. (2024). A Franka Panda robotic arm is required to push cubes in different colors to goal positions specified by images. The agent perceives the environment from two views (front and side) and performs actions in the form of deltas in the end effector coordinates (Axre: Ayee 2 Demonstration data for each task (number of objects) was collected by deploying an ECRL (Haramati et al. (024) state-basee agent trained on the corresponding number of objects. We collect 2000 trajectories per task, each containing 30, 50, 100 transitions for 2,

---

Published as a conference paper at ICLR 2025
objects respectively.
For object-centric image representations, we train a single DLP model on a total of 600, 000 images collected by random policy from 2 views (300, 000 transitions) distinct colors. DLP was able to generalize well to images with fewer objects.
environment with 6 cubes in
To get a wide picture of the goal-reaching performance, we consider the following metries: Success- A trajectory is considered a success if at the end of it, all N objects are at threshold distance from their desired goal. The threshold is slightly smaller than the effective radius of a cube. Success Fraction The fraction of objects that meet the success condition.
Maximum Object Distance The largest distance of an object from its desired goal. Averuge Object Distance - The average distance of all objects from their desired goal.
PushT: IsaacCym- (Makoviychuk et al. 2021) tabletop manipulation environment introduced in Haramati et al. (2024). A Franka Panda robotic arm is required to push T-shaped blocks to goal orientations specified by images. The object position is not considered part of the task in this setting. The agent perceives the environment from two views (front and side) and performs actions in the form of deltas in the end effector coordinates a = (Areer Ayeer A AZee). Demonstration data for each task (number of objects) was collected by deploying agent trained on the corresponding number of objects. We collect 2000 trajectories per task, each containing 50, 100, 150 transitions for 1, 2, 3 objects respectively.
ECRL (Haramati et al. 2024) state- base-
For object-centric image representations, we train a single DL.P model on a total of 600, 000 images collected by random policy from 2 views (300, 000 transitions) on an environment with 3 T-blocks in distinct colors. DLP was able to generalize well to images with different numbers of objects. Since any orientation threshold we choose to define success would be arbitrary, we use the following metric to asses performance:
Average Orientation Distance The average distance of all objects from their desired goal orientation in radians. Since the distance can be considered with respect to both directions of rotation, we always take the smaller of the two. Thus, the largest distance would be т and the average of a random policy T/2 (orientations uniformly randomly sampled).
Initially introduced in Gupta ct al. (2020), the agent is required to complete a set of 4 out of possible tasks in kitchen environment: (1) Turn bottom burner by switching a knob; (2) Turn on top burner by switching a knob; (3) Turn on a light switch; (4) Open a sliding cabinet door; (5) Open a hinge cabinet door; (6) Open microwave door; (7) Move kettle from the bottom to top burner. The action space includes the velocities of the 7 DOF robot joints as well its right and left gripper. totaling in 9 dimensions. A full documentation can be found in: https: robot ics. farama. org/ envs/ franka www.ehteteeeeeeccc use the conl-conditioned image-based variant from Lee et al. (2024), where the environment is perceived from a single view and the goal is specified by the last image in the demonstration trajectory. The demonstration dataset contains 566 colleced trajectories of the robot completing out of the 7 tasks in varying order with the longest trajectory length being 409 timesteps.
FrankaKit chen:
tattssiiicchrrrti chen// We
For object-centric image representations, we train a single DLP model on the demonstration data. Performance is measured by the number of goals reached in trajectory (Goals Reached), and the maximum value is 4.
C IMPLEMENTATION DETAILS
C.I EC-DIFFUSER
We build the Transformer network of EC-Diffuser on top of the Particle Interaction Transformer (PINT) modules from Daniel & Tamar (2024). We remove the positional embedding for the particles to ensure the Transformer is types of positional embedding for views, action particles, and timesteps. We leverage the Diffuuer nner et al. (2022) codebase to train our model: hatps:: github. com/ aannerm/ diffuser The hyper- paramttarseeters use in training the EC-Diffuser model is shown in Table
w.r.t the particles. We add different
We additionally provide details on the model sizes and compute resources used in our experiments in T For GPUs, use both NVIDIA RTX A5500 (20GB) and NVIDIA A40 (40GB), though our model training requires only around 8GB of memory. All baseline models have networks of comparable size and are trained on the same hardware.
2

---

Published as a conference paper at ICLR 2025
Batch size
32
Learning rate
8c-5
Diffusion steps
5. 100 (generalization tasks)
Horizon
Number of heads
8
Number of layers
6, 12 (generalization tasks)
Hidden dimensions
256, 512 (generalization tasks)
Table 5: Hyper-parametern used for the EC-Diffuser model.
Model
Parameters
Concurrent GPUs
GPU Hours
EC-Diffuser
6M
12
CC-Diffuser (generalization)
60M
288
Table 6: Compute resources used for the EC-Diffuser model.
C.2 BASELINES
VQ-BeT (Lee al. 2024): A SOTA BC method that uses Transformer-basedd architecture to predict actions in the quantized latent space of a VQ-VAE. When learning from images, they use a pretrained ResNet18 backbone to acquire an image representation. They experiment with both freezing and finetuning this backbone and report improved performance when finetuning. At the time of writing this paper, code implementing training with finetuning the ResNet was not available. We therefore experimented with either using a frozen ResNet or VQ-VAE encoder pretrained on environment images as image representations, We report results from the best performing variant in each environment. We use the official code implementation that can be found in: https
com/ /ay.:a00011/vaa_eeet.ooooo
Diffuser (anner al. 2022): A diffusson--ased decision-making algorithm that we built our method on, thus providing natural baseline. Diffuser trains a U-Net diffusion model to simultancously generate entire state-action sequences. Being in the BC setting, use the variant that uses unguided sampling. As the original paper does not deal with pixel observations, we provide the model with pretrained image representations. We use VQ-VAE VAE encoder to extract latent representation of images and flatten it for compatibility with the ID U-Net architecture. For FrankaKitchen, we use the representation provided by the frozen ResNet18. We use the official code implementation that be found in: https:://// / liggggttuuuuu hub. com annerm/ di fuser
This method implements naive adaptation of ECRL (Haramati al. 2024) to the BC setting by training the Entity Interaction Transformer (EIT) architecture as a BC policy with an /1 loss on the actions. It uses the DLP representations of images, as our method. We use the official code implementation of the ETT that can be found in: https: / / whw.tb..t..uggg.ggg.oggggg...... com/ DanH i ECRL
EIT+BC:
An entity- diffusion policy inspired by Chi et al., (2023). It uses the DLP representations of images and has a similar architecture to ours but generates action-only instead of state-action sequences. The difference in the architecture we use this method that it uses a encoder-decoder Transformer module. The particles are first encoded with Transformer encoder with self-attentions. Then, in the decoder, we interleave self and cross attention between the actions and the particle embedding to obtain denoised actions. The are described in Table 7 For the implementation this method we use the same codebase as for our method.
EC Diffuusion Policy:
used for this method
C.3 PRE-TRAINED REPRESENTATION MODELS
Data: For the IsaacGym environments, similarly to Haramati et al. (2024), we collect 600k images
from viewpoints by interacting with the environment using a random policy for 300k timesteps. For all methods,
RGB images at resolution of 128 128, i.c., I E R 1228xx x28x3 For
FrankaKitchen, we use the offline demonstration dataset collected by Lee et al. (2024) that contains 566 trajectories and around 200k images.

---

Published as a conference paper at ICLR 2025
Batch size
Learning rate
8c-5
Diffusion steps
Horizon
5, 16 (FrankaKitchen)
Number of heads encoder
Number of layers encoder
6. 12 (generalization tasks)
Number heads decoder
Number of layers decoder
6, 12 (generalization tasks)
Hidden dimensions
256. 512 (generalization tasks)
Table Hyper-parameters used for the EC Diffusion Policy model.
Deep Latent Particles (DLP) (Daniel & Tamar 2022): We follow Haramati al. (2024) and train DLPv2 using the publicly available codebase: ht tittps:: / /i thub. com/ t. da techeechaee /ddl on the image datasets. Similarly to Haramati al. 2024), we assign the background particle features dimension of 1, and discard the background particle for the BC stage. The motivation for this is to limit the background capacity to capture changing parts of the scene such as the objects or the agent. We use the default recommended parar-pprsmeters and report the data-specific prppeepprraetees in Table g
Batch size
64
Posterior kp M
24 (IsaacsGym), 40 (Franka
Prior kp proposals L
32 (IsaacsGym), 64 (FrankaKitchen)
Reconstruction loss
MSE
BKL.
Prior patch size
16
Glimpse size S
32
Feature dim n
Background feature dim Tbg
Epochs
60 (IsaacsGym), 250 (FrankaKitchen)
Table Hyper-parameters used for the Deep Latent Particles (DLP) object-centric model.
Vector-Quantizer Variational AutoEncoder (VQ-VAE) (Van Den Oord et al. 2017): We fol-
low Haramati et al. (2024) and train VQ-VAE models using their publicly available codebase: http//DannrmmmiBBB İECCL This aims to provide an unstructured representa- tion in contrast to DLP. We use the default recommended hyper-parameters and report the data-specific hyper-parameters in Table
ht https: / /itthb com
Batch size
Learning rate
20-4
Reconstruction loss
MSE
3KL.
0.
Prior patch size
16
N embed
1024 (PushCube), 2048 (PushT)
Embed dim
Latent dim
256
Epochs
150
Table 9: Hyper--arametted used for VQ-VAE VAE model.
D ADDITIONAL RESULTS
D.1 ADDITIONAL METRICS
First, we report the success fraction (i.e., the proportion of objects that meet the success condition) for PushCube across all methods in Table 10 This aims supplement the findings presented
4

---

Published as a conference paper at ICLR 2025
in the main experimental table by offering per-object success data. For EC- -C-Diffuser, the success fraction remains consistent as the number of objects increases, highlighting our method's 's superior performance in multi-object manipulation. Additionally, we provide more comprehensive results on PushCube generalization in Table [T] which includes additional metrics (i.e. Maximum Object
Distance, Average Object Distance) from the task.
Eav (Metrie)
# Obj
VQ-BeT
Diffuser
EIT+BC (DLP)
EC Diffusion Policy (DLP)
EC-Diffuser (DLP)
0,929 0,032
0.367 + 0.027
0.890 + 0.019
01887 + 0.031
0.948 + 0.015
PushCube
0.207 0.020
0.083 = 0.010
0.342 + 0.140
0.443 0.086
0.948 0.023
(Success Rate
0.097 0.015
0.054 + 0.012
0.396 + 0.239
0.807 0.121
0.950 0.016
Table 10: Success fractions for different methods in the PushCube environments for varying number of objects.
Methods are trained for 1000 epochs, and best performing checkpoints are reported. The best values are in bold. The values are computed as the mean over 96 randomly initialized configurations, and standard deviations are computed across 5 seeds.
Number of Objects
Success Rate t
Success Fraction t
Max Obj Dist +
Avg Obj Dist
0.993 = 0.006
0,993 + 0.01
0,011 0.001
0.011 0.001
2
0.968 = 0.010
0,981 + 0.003
0.023 + 0.005
0.016 = 0.002
(training)
0.833 = 0.031
0.886 0.051
0.056 1 0.020
0.030 0.009
4
0.625 = 0.018
0.858 + 0.002
0.014
0.045 0.003
5
0.448 = 0.057
0.767 0.032
0.198 1 0.025
0.070 = 0.011
6
0.240 = 0.075
0.711 + 0.070
0.253 + 0.056
0.089 = 0.023
Table 11: Our method's 's compositional generalization performance on different numbers of cubes the PushCube environment, trained on 3 objects random colors chosen out options. The values are computed as the mean over 96 randomly initialized configurations, and standard deviations are calculated across 5 seeds.
D.2 ADDITIONAL OBJECT-CENTRIC REPRESENTATIONS
We provide experimental results with different types of input to the EC-Diffuser on the PushCube and FrankaKit chen environments. The results are shown in Table 12 and 13 respectivell All methods use the same architecture and training procedure as EC-Diffuser.
PushCube: Slot Attention We train a Slot Attention model (Locatello et al. 2020 to produce set of latent "slot" vectors. Here, we employ the same slot attention model used ECRL (Haramati et al., 2024), which produces 10 slots per image. We treat each slot as an entity and pass them into the same EC-Transformer model as used in our method. This variant achieves good results for 1 Cube, but its performance deteriorates quickly as the number of cubes increases, although it remains better than that the non-object-centric baselines. We attribute this to the fact that, as shown in Figure 15] the slot model occasionally has trouble with individuating nearby objects and represents them in a single slot. Ground-trufi State - we extract the object location information from the simulator and append one-hot identifier to distinguish each object. Each (position, one-hot) vector treated as an entity to be passed into the EC- Transformer. This variant is meant to shed light on the efficacy and generality of our approach by emulating a "perfect" entity-level factorization of images. With these entity-centric set-based state representations EC-Diffuser achieves slightly better performance than using the DLP representation, as expected. Single View This variant only inputs the DLP representation of the front-view image into the EC-Transformer model. We see drop in performance in this case, highlighting the importance of the ability of EC-Diffuser to effectively leverage multi-view perception in order to mitigate the effect of occlusion and leverage complimenting observational information.
Barnkaiitchen We additionally train EC-Diffuser on top of Slot Attention model trained from the FrankaKit chen demonstration data. We use 10 slots, each with latent dimension of 64. We train the model for 100 epochs on multiple seeds, observe the loss has converged and take the best performing seed. EC-Diffuser with Slot Attention achieves state of the art performance (3.340), surpassing both EC-Diffuser with DLP (achieving 3.031) and VQ-BeT (with reported performance of 2.60). Based on the slot decompositions FrankaKitchen (visualized in Figure [6). is difficult to conclude that its superior performance is due to its ability to capture the objects
5

---

Published as a conference paper at ICLR 2025
Inpul Variation
Cube
2 Cubes
3 Cubes
Slot Attention
0.909 + 0.021 009099 + 0.021
0.364 + 0.027 /0446 + 0.034
0.255 + 0.055 /0.493 0.061
Single View
0.934 = 0.015 0,934 + 0.015
0.752 = 0.054/0.822 = 0.041
0.685 = 0.021 0.833 = 0.016
DLP (ours)
0.948 + 0.015 0.948 + 0.015
0.917 + 0.030 / 0.948 = 0.023
0.894 = 0.025 0.950 0.016
Ground-truth State
0.985 = 0.013 / 0,985 + 0.013
0.963 = 0.015/0.958 0.023
0.916 = 0.025 0.930 0.017
Table 12: Quantitative performance for different input variations on PushCube tasks. All methods the same architecture and training procedure as EC-Diffuser. Results are reported as the mean success rate and success fraction (number of object successes total number of objects) for cach task.
of interest. Nevertheless, these results demonstrate the generality of our method with respect to compatibility with different object-centric representations.
Env (Metric)
EC-Diffuser with Slots
EC-Diffuser with DLP
FrankaKitcher
3.340 + 0.097
3.031 0.087
(Goals Reached f)
Table 13: Quantitative performance for Frankakitchen of EC-Diffuser with different object-centric inputs.
D.3 EC-DIFFUSER ON REAL WORLD DATA
We provide preliminary EC-Diffuser results on real world data in this section. Specifically, we use the real-world Language-Taan
dataset (Lynch et al. 2023). We subsample 3000 episodes from the real robot dataset, each padded to 45 images, and we randomly select 2700 episodes for training and 300 for validation. We train the DLP model and EC-Diffuser on the training set. We provide visualization of EC-Diffuser's particle state generation in Figure 8 As shown in the figure, EC- Diffuser can effectively generate high quality rollouts, which shows promise in applying EC-Diffuser to real world problems. We provide DLP decompositions of images from this dataset in Figure 14
Observation
Generated Frames
Goal
Figure 8: Visualization of EC-Diffuser's generation of DL.P states of language-Table The objects of interest and their respective goals in the generated frames highlighted with an orange circle.

---

Published as a conference paper at ICLR 2025
New Color
Star-Shaped
Rectangular-Shaper
T-Shaped
Figure 9: Generalization to unseen objects. EC-Diffuser trained on PushCube (only cube data in DLP training). From left to right: new colors, star-shaped objects, rectangular-shaped objects and T-shaped objects.
First Observation
EC-Diffuser Generation
Front View
Side View
Figure 10: We decode the DLP states generated by EC-Diffuser for T-shaped objects. Note that the DLP and BC-Diffuser both trained only cubes, but interestingly, the generated shapes combination of cubes resemble T-shaped blocks.
D.4 GENERALIZATION UNSEEN OBJECTS
We provide additional cases of generalization to unseen shapes and colors as described in Figure We report the performance in Table 14] We see that EC-Diffuser coupled with DLP is able to generalize zero-shot with little no drop performance new colors as well as new shapes (star, rectangular cuboid). When replacing cubes with T-shaped blocks there is significant drop in performance although success rate is better than random, suggesting some zero-shot generalization capabilities in this case well. Additionally, we visualize the DLP state generation from the EC-Diffusee (both the DLP encoder and EC-Diffuser trained on PushCube) for the T-shaped blocks in Figure TO We can that the DLP represents the T-block the objects it is trained on. This can be seen as a form of compositional generalization. While we find this is an interesting capability translate to better action or future state generation. EC-Diffuser is still trained for the dynamics of individual cubes and cannot account for dynamics of cubes that are "mended" together form T-block.
collection of cubes, composing the overall scene from
its own, this generalization is only visual and does not always
We 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.
Cube Variation
Cube
2 Cubes
Cubes
New Color
0.958
0.947
0.909
Star-Shaped
0.979
0,916
0.885
Rectangluar- Shaped
0.989
0,906
0.844
T-Shaped
0.531
0.339
0.139
Cube (training)
0.968
0.958
0.895
Table 14: Success rate for generalization to different object types the PushCube tasks, computed as the 96 randomly initialized configurations.
7

---

Published as a conference paper at ICLR 2025
D.5 DLP DECOMPOSITIONS
We provide the visualizations of DLP decomposition by overlaying the particles on top of the original image well the showing reconstructions of the foreground and background from the particles. Figure T1 shows the visualization for PushCube. Figure T2] shows the visualization for PushT. Figure 13 shows the visualization for FrankaKiz chen. addition.
train DLP on images from the real-world Language Table dataset (Lynch et al., (023), and provide visualization of DLP's output Figure DLP provides accurate decompositions of the scene, indicating could be paired with EC-Diffuser for real robotic manipulation environments.
Image
DLP
Foreground
Background
Visualization DLP decomposition PushCube. From top to bottom: original image; DLP overlaid on top of the original image; image reconstruction from the foreground particles; image reconstruction from the background particle.
Figure 11:
Image
DLP
Foreground
Background
Visualization of DLP decomposition of PushT. From top to bottom: original image; DLP overlaid top of the original image; image reconstruction from the foreground particles; image reconstruction from the background particle.
Figure 12:
D.6 SLOT ATTENTION DECOMPOSITIONS
We provide the visualizations of Slot Attention decomposition in this section.
8

---

Published as a conference paper at ICLR 2025
Image
DLP
Foreground
Background
Figure 13: Visualization of DLP decomposition of Frankak thehe.. From top bottom: original image; DLP overlaid top the original image; image reconstruction from the foreground particles; image recon- struction from background particle.
Image
DLP
Foreground
Background
Figure 14: Visualization of DLP decomposition of Language Tangue From top to bottom: original image; DLP overlaid on top of the original image; image reconstruction from the foreground particles; image recon- struction from the background particle.
9

---

Published as a conference paper at ICLR 2025
Image
Recon
Alpha Masks
Slot
Slot 2
Slot
Slot
Slot
Slot
Slot
Figure 15: Visualization Slot decomposition of PushCube. From left right: original image; reconstruction from slots; alpha masks of the slots; per-slot reconstractions.
Image
Recon
Alpha Masks
Slot
Slot
Slot
Slot
Figure 16: Visualization of Slot decomposition of FrankaKitcher From left right: original image; reconstruction from the slots; alpha masks the slots per-slot reconstructions.
10


Reviews:
[{"review_id": "bdtZIVCnXm", "reviewer": "Reviewer_i8yN", "summary": "The paper proposes entity-centric diffusion transformer architecture to generate the sequence of actions and latent states conditioned on a trajectory of image observations across multiple views. The paper claims that object centric representations facilitate composite generalization of behavior cloning agents. Experiments are conducted across three tasks namely, PushCube, PushT and FrankaKitchen to show the effectiveness of the proposed method.", "questions": "Is the DLP pretraining done separately for each task in IsaacGym environment? \n\nFurther, if PushCube and PushT image observations were used together for pretraining DLP, can it handle the composition where both objects are present together?\n\nCan the method generalize in manipulating novel objects? For example, what is the success rate when T-cubes are added to the Task 1 environment instead of normal cubes.\n\nSince competing methods VQ-BET and EIT-BC work well for single objects, for multiple objects, can they be used sequentially manipulating one object at a given time? What is the performance of this baseline?\n\nWhat is the performance of the method when using only the frontal view to train the diffusion model? Are two views essential?\n\nWhy is classifier-free guidance not used in the diffusion model? Can classifier-free guidance improve the success rate?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The proposed method leverages object-centric representations to generate action-state sequences for Behavior Cloning and shows that it generalizes to manipulating multiple objects. Experiments are conducted across three different tasks.", "weaknesses": "The novelty of the proposed method is limited as it combines existing methods on entity-centric representation with diffusion policy. It builds upon the existing method [1] to use diffusion models instead of transformers for generating the action-state sequences for object manipulation. Diffusion models have already been shown to be useful for object manipulation tasks in prior works [2, 3].\n\n[1] Haramati et. al, Entity-Centric Reinforcement Learning for Object manipulation from pixels.\n\n[2] Mishra et. al. ReorientDiff: Diffusion Model based Reorientation for Object Manipulation\n\n[3] Chi et. al. Diffusion policy: Visuomotor policy learning via action diffusion. \n\nMissing experiments on evaluating the learned policies on real world robots.\n\nFigure 4 and Figure 5 overlaps with the captions of the Tables above them.", "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": "lAx9SbOdhX", "reviewer": "Reviewer_1WDY", "summary": "This paper propose a diffusion-based behavioral cloning method. It uses the object-centric representations instead of pixel-level representation. Deep Latent Particle is used to get representation, followed by an entity-centric Transformer at the particle level to predict action sequences.", "questions": "- What's the motivation of using diffusion for this task? What are the advantages of using a diffusion model over other transformer-based methods?\n\n- Is it possible to also test the model on other environments, such as Multimodal Ant, BlockPush, UR3 BlockPush, or even nuScenes self-driving, to better demonstrate the advantage of the proposed method?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This idea of predicting object manipulation actions through a diffusion-based architecture is interesting. By incorporating object-level information, the proposed method achieves a performance improvement over baseline methods.\n\n- This paper also compare with the state-of-the art non-diffusion baseline, i.e., VQ-BeT, demonstrating the effectiveness of the proposed method.", "weaknesses": "- The method depends on the capabilities of the image representation algorithm DLP, and experiments are conducted only in synthetic environments. It is unclear if it will perform well in more complex settings. For example, the method is only compared with one non-diffusion baseline VQ-BeT, which is preformed extremely bad on PushCube and PushT (2 out of 3 test environments used in this paper). I wonder why it is only tested on these three environments. Is it possible to also test the model on other environments that baselines achieve very good results, such as Multimodal Ant, BlockPush, UR3 BlockPush, or even nuScenes self-driving, to better demonstrate the advantage of the proposed method?\n\n- What are the advantages of using a diffusion model over other transformer-based methods? Could you provide an intuitive explanation for why it is suitable for this task?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QjBwHQMjf3", "reviewer": "Reviewer_3tkq", "summary": "The authors propose EC-Diffuser, a diffusion-based behavioral cloning method designed to enhance multi-object goal-conditioned manipulation tasks from high-dimensional pixel observations. This approach leverages a pre-trained object-centric encoder to map each observation into an unstructured set of latent vectors, where each vector corresponds to a specific object or the background within the observation. The diffusion model is trained to denoise these latent trajectories paired with corresponding actions, using the initial and goal representations as conditional information. To address the unordered nature of these latent representations, the authors employ a Transformer-based diffusion model due to its permutation equivariance.\n\nThe authors validate their approach through experiments on several simulated multi-object manipulation environments, demonstrating that it outperforms baseline methods. They also provide ablation studies to highlight the contributions of different components in their method.", "questions": "* __Table 1 results interpretation.__ In Table 1, regarding the performance of the __EC Diffusion Policy__ on the __PushCube__ environment, there is an unusual pattern where the success rate decreases when moving from one to two objects but then increases from two to three objects. In contrast, all other methods exhibit a consistent decrease in performance as the number of objects increases. Could the authors provide an explanation for this anomaly in their results?\n\n* __Real-world limitations.__ It appears that the object-centric encoder is the main bottleneck when it comes to extending this model to real-world scenarios. Is this assessment correct? If so, do the authors believe that object-centric approaches are preferable for real-world applications, and how might they address the challenges posed by the encoder in practical implementations?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* __Writing.__ The paper is well-written and easy to follow. The motivation behind addressing the challenges in multi-object manipulation from high-dimensional pixel observations is clearly articulated. \n\n* __Novelty.__ The authors present a novel integration of existing approaches by combining object-centric representations (Deep Latent Particles) with a diffusion-based behavioral cloning method. The use of an entity-centric Transformer to handle the unordered nature of latent representations is a noteworthy contribution.\n\n* __Positioning.__ The paper is well-situated within the existing body of research, providing a comprehensive review of related work and clearly distinguishing its contributions from prior methods.\n\n* __Experimental design.__ The authors conduct thorough experiments across several simulated multi-object manipulation environments. They adapt and compare against multiple competitive baselines, ensuring a fair and rigorous evaluation of their method. They also thoroughly compare their approach with both structured and unstructured latent representation-based methods, providing a comprehensive evaluation.\n\n* __Performance.__ Experimental results demonstrate that the proposed approach consistently outperforms baseline methods.", "weaknesses": "* __Alternative object-centric encoders.__ While the authors utilize DLPv2 as a powerful object-centric encoder, the paper would benefit from experimenting with other similar approaches. Since the DLP module appears to be easily replaceable with other object-centric encoders, exploring alternatives could provide insights into the generality and robustness of the proposed method. This would also demonstrate whether the observed benefits are specific to DLPv2 or applicable across different encoding techniques.\n\n* __More generalization experiments.__ Although the authors present a wide range of experiments showcasing the advantages of their approach, additional generalization studies could strengthen the paper. Specifically, it would be interesting to see experiments involving novel object shapes, textures, or views where the use of object-centric latent representations might not be as beneficial at test time compared to regular unstructured latents. \n\nMinor Presentation Issues:\n\n* There is a rendering issue with Figure 5, where the figure overlaps with the text.\n\n* In Tables 2 and 3, the authors use __bold text__ to indicate the training setup. It might be clearer to mention this in parentheses in the column header. This is a minor stylistic suggestion but could improve the overall presentation.", "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": "jU8lYAlbFL", "reviewer": "Reviewer_WGqn", "summary": "The paper presents a method that tackles the problems of behavior cloning and object manipulation. The method contains three main components: an object-centric encoder (DLP encoder), an entity-centric transformer (for temporal information aggregation), and a goal-conditioned diffusion-based decoder (for simultaneous prediction of states and actions). These components were proposed by previous methods, and EC-Diffuser appears to be a combination of several approaches. The results are superior compared to previous work on three simulation benchmarks. Furthermore, it can generalize to scenarios with a higher number of objects (trained on fewer objects, tested on more objects). The ablation studies (Table 5) show the effect of each component, revealing that the DLP encoder and goal-conditioned diffusion parts are the most important ones. Overall, the proposed method appears to be the new state-of-the-art. However, it seems to be a combination of previous methods: DLP (for object encoding), transformer-based sequence modeling (VQ-Bet), and diffusion-based decoding (diffuser). By combining these components, EC-Diffuser gains advantages over previous work through its use of the DLP encoder (better object-centric representations compared to VQ-Bet), improved entity-centric transformer (temporal information aggregation compared to DLPv2, which uses autoregressive prediction), and diffusion decoder (better handling of multimodality and uncertainty compared to VQ-Bet).\n\n**Important note:** This review's technical content and analysis are my original work. Large Language Models were used solely to improve grammar and writing style, without altering any meanings or substantive feedback.", "questions": "* As stated in the limitations, the model cannot capture the object-centric notion in scenes from the FrankaKitchen environment. The paper suggests that the DLP encoder could be modified to incorporate additional information, such as ResNet features. Could Slot Attention be used to obtain better object-centric features in FrankaKitchen?\n* Could EC-Diffuser be applied to real-world datasets like NuScenes, setting aside its inference cost considerations? (See Weaknesses section as well.)\n* Can EC-Diffuser generalize to different time horizons?\n    - For instance, if the training set uses a time horizon H, could the model be tested on longer time horizons such as 2H?\n    - Are there any architectural considerations that might enable or limit generalization to longer time horizons?\n* Can the diffusion model's conditioning mechanism (goal action) be replaced with cross-attention instead of AdaLN?\n    - Could the conditioning be improved this way, since any input could be tokenized and fed to the cross-attention layers? The authors' discussion and ideas would be informative.\n    - Most state-of-the-art diffusion transformers (DiT) don't use cross-attention but instead use Multimodal DiTs, where both input and condition are tokenized and concatenated (as seen in SD3, SD3.5, etc.). Therefore, the claim about using AdaLN would also be reasonable.\n* Minor formatting error: Table 1 and Figure 4 overlap, appearing to be a vertical spacing error.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Experiments:\n   - The paper evaluates the method on three different benchmarks and provides ablation studies that demonstrate the effect of each component\n   - Section 5.4 demonstrates the necessity of state generation, as it helps match objects across different frames. This evaluation approach was both effective and informative\n* Performance:\n  - The results demonstrate that EC-Diffuser performs exceptionally well on simulation benchmarks and can generalize to different scenarios (as shown in the object-number generalization experiments in Section 5.2)\n* Code is available, and video results are presented on the website for better evaluation\n* Overall, the paper clearly explains the problem, tasks, and methodology", "weaknesses": "* Novelty:\n     - As previously stated, EC-Diffuser appears to be a combination of previous work\n     - The model incorporates the best components from several predecessors, such as:\n          - Object-centric encoder from DLP (while its usage is fair, the results should perhaps not depend so heavily on the DLP encoder; without it, the performance is comparable to VQ-Bet)\n          - Transformer encoder from VQ-Bet, which uses a GPT-like transformer encoder on observations encoded by a VQ-VAE\n          - Diffusion decoder that predicts both states and actions, similar to Diffuser which also employs this architecture for denoising both state and action\n* Real-world evaluation:\n    - Can this model be evaluated on real-world datasets such as NuScenes (as VQ-Bet does)?\n    - Can EC-Diffuser be adapted for use as GPT-driver?\n     - Is it possible to generalize to real-world setups, or would EC-Diffuser be limited to simulation data?\n     - What are the necessary changes to make EC-Diffusers work on real-world data?\n* Inference runtime cost:\n    - As stated in the limitations and future work section, diffusion inference is time-consuming due to its iterative denoising requirement\n     - While real-time inference might not be crucial for simulation data, it is essential for tasks such as autonomous driving where real-time decisions are critical", "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 (Poster)